[ { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[codz]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.buildx-cache/\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# UV\n# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n#uv.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#poetry.toml\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.\n# https://pdm-project.org/en/latest/usage/project/#working-with-version-control\n#pdm.lock\n#pdm.toml\n.pdm-python\n.pdm-build/\n\n# pixi\n# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.\n#pixi.lock\n# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one\n# in the .venv directory. It is recommended not to include this directory in version control.\n.pixi\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# Abstra\n# Abstra is an AI-powered process automation framework.\n# Ignore directories containing user credentials, local state, and settings.\n# Learn more at https://abstra.io/docs\n.abstra/\n\n# Visual Studio Code\n# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore \n# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore\n# and can be added to the global gitignore or merged into this file. However, if you prefer, \n# you could uncomment the following to ignore the entire vscode folder\n.vscode/\n\n# Ruff stuff:\n.ruff_cache/\n\n# PyPI configuration file\n.pypirc\n\n# Cursor\n# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to\n# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data\n# refer to https://docs.cursor.com/context/ignore-files\n\ntest/\ncheckpoints/\n.cursorignore\n.cursorindexingignore\n\n# Marimo\nmarimo/_static/\nmarimo/_lsp/\n__marimo__/\n\n.DS_Store" }, { "path": ".gitmodules", "content": "[submodule \"camera_pose_annotation/base\"]\n\tpath = camera_pose_annotation/base\n\turl = https://github.com/SpatialVID/base.git\n" }, { "path": "Dockerfile.cuda", "content": "# This Dockerfile builds FFmpeg with NVIDIA GPU support and libvmaf from source\n# It uses a two-stage build to create a smaller runtime image\n# This file is adapted from https://github.com/Netflix/vmaf/blob/master/Dockerfile.cuda\n\nARG CUDA_BASE_IMAGE=docker.io/nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04\nARG RUN_TIME_IMG=docker.io/nvidia/cuda:12.6.3-runtime-ubuntu22.04\n# ARG CUDA_BASE_IMAGE=swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04\n# ARG RUN_TIME_IMG=swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/nvidia/cuda:12.6.3-runtime-ubuntu22.04\n\nFROM $CUDA_BASE_IMAGE as builder\n\n\nARG VMAF_TAG=master\nARG FFMPEG_TAG=master\n\nRUN DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y libopenjp2-7-dev \\\n ninja-build cmake git python3 python3-pip nasm xxd pkg-config curl unzip nvidia-cuda-toolkit\n\nRUN git clone https://github.com/Netflix/vmaf.git && cd vmaf && git checkout $VMAF_TAG\n\nRUN git clone https://github.com/FFmpeg/FFmpeg.git && cd FFmpeg && git checkout $FFMPEG_TAG\n\nRUN git clone https://github.com/FFmpeg/nv-codec-headers.git && cd nv-codec-headers && make && make install\n\n# install vmaf\nRUN python3 -m pip install meson\nRUN cd vmaf && meson libvmaf/build libvmaf -Denable_cuda=true -Denable_avx512=true --buildtype release && \\\n ninja -vC libvmaf/build && \\\n ninja -vC libvmaf/build install\n\nENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/x86_64-linux-gnu/\nRUN ldconfig\n\n# install ffmpeg\nRUN cd FFmpeg && ./configure \\\n --enable-libnpp \\\n --enable-nonfree \\\n --enable-nvdec \\\n --enable-nvenc \\\n --enable-cuvid \\\n --enable-cuda \\\n --enable-cuda-nvcc \\\n --enable-libvmaf \\\n --enable-ffnvcodec \\\n --disable-stripping \\\n --extra-cflags=\"-I/usr/local/cuda/include\" \\\n --extra-ldflags=\"-L/usr/local/cuda/lib64 -L/usr/local/cuda/lib64/stubs/\" \n\nRUN cd FFmpeg && make -j && make install\n\nRUN mkdir /data\n\n\n# Create a smaller runtime image\nFROM ${RUN_TIME_IMG} as runtime\nENV DEBIAN_FRONTEND=noninteractive\nRUN apt-get update && apt-get install -y --no-install-recommends \\\n ca-certificates python3 python3-pip python3-venv libnuma-dev libsm6 libxext6 libxrender1 libgl1 git vim && rm -rf /var/lib/apt/lists/*\n\nWORKDIR /workspace\n\n# Copy FFmpeg and libvmaf from builder (installed under /usr/local)\nCOPY --from=builder /usr/local /usr/local\n# copy libraries installed by the builder stage if present\nCOPY --from=builder /usr/lib/ /usr/lib/\n\n# Link python\nRUN ln -sf /usr/bin/python3 /usr/bin/python\nRUN python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel\n\n# Copy repository\nCOPY . /workspace\nRUN apt-get update\n\n# Install Python requirements (may still fail for some packages requiring system libs)\nRUN python3 -m pip --no-cache-dir install -r requirements/requirements.txt\nRUN python3 -m pip --no-cache-dir install -r requirements/requirements_scoring.txt || true\nRUN python3 -m pip --no-cache-dir install -r requirements/requirements_annotation.txt || true\n\n# Entrypoint\nCOPY docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh\nRUN chmod +x /usr/local/bin/docker-entrypoint.sh\n\nENV FFMPEG_PATH=/usr/local/bin/ffmpeg\nENTRYPOINT [\"/usr/local/bin/docker-entrypoint.sh\"]\nCMD [\"bash\", \"ldconfig\"]\n" }, { "path": "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": "README.md", "content": "

SpatialVID: A Large-Scale Video Dataset with Spatial Annotations

\n
\n Jiahao Wang1* \n Yufeng Yuan1* \n Rujie Zheng1* \n Youtian Lin1 \n Jian Gao1 \n Lin-Zhuo Chen1 \n
\n
\n Yajie Bao1 \n Yi Zhang1 \n Chang Zeng1 \n Yanxi Zhou1 \n Xiaoxiao Long1 \n Hao Zhu1 \n
\n
\n Zhaoxiang Zhang2 \n Xun Cao1 \n Yao Yao1†\n
\n
\n 1Nanjing University  2Institute of Automation, Chinese Academy of Science \n
\n
\n *Equal Contribution  †Corresponding Author\n
\n
\n CVPR 2026\n
\n
\n
\n  \n  \n  \n  \n \n
\n

\n \n

\n\n## 🎉NEWS\n+ [2026.02.21] 🎉 SpatialVID is accepted by CVPR 2026!\n+ [2025.10.11] 🐳 Docker support is now available, featuring a pre-configured environment with NVIDIA GPU-accelerated FFmpeg.\n+ [2025.09.29] 🚀 Depth data for the SpatialVID-HQ dataset is now officially available.\n+ [2025.09.24] 🤗 Raw metadata access is now available via a [gated HuggingFace dataset](https://huggingface.co/datasets/SpatialVID/SpatialVID-RAW) to better support community research!!\n+ [2025.09.24] 🔭 Enhanced instructions for better camera control are updated.\n+ [2025.09.18] 🎆 SpatialVID dataset is now available on both HuggingFace and ModelScope.\n+ [2025.09.14] 📢 We have also uploaded the SpatialVID-HQ dataset to ModelScope offering more diverse download options.\n+ [2025.09.11] 🔥 Our paper, code and SpatialVID-HQ dataset are released!\n \n**[✍️ Note]** Each video clip is paired with a dedicated annotation folder (named after the video’s id). The folder contains 5 key files, and details regarding these files can be found in [Detailed Explanation of Annotation Files](https://huggingface.co/datasets/SpatialVID/SpatialVID#3-detailed-explanation-of-annotation-files).\n\n## Abstract\n\nSignificant progress has been made in spatial intelligence, spanning both spatial reconstruction and world exploration. However, the scalability and real-world fidelity of current models remain severely constrained by the scarcity of large-scale, high-quality training data. While several datasets provide camera pose information, they are typically limited in scale, diversity, and annotation richness, particularly for real-world dynamic scenes with ground-truth camera motion. To this end, we collect **SpatialVID**, a dataset consisting of a large corpus of in-the-wild videos with diverse scenes, camera movements and dense 3D annotations such as per-frame camera poses, depth, and motion instructions. Specifically, we collect more than **21,000 hours** of raw videos, and process them into **2.7 million clips** through a hierarchical filtering pipeline, totaling **7,089 hours** of dynamic content. A subsequent annotation pipeline enriches these clips with detailed spatial and semantic information, including camera poses, depth maps, dynamic masks, structured captions, and serialized motion instructions. Analysis of SpatialVID's data statistics reveals a richness and diversity that directly foster improved model generalization and performance, establishing it as a key asset for the video and 3D vision research community.\n\n\n## Preparation\n\nThis section describes how to set up the environment manually. For a simpler, containerized setup, please refer to the **[Docker Setup and Usage](#docker-setup-and-usage)** section.\n\n### Environment\n\n1. Necessary packages\n\n ```bash\n git clone --recursive https://github.com/NJU-3DV/SpatialVID.git\n cd SpatialVid\n conda create -n SpatialVID python=3.10.13\n conda activate SpatialVID\n pip install -r requirements/requirements.txt\n ```\n2. Package needed for scoring\n\n ```bash\n pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/\n pip install -r requirements/requirements_scoring.txt\n ```\n\n Ignore the warning about `nvidia-nccl-cu12` and `numpy` version, it is not a problem.\n\n About FFMPEG, please refer to the [`INSTALL.md`](scoring/motion/INSTALL.md) for detailed instructions on how to install ffmpeg. After installation, replace the `FFMPEG_PATH` variable in the [`scoring/motion/inference.py`](scoring/motion/inference.py) and [`utils/cut.py`](utils/cut.py) with the actual path to your ffmpeg executable, default is `/usr/local/bin/ffmpeg`.\n\n ⚠️ If your videos are in av1 codec instead of h264, you need to install ffmpeg (already in our requirement script), then run the following to make conda support av1 codec:\n\n ```bash\n pip uninstall opencv-python\n conda install -c conda-forge opencv==4.11.0\n ```\n\n If unfortunately your conda environment still cannot support av1 codec, you can use the `--backend av` option in the scoring scripts to use PyAV as the video reading backend.\n But note that using PyAV for frame extraction may lead to slight inaccuracies in frame positioning.\n\n3. Package needed for annotation\n\n ```bash\n pip install -r requirements/requirements_annotation.txt\n ```\n\n Compile the extensions for the camera tracking module:\n\n ```bash\n cd camera_pose_annotation/base\n python setup.py install\n ```\n\n4. [Optional] Package needed for visualization\n\n ```bash\n pip install plotly\n pip install -e viser\n ```\n\n### Model Weight\n\nDownload the model weights used in our experiments:\n\n```bash\nbash scripts/download_checkpoints.sh\n```\n\nOr you can manually download the model weights from the following links and place them in the appropriate directories.\n\n| Model | File Name | URL |\n| ------------------- | ----------------------- | --------------------------------------------------------------------------------------------------------------- |\n| Aesthetic Predictor | aesthetic | [🔗](https://github.com/christophschuhmann/improved-aesthetic-predictor/raw/main/sac+logos+ava1-l14-linearMSE.pth) |\n| MegaSAM | megasam_final | [🔗](https://github.com/mega-sam/mega-sam/blob/main/checkpoints/megasam_final.pth) |\n| RAFT | raft-things | [🔗](https://drive.google.com/uc?id=1MqDajR89k-xLV0HIrmJ0k-n8ZpG6_suM) |\n| Depth Anything | Depth-Anything-V2-Large | [🔗](https://huggingface.co/depth-anything/Depth-Anything-V2-Large) |\n| UniDepth | unidepth-v2-vitl14 | [🔗](https://huggingface.com/lpiccinelli/unidepth-v2-vitl14) |\n| SAM | sam2.1-hiera-large | [🔗](https://huggingface.co/facebook/sam2.1-hiera-large) |\n\n\n## Quick Start\n\nThe whole pipeline is illustrated in the figure below:\n\n

\n \n

\n\n1. Scoring\n\n ```bash\n bash scripts/scoring.sh\n ```\n\n Inside the [`scoring.sh`](scripts/scoring.sh) script, you need to set the following variables:\n\n - `ROOT_VIDEO` is the directory containing the input video files.\n - `OUTPUT_DIR` is the directory where the output files will be saved.\n2. Annotation\n\n ```bash\n bash scripts/annotation.sh\n ```\n\n Inside the [`annotation.sh`](scripts/annotation.sh) script, you need to set the following variables:\n\n - `CSV` is the CSV file generated by the scoring script, default is `$OUTPUT_DIR/results.csv`.\n - `OUTPUT_DIR` is the directory where the output files will be saved.\n3. Caption\n\n ```bash\n bash scripts/caption.sh\n ```\n\n Inside the [`caption.sh`](scripts/caption.sh) script, you need to set the following variables:\n\n - `CSV` is the CSV file generated by the annotation script, default is `$OUTPUT_DIR/results.csv`.\n - `SRC_DIR` is the annotation output directory, default is the same as the `OUTPUT_DIR` in the annotation step.\n - `OUTPUT_DIR` is the directory where the output files will be saved.\n - The API keys for the LLM models used in the captioning step. You can replace them with your own API keys.\n4. Visualization\n\n - You can visualize the `poses.npy` in the `reconstruction` folder of each annotated clip using the [`visualize_pose.py`](viser/visualize_pose.py) script.\n - You can visualize the final annotation result(`sgd_cvd_hr.npz`) using the [`visualize_megasam.py`](viser/visualize_megasam.py) script.\n \n Note that if you want to visualize any clip in our dataset, you need to use the script [`pack_clip_assets.py`](utils/pack_clip_assets.py) to unify the depth, RGB frames, intrinsics, extrinsics, etc. of that clip into a single npz file first. And then you can use the visualization script to visualize it.\n\n\n## Docker Setup and Usage\n\nWe provide a Dockerfile to create a fully configured environment that includes all dependencies, including a custom-built FFmpeg with NVIDIA acceleration. This is the recommended way to ensure reproducibility and avoid environment-related issues.\n\nBefore you begin, ensure your system environment is similar to the configuration below. Version matching is crucial for a successful compilation.\nThe GPU needs to support HEVC; refer to the [NVIDIA NVDEC Support Matrix](https://en.wikipedia.org/wiki/NVIDIA_Video_Coding_Engine#NVDEC).\n\n\n### Prerequisites: Setting up the Host Environment\n\nBefore building and running the Docker container, your host machine must be configured to support GPU access for Docker.\n\n1. **NVIDIA Drivers**: Ensure you have the latest NVIDIA drivers installed. You can verify this by running `nvidia-smi`.\n\n2. **Docker Engine**: Install Docker on your system. Follow the official instructions at [docs.docker.com/engine/install/](https://docs.docker.com/engine/install/).\n\n3. **NVIDIA Container Toolkit**: This toolkit allows Docker containers to access the host's NVIDIA GPU. Install it using the following commands (for Debian/Ubuntu):\n To run docker containers with GPU support you have to install the [nvidia container toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html). \n ```bash\n # Add the GPG key\n curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg\n \n # Add the repository\n curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \\\n sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \\\n sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list\n \n # Update package lists and install the toolkit\n sudo apt-get install -y \\\n nvidia-container-toolkit=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \\\n nvidia-container-toolkit-base=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \\\n libnvidia-container-tools=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \\\n libnvidia-container1=${NVIDIA_CONTAINER_TOOLKIT_VERSION}\n\n # Configure Docker to use the NVIDIA runtime\n sudo nvidia-ctk runtime configure --runtime=containerd\n \n # Restart the Docker daemon to apply the changes\n sudo systemctl restart containerd\n ```\n For other operating systems, please refer to the [official NVIDIA documentation](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).\n\n4. **Docker Image Pre-pulls [optional]**: To accelerate the build process, we provide a script to pre-pull necessary Docker images from a mirror registry.\n\n ```bash\n bash scripts/build_gpu_docker.sh\n ```\n\n### Build and Run the Container\n\nYou can also build and run the image using standard Docker commands from the root of the repository.\n\n1. **Build the GPU image**:\n ```bash\n docker build -f Dockerfile.cuda \\\n --build-arg NUM_JOBS=8 \\\n -t spatialvid-gpu .\n ```\n\n2. **Run the container**:\n ```bash\n docker run --gpus all --rm -it \\\n -v $(pwd):/workspace \\\n -w /workspace \\\n -e NVIDIA_DRIVER_CAPABILITIES=compute,video,utility \\\n spatialvid-gpu bash\n ```\n\n3. **Verify the environment (inside the container)**:\n Once inside the container, you can verify that FFmpeg and PyTorch are correctly installed and can access the GPU.\n ```bash\n # Check the custom FFmpeg build\n /usr/local/bin/ffmpeg -version\n\n # Check PyTorch and CUDA availability\n python3 -c \"import torch; print(f'PyTorch: {torch.__version__}, CUDA: {torch.version.cuda}, GPU Available: {torch.cuda.is_available()}')\"\n ```\n\n## Dataset Download\n\nOur dataset is available on [HuggingFace](https://huggingface.co/SpatialVID) and [ModelScope](https://www.modelscope.cn/organization/SpatialVID).\n\nApart from downloading the dataset using terminal commands, we provide scripts to download the SpatialVID/SpatialVID-HQ dataset from HuggingFace. Please refer to the [`download_SpatialVID.py`](utils/download_SpatialVID.py) script for more details.\n\nWe also provide our script to download the raw videos from YouTube. You can refer to the [`download_YouTube.py`](utils/download_YouTube.py) script for more details.\n\n## License\n\nPlease refer to the [LICENSE](LICENSE) file for more details about the license of our code.\n\n⚠️ SpatialVID dataset is released under the [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License](https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode) (CC-BY-NC-SA-4.0). Users must attribute the original source, use the resource only for non-commercial purposes, and release any modified/derived works under the same license. If you are the copyright owner of any video in our dataset and you need it to be removed, please contact us, and we will remove the video samples from our dataset / Github / project webpage / technical presentation as soon as possible.\n\n## References\n\nThanks to the developers and contributors of the following open-source repositories, whose invaluable work has greatly inspire our project:\n\n- [Open-Sora](https://github.com/hpcaitech/Open-Sora): An initiative dedicated to efficiently producing high-quality video.\n- [MegaSaM](https://github.com/mega-sam/mega-sam): An accurate, fast and robust casual structure and motion from casual dynamic videos.\n- [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2): A model for monocular depth estimation.\n- [UniDepthV2](https://github.com/lpiccinelli-eth/UniDepth): A model for universal monocular metric depth estimation.\n- [SAM2](https://github.com/facebookresearch/sam2): A model towards solving promptable visual segmentation in images and videos.\n- [Viser](https://viser.studio/latest/): A library for interactive 3D visualization in Python.\n\nOur repository is licensed under the Apache 2.0 License. However, if you use MegaSaM or other components in your work, please follow their license.\n\n## Citation\n\n```bibtex\n@article{wang2025spatialvid,\n title={Spatialvid: A large-scale video dataset with spatial annotations},\n author={Wang, Jiahao and Yuan, Yufeng and Zheng, Rujie and Lin, Youtian and Gao, Jian and Chen, Lin-Zhuo and Bao, Yajie and Zhang, Yi and Zeng, Chang and Zhou, Yanxi and others},\n journal={arXiv preprint arXiv:2509.09676},\n year={2025}\n}\n```\n" }, { "path": "camera_pose_annotation/.gitignore", "content": "# files\ndata/*\n*.log\n*.txt\n*.bz2\n*.zip\n*.ipynb\ndata_videos\n!requirements.txt\n!requirements_megasam.txt\n\n#python\n*.pyc\n__pycache__/\n\n# dir\noutputs/\noutputs_303/\ndata_videos/\ncheckpoints/*\n!checkpoints/megasam_final.pth\nDROID-SLAM/\n.vscode/" }, { "path": "camera_pose_annotation/README.md", "content": "# Camera Pose Annotation\n\n## Depth Estimation\nUse both [Depth-Anything V2](depth_estimation/Depth-Anything) and [UniDepth V2](depth_estimation/UniDepth) to estimate depth maps from images.\n\nDownload the pre-trained models from the respective repositories. Skip this step if you already follow the installation instructions in [README](../README.md).\n- [Depth-Anything V2](https://huggingface.co/depth-anything/Depth-Anything-V2-Large)\n- [UniDepth V2](https://huggingface.co/lpiccinelli/unidepth-v2-vitl14)\n\nTo inference depth using Depth-Anything V2, run the following command:\n\n```bash\ntorchrun --standalone --nproc_per_node ${GPU_NUM} camera_pose_annotation/depth_estimation/Depth-Anything/inference_batch.py \\\n ${CSV} \\\n --encoder vitl \\\n --checkpoints_path checkpoints \\\n --OUTPUT_DIR ${OUTPUT_DIR} \\\n --bs 16 \\\n --num_workers ${GPU_NUM}\n```\n\nTo inference depth using UniDepth V2, run the following command:\n\n```bash\ntorchrun --standalone --nproc_per_node ${GPU_NUM} camera_pose_annotation/depth_estimation/UniDepth/inference_batch.py \\\n ${CSV} \\\n --OUTPUT_DIR ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints \\\n --bs 32 \\\n --num_workers ${GPU_NUM}\n```\n\n## Camera Tracking\nUsing a DROID-SLAM based method to track camera poses from videos.\n\nTo inference a single video, run the following command:\n\n```bash\npython camera_pose_annotation/camera_tracking/camera_tracking.py \\\n --dir_path ${DIR_PATH} \\\n --weights checkpoints/megasam_final.pth \\\n --disable_vis\n```\n\nTo inference videos in batch, run the following command:\n\n```bash\npython camera_pose_annotation/camera_tracking/inference_batch.py ${CSV} \\\n --OUTPUT_DIR ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints --gpu_id ${CUDA_VISIBLE_DEVICES} \\\n --num_workers $((GPU_NUM * 2))\n```\n\n## CVD (Camera View Depth) Optimization\n### Optical Flow\nInfer optical flow using RAFT model.\n\nDownload the [`raft_things.pth`](https://drive.google.com/uc?id=1MqDajR89k-xLV0HIrmJ0k-n8ZpG6_suM).\n\nTo inference a single video, run the following command:\n\n```bash\npython camera_pose_annotation/cvd_opt/preprocess/preprocess_flow.py \\\n --dir_path ${DIR_PATH} \\\n --model checkpoints/raft-things.pth \\\n --mixed_precision\n```\n\nTo inference videos in batch, run the following command:\n\n```bash\npython camera_pose_annotation/cvd_opt/preprocess/inference_batch.py ${CSV} \\\n --OUTPUT_DIR ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints --gpu_id ${CUDA_VISIBLE_DEVICES} \\\n --num_workers $((GPU_NUM * 2))\n```\n\n### Optimization\nUsing the optical flow to optimize the estimated depth maps.\n\nTo inference a single video, run the following command:\n\n```bash\npython camera_pose_annotation/cvd_opt/cvd_opt.py \\\n --dir_path ${DIR_PATH} \\\n --w_grad 2.0 --w_normal 5.0\n```\n\nTo inference videos in batch, run the following command:\n\n```bash\npython camera_pose_annotation/cvd_opt/inference_batch.py ${CSV} \\\n --OUTPUT_DIR ${OUTPUT_DIR} \\\n --gpu_id ${CUDA_VISIBLE_DEVICES} \\\n --num_workers $((GPU_NUM * 2))\n```\n\n## Dynamic Mask\nGiven the limitations of MegaSaM in predicting motion probabilities, we opt to enhance its performance using SAM2.\n\nSpecifically, an adaptive thresholding mechanism, calibrated to the system’s motion probability distribution, is first employed to generate initial masks. Subsequently, contour detection is performed to mitigate redundant segmentation of overlapping regions; for each identified contour, four evenly spaced anchor points are sampled along its perimeter to serve as dedicated prompts for the SAM2 model.\n\nDownload the pre-trained [SAM2 model](https://huggingface.co/facebook/sam2.1-hiera-large).\n\nRun the following command:\n\n```bash\npython camera_pose_annotation/dynamic_mask/inference_batch.py ${CSV} \\\n --OUTPUT_DIR ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints --gpu_num ${GPU_NUM} \\\n --num_workers $((GPU_NUM * 2))\n```\n" }, { "path": "camera_pose_annotation/__init__.py", "content": "" }, { "path": "camera_pose_annotation/camera_tracking/__init__.py", "content": "" }, { "path": "camera_pose_annotation/camera_tracking/camera_tracking.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Test camera tracking on a single scene.\"\"\"\n\n# pylint: disable=invalid-name\n# pylint: disable=g-importing-member\n# pylint: disable=g-bad-import-order\n# pylint: disable=g-import-not-at-top\n# pylint: disable=redefined-outer-name\n# pylint: disable=undefined-variable\n# pylint: disable=undefined-loop-variable\n\nimport sys\n\nsys.path.append(\"camera_pose_annotation/base/droid_slam\")\nfrom droid import Droid\nfrom lietorch import SE3\nimport argparse\nimport glob\nimport os\nimport cv2\nimport torch\nimport numpy as np\nfrom tqdm import tqdm\nimport torch.nn.functional as F\n\n\ndef image_stream(\n image_list,\n mono_disp_list,\n scene_name,\n use_depth=False,\n aligns=None,\n K=None,\n stride=1,\n):\n \"\"\"image generator.\"\"\"\n del scene_name, stride\n\n fx, fy, cx, cy = (\n K[0, 0],\n K[1, 1],\n K[0, 2],\n K[1, 2],\n ) # np.loadtxt(os.path.join(dir_path, 'calibration.txt')).tolist()\n\n for t, (image_file) in enumerate(image_list):\n image = cv2.imread(image_file)\n # depth = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH) / 5000.\n # depth = np.float32(np.load(depth_file)) / 300.0\n # depth = 1. / pt_data[\"depth\"]\n\n mono_disp = mono_disp_list[t]\n # mono_disp = np.float32(np.load(disp_file)) #/ 300.0\n depth = np.clip(\n 1.0 / ((1.0 / aligns[2]) * (aligns[0] * mono_disp + aligns[1])),\n 1e-4,\n 1e4,\n )\n depth[depth < 1e-2] = 0.0\n\n # breakpoint()\n h0, w0, _ = image.shape\n h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))\n w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))\n\n image = cv2.resize(image, (w1, h1), interpolation=cv2.INTER_AREA)\n image = image[: h1 - h1 % 8, : w1 - w1 % 8]\n\n image = torch.as_tensor(image).permute(2, 0, 1)\n\n depth = torch.as_tensor(depth)\n depth = F.interpolate(\n depth[None, None], (h1, w1), mode=\"nearest-exact\"\n ).squeeze()\n depth = depth[: h1 - h1 % 8, : w1 - w1 % 8]\n\n mask = torch.ones_like(depth)\n\n intrinsics = torch.as_tensor([fx, fy, cx, cy])\n intrinsics[0::2] *= w1 / w0\n intrinsics[1::2] *= h1 / h0\n\n if use_depth:\n yield t, image[None], depth, intrinsics, mask\n else:\n yield t, image[None], intrinsics, mask\n\n\ndef save_full_reconstruction(\n droid, full_traj, rgb_list, senor_depth_list, motion_prob, scene_name, save_path\n):\n \"\"\"Save full reconstruction.\"\"\"\n from pathlib import Path\n\n t = full_traj.shape[0]\n images = np.array(rgb_list[:t]) # droid.video.images[:t].cpu().numpy()\n disps = 1.0 / (np.array(senor_depth_list[:t]) + 1e-6)\n\n poses = full_traj # .cpu().numpy()\n intrinsics = droid.video.intrinsics[:t].cpu().numpy()\n Path(f\"{save_path}\").mkdir(parents=True, exist_ok=True)\n np.save(f\"{save_path}/images.npy\", images)\n np.save(f\"{save_path}/disps.npy\", disps)\n np.save(f\"{save_path}/poses.npy\", poses)\n np.save(f\"{save_path}/intrinsics.npy\", intrinsics * 8.0)\n np.save(f\"{save_path}/motion_prob.npy\", motion_prob)\n\n intrinsics = intrinsics[0] * 8.0\n poses_th = torch.as_tensor(poses, device=\"cpu\")\n cam_c2w = SE3(poses_th).inv().matrix().numpy()\n\n K = np.eye(3)\n K[0, 0] = intrinsics[0]\n K[1, 1] = intrinsics[1]\n K[0, 2] = intrinsics[2]\n K[1, 2] = intrinsics[3]\n\n max_frames = min(1000, images.shape[0])\n if not os.path.exists(save_path):\n os.makedirs(save_path)\n\n np.savez(\n os.path.join(save_path, f\"{scene_name}_droid.npz\"),\n images=np.uint8(images[:max_frames, ::-1, ...].transpose(0, 2, 3, 1)),\n depths=np.float32(1.0 / disps[:max_frames, ...]),\n intrinsic=K,\n cam_c2w=cam_c2w[:max_frames],\n )\n\n\ndef parse_args():\n \"\"\"Parse command line arguments.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--dir_path\", help=\"path to the dataset\")\n parser.add_argument(\"--weights\", default=\"droid.pth\")\n parser.add_argument(\"--buffer\", type=int, default=1024)\n parser.add_argument(\"--image_size\", default=[240, 320])\n parser.add_argument(\"--disable_vis\", action=\"store_true\")\n\n parser.add_argument(\"--beta\", type=float, default=0.3)\n parser.add_argument(\n \"--filter_thresh\", type=float, default=2.0\n ) # motion threhold for keyframe\n parser.add_argument(\"--warmup\", type=int, default=8)\n parser.add_argument(\"--keyframe_thresh\", type=float, default=2.0)\n parser.add_argument(\"--frontend_thresh\", type=float, default=12.0)\n parser.add_argument(\"--frontend_window\", type=int, default=25)\n parser.add_argument(\"--frontend_radius\", type=int, default=2)\n parser.add_argument(\"--frontend_nms\", type=int, default=1)\n\n parser.add_argument(\"--stereo\", action=\"store_true\")\n parser.add_argument(\"--depth\", action=\"store_true\")\n parser.add_argument(\"--upsample\", action=\"store_true\")\n parser.add_argument(\"--scene_name\", help=\"scene_name\")\n\n parser.add_argument(\"--backend_thresh\", type=float, default=16.0)\n parser.add_argument(\"--backend_radius\", type=int, default=2)\n parser.add_argument(\"--backend_nms\", type=int, default=3)\n\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n\n scene_name = os.path.basename(args.dir_path)\n\n rgb_list = []\n senor_depth_list = []\n img_path = os.path.join(args.dir_path, \"img\")\n\n img_list = sorted(glob.glob(os.path.join(img_path, \"*.jpg\")))\n img_list += sorted(glob.glob(os.path.join(img_path, \"*.png\")))\n\n # NOTE Mono is inverse depth, but metric-depth is depth!\n mono_disp_paths = sorted(\n glob.glob(os.path.join(args.dir_path, \"depth-anything\", \"*.npy\"))\n )\n metric_depth_paths = sorted(\n glob.glob(os.path.join(args.dir_path, \"unidepth\", \"*.npz\"))\n )\n\n img_0 = cv2.imread(img_list[0])\n scales = []\n shifts = []\n mono_disp_list = []\n fovs = []\n for t, (mono_disp_file, metric_depth_file) in enumerate(\n zip(mono_disp_paths, metric_depth_paths)\n ):\n da_disp = np.float32(np.load(mono_disp_file)) # / 300.0\n uni_data = np.load(metric_depth_file)\n metric_depth = uni_data[\"depth\"]\n\n fovs.append(uni_data[\"fov\"])\n\n da_disp = cv2.resize(\n da_disp,\n (metric_depth.shape[1], metric_depth.shape[0]),\n interpolation=cv2.INTER_NEAREST_EXACT,\n )\n mono_disp_list.append(da_disp)\n gt_disp = 1.0 / (metric_depth + 1e-8)\n\n # avoid some bug from UniDepth\n valid_mask = (metric_depth < 2.0) & (da_disp < 0.02)\n gt_disp[valid_mask] = 1e-2\n\n # avoid cases sky dominate entire video\n sky_ratio = np.sum(da_disp < 0.01) / (da_disp.shape[0] * da_disp.shape[1])\n if sky_ratio > 0.5:\n non_sky_mask = da_disp > 0.01\n gt_disp_ms = gt_disp[non_sky_mask] - np.median(gt_disp[non_sky_mask]) + 1e-8\n da_disp_ms = da_disp[non_sky_mask] - np.median(da_disp[non_sky_mask]) + 1e-8\n scale = np.median(gt_disp_ms / da_disp_ms)\n shift = np.median(gt_disp[non_sky_mask] - scale * da_disp[non_sky_mask])\n else:\n gt_disp_ms = gt_disp - np.median(gt_disp) + 1e-8\n da_disp_ms = da_disp - np.median(da_disp) + 1e-8\n scale = np.median(gt_disp_ms / da_disp_ms)\n shift = np.median(gt_disp - scale * da_disp)\n\n gt_disp_ms = gt_disp - np.median(gt_disp) + 1e-8\n da_disp_ms = da_disp - np.median(da_disp) + 1e-8\n\n scale = np.median(gt_disp_ms / da_disp_ms)\n shift = np.median(gt_disp - scale * da_disp)\n\n scales.append(scale)\n shifts.append(shift)\n\n print(\"************** UNIDEPTH FOV \", np.median(fovs))\n ff = img_0.shape[1] / (2 * np.tan(np.radians(np.median(fovs) / 2.0)))\n K = np.eye(3)\n K[0, 0] = ff * 1.0 # pp_intrinsic[0] * (img_0.shape[1] / (pp_intrinsic[1] * 2))\n K[1, 1] = ff * 1.0 # pp_intrinsic[0] * (img_0.shape[0] / (pp_intrinsic[2] * 2))\n K[0, 2] = (\n img_0.shape[1] / 2.0\n ) # pp_intrinsic[1]) * (img_0.shape[1] / (pp_intrinsic[1] * 2))\n K[1, 2] = (\n img_0.shape[0] / 2.0\n ) # (pp_intrinsic[2]) * (img_0.shape[0] / (pp_intrinsic[2] * 2))\n\n ss_product = np.array(scales) * np.array(shifts)\n med_idx = np.argmin(np.abs(ss_product - np.median(ss_product)))\n\n align_scale = scales[med_idx] # np.median(np.array(scales))\n align_shift = shifts[med_idx] # np.median(np.array(shifts))\n normalize_scale = (\n np.percentile((align_scale * np.array(mono_disp_list) + align_shift), 98) / 2.0\n )\n\n aligns = (align_scale, align_shift, normalize_scale)\n\n for t, image, depth, intrinsics, mask in tqdm(\n image_stream(\n img_list,\n mono_disp_list,\n scene_name,\n use_depth=True,\n aligns=aligns,\n K=K,\n )\n ):\n rgb_list.append(image[0])\n senor_depth_list.append(depth)\n # breakpoint()\n if t == 0:\n args.image_size = [image.shape[2], image.shape[3]]\n droid = Droid(args, device=0)\n\n droid.track(t, image, depth, intrinsics=intrinsics, mask=mask)\n\n # last frame\n droid.track_final(t, image, depth, intrinsics=intrinsics, mask=mask)\n\n traj_est, depth_est, motion_prob = droid.terminate(\n image_stream(\n img_list,\n mono_disp_list,\n scene_name,\n use_depth=True,\n aligns=aligns,\n K=K,\n ),\n _opt_intr=True, # default is opt_focal\n full_ba=True,\n scene_name=scene_name,\n )\n\n save_full_reconstruction(\n droid,\n traj_est,\n rgb_list,\n senor_depth_list,\n motion_prob,\n args.scene_name,\n os.path.join(args.dir_path, \"reconstructions\"),\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/camera_tracking/inference_batch.py", "content": "\"\"\"\nBatch inference for camera tracking using multiple GPUs.\n\nThis module provides functionality for:\n- Parallel camera tracking processing across multiple videos\n- Multi-GPU support with automatic device assignment\n- Subprocess management for camera tracking pipeline\n- Progress tracking and error handling\n\"\"\"\n\nimport pandas as pd\nimport os\nimport argparse\nimport concurrent.futures\nfrom multiprocessing import Manager\nimport subprocess\nimport queue\nfrom tqdm import tqdm\n\n\ndef process_single_row(row, index, args, worker_id=0):\n \"\"\"\n Process a single video for camera tracking.\n \"\"\"\n dir_path = os.path.join(args.dir_path, row[\"id\"])\n device_id = worker_id % args.gpu_num\n\n cmd = (\n f\"CUDA_VISIBLE_DEVICES={args.gpu_id[device_id]} python camera_pose_annotation/camera_tracking/camera_tracking.py \"\n f\"--dir_path {dir_path} \"\n f\"--weights {args.checkpoints_path}/megasam_final.pth \"\n f\"--disable_vis\"\n )\n process = subprocess.Popen(\n cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE\n )\n stdout, stderr = process.communicate()\n if process.returncode != 0:\n print(f\"Error tracking camera for {row['id']}: {stderr.decode()}\")\n\n\ndef worker(task_queue, args, worker_id, pbar):\n \"\"\"\n Worker function for parallel camera tracking processing.\n \"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n process_single_row(row, index, args, worker_id)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for camera tracking batch inference.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\"--dir_path\", type=str, default=\"./outputs\")\n parser.add_argument(\"--checkpoints_path\", type=str, default=\"./checkpoints\")\n parser.add_argument(\n \"--gpu_id\", type=str, default=\"0\", help=\"Comma-separated list of GPU IDs to use\"\n )\n parser.add_argument(\n \"--num_workers\",\n type=int,\n default=4,\n help=\"Number of workers for parallel processing\",\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n\n # Parse GPU configuration\n args.gpu_num = len(args.gpu_id.split(\",\"))\n args.gpu_id = [int(gpu) for gpu in args.gpu_id.split(\",\")]\n\n df = pd.read_csv(args.csv_path)\n\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(df.iterrows(), total=len(df), desc=\"Processing rows\"):\n process_single_row(row, index, args)\n else:\n # Parallel processing with multiple workers\n manager = Manager()\n task_queue = manager.Queue()\n\n # Add all tasks to queue\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n with tqdm(total=len(df), desc=\"Processing rows\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for id in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, args, id, pbar))\n\n for future in concurrent.futures.as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/cvd_opt/__init__.py", "content": "" }, { "path": "camera_pose_annotation/cvd_opt/cvd_opt.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Consistent video depth optimization.\"\"\"\n\n# pylint: disable=invalid-name\n# pylint: disable=g-importing-member\n# pylint: disable=redefined-outer-name\n\nimport argparse\nimport os\nfrom pathlib import Path\nimport pandas as pd\n\nfrom geometry_utils import NormalGenerator\nimport kornia\nfrom lietorch import SE3\nimport numpy as np\nimport torch\n\nimport zipfile\nimport tempfile\nimport OpenEXR\nimport Imath\n\n\ndef save_depth(path, depths):\n with zipfile.ZipFile(path, \"w\", zipfile.ZIP_DEFLATED) as z:\n for index, depth in enumerate(depths):\n height, width = depth.shape\n header = OpenEXR.Header(width, height)\n header[\"channels\"] = {\"Z\": Imath.Channel(Imath.PixelType(Imath.PixelType.HALF))}\n with tempfile.NamedTemporaryFile(suffix=\".exr\") as f:\n exr = OpenEXR.OutputFile(f.name, header)\n exr.writePixels({\"Z\": depth.astype(np.float16).tobytes()})\n exr.close()\n z.write(f.name, f\"{index:05d}.exr\")\n\n\ndef gradient_loss(gt, pred, u):\n \"\"\"Gradient loss.\"\"\"\n del u\n diff = pred - gt\n v_gradient = torch.abs(diff[..., 0:-2, 1:-1] - diff[..., 2:, 1:-1]) # * mask_v\n h_gradient = torch.abs(diff[..., 1:-1, 0:-2] - diff[..., 1:-1, 2:]) # * mask_h\n\n pred_grad = torch.abs(pred[..., 0:-2, 1:-1] - (pred[..., 2:, 1:-1])) + torch.abs(\n pred[..., 1:-1, 0:-2] - pred[..., 1:-1, 2:]\n )\n gt_grad = torch.abs(gt[..., 0:-2, 1:-1] - (gt[..., 2:, 1:-1])) + torch.abs(\n gt[..., 1:-1, 0:-2] - gt[..., 1:-1, 2:]\n )\n\n grad_diff = torch.abs(pred_grad - gt_grad)\n nearby_mask = (torch.exp(gt[..., 1:-1, 1:-1]) > 1.0).float().detach()\n # weight = (1. - torch.exp(-(grad_diff * 5.)).detach())\n weight = 1.0 - torch.exp(-(grad_diff * 5.0)).detach()\n weight *= nearby_mask\n\n g_loss = torch.mean(h_gradient * weight) + torch.mean(v_gradient * weight)\n return g_loss\n\n\ndef si_loss(gt, pred):\n log_gt = torch.log(torch.clamp(gt, 1e-3, 1e3)).view(gt.shape[0], -1)\n log_pred = torch.log(torch.clamp(pred, 1e-3, 1e3)).view(pred.shape[0], -1)\n log_diff = log_gt - log_pred\n num_pixels = gt.shape[-2] * gt.shape[-1]\n data_loss = torch.sum(log_diff**2, dim=-1) / num_pixels - torch.sum(\n log_diff, dim=-1\n ) ** 2 / (num_pixels**2)\n return torch.mean(data_loss)\n\n\ndef sobel_fg_alpha(disp, mode=\"sobel\", beta=10.0):\n sobel_grad = kornia.filters.spatial_gradient(disp, mode=mode, normalized=False)\n sobel_mag = torch.sqrt(\n sobel_grad[:, :, 0, Ellipsis] ** 2 + sobel_grad[:, :, 1, Ellipsis] ** 2\n )\n alpha = torch.exp(-1.0 * beta * sobel_mag).detach()\n\n return alpha\n\n\nALPHA_MOTION = 0.25\nRESIZE_FACTOR = 0.5\n\n\ndef consistency_loss(\n cam_c2w,\n K,\n K_inv,\n disp_data,\n init_disp,\n uncertainty,\n flows,\n flow_masks,\n ii,\n jj,\n compute_normals,\n fg_alpha,\n w_ratio=1.0,\n w_flow=0.2,\n w_si=1.0,\n w_grad=2.0,\n w_normal=4.0,\n):\n \"\"\"Consistency loss.\"\"\"\n _, H, W = disp_data.shape\n # mesh grid\n xx = torch.arange(0, W).view(1, -1).repeat(H, 1)\n yy = torch.arange(0, H).view(-1, 1).repeat(1, W)\n xx = xx.view(1, 1, H, W) # .repeat(B ,1 ,1 ,1)\n yy = yy.view(1, 1, H, W) # .repeat(B ,1 ,1 ,1)\n grid = torch.cat((xx, yy), 1).float().cuda().permute(0, 2, 3, 1) # [None, ...]\n\n loss_flow = 0.0 # flow reprojection loss\n loss_d_ratio = 0.0 # depth consistency loss\n\n flows_step = flows.permute(0, 2, 3, 1)\n flow_masks_step = flow_masks.permute(0, 2, 3, 1).squeeze(-1)\n\n cam_1to2 = torch.bmm(\n torch.linalg.inv(torch.index_select(cam_c2w, dim=0, index=jj)),\n torch.index_select(cam_c2w, dim=0, index=ii),\n )\n\n # warp disp from target time\n pixel_locations = grid + flows_step\n resize_factor = torch.tensor([W - 1.0, H - 1.0]).cuda()[None, None, None, ...]\n normalized_pixel_locations = 2 * (pixel_locations / resize_factor) - 1.0\n\n disp_sampled = torch.nn.functional.grid_sample(\n torch.index_select(disp_data, dim=0, index=jj)[:, None, ...],\n normalized_pixel_locations,\n align_corners=True,\n )\n\n uu = torch.index_select(uncertainty, dim=0, index=ii).squeeze(1)\n\n grid_h = torch.cat([grid, torch.ones_like(grid[..., 0:1])], dim=-1).unsqueeze(-1)\n # depth of reference view\n ref_depth = 1.0 / torch.clamp(\n torch.index_select(disp_data, dim=0, index=ii), 1e-3, 1e3\n )\n\n pts_3d_ref = ref_depth[..., None, None] * (K_inv[None, None, None] @ grid_h)\n rot = cam_1to2[:, None, None, :3, :3]\n trans = cam_1to2[:, None, None, :3, 3:4]\n\n pts_3d_tgt = (rot @ pts_3d_ref) + trans # [:, None, None, :, None]\n depth_tgt = pts_3d_tgt[:, :, :, 2:3, 0]\n disp_tgt = 1.0 / torch.clamp(depth_tgt, 0.1, 1e3)\n\n # flow consistency loss\n pts_2D_tgt = K[None, None, None] @ pts_3d_tgt\n\n flow_masks_step_ = flow_masks_step * (pts_2D_tgt[:, :, :, 2, 0] > 0.1)\n pts_2D_tgt = pts_2D_tgt[:, :, :, :2, 0] / torch.clamp(\n pts_2D_tgt[:, :, :, 2:, 0], 1e-3, 1e3\n )\n\n disp_sampled = torch.clamp(disp_sampled, 1e-3, 1e2)\n disp_tgt = torch.clamp(disp_tgt, 1e-3, 1e2)\n\n ratio = torch.maximum(\n disp_sampled.squeeze() / disp_tgt.squeeze(),\n disp_tgt.squeeze() / disp_sampled.squeeze(),\n )\n ratio_error = torch.abs(ratio - 1.0) #\n\n loss_d_ratio += torch.sum(\n (ratio_error * uu + ALPHA_MOTION * torch.log(1.0 / uu)) * flow_masks_step_\n ) / (torch.sum(flow_masks_step_) + 1e-8)\n\n flow_error = torch.abs(pts_2D_tgt - pixel_locations)\n loss_flow += torch.sum(\n (flow_error * uu[..., None] + ALPHA_MOTION * torch.log(1.0 / uu[..., None]))\n * flow_masks_step_[..., None]\n ) / (torch.sum(flow_masks_step_) * 2.0 + 1e-8)\n\n # prior mono-depth reg loss\n loss_prior = si_loss(init_disp, disp_data)\n KK = torch.inverse(K_inv)\n\n # multi gradient consistency\n disp_data_ds = disp_data[:, None, ...]\n init_disp_ds = init_disp[:, None, ...]\n K_rescale = KK.clone()\n K_inv_rescale = torch.inverse(K_rescale)\n pred_normal = compute_normals[0](\n 1.0 / torch.clamp(disp_data_ds, 1e-3, 1e3), K_inv_rescale[None]\n )\n init_normal = compute_normals[0](\n 1.0 / torch.clamp(init_disp_ds, 1e-3, 1e3), K_inv_rescale[None]\n )\n\n loss_normal = torch.mean(\n fg_alpha * (1.0 - torch.sum(pred_normal * init_normal, dim=1))\n ) # / (1e-8 + torch.sum(fg_alpha))\n\n loss_grad = 0.0\n for scale in range(4):\n interval = 2**scale\n disp_data_ds = torch.nn.functional.interpolate(\n disp_data[:, None, ...],\n scale_factor=(1.0 / interval, 1.0 / interval),\n mode=\"nearest-exact\",\n )\n init_disp_ds = torch.nn.functional.interpolate(\n init_disp[:, None, ...],\n scale_factor=(1.0 / interval, 1.0 / interval),\n mode=\"nearest-exact\",\n )\n uncertainty_rs = torch.nn.functional.interpolate(\n uncertainty,\n scale_factor=(1.0 / interval, 1.0 / interval),\n mode=\"nearest-exact\",\n )\n loss_grad += gradient_loss(\n torch.log(disp_data_ds), torch.log(init_disp_ds), uncertainty_rs\n )\n\n return (\n w_ratio * loss_d_ratio\n + w_si * loss_prior\n + w_flow * loss_flow\n + w_normal * loss_normal\n + loss_grad * w_grad\n )\n\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--w_grad\", type=float, default=2.0, help=\"w_grad\")\n parser.add_argument(\"--w_normal\", type=float, default=6.0, help=\"w_normal\")\n parser.add_argument(\"--dir_path\", type=str, default=\".\", help=\"directory path\")\n parser.add_argument(\"--only_depth\", action=\"store_true\", help=\"only save optimize depth\")\n\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n\n scene_name = os.path.basename(args.dir_path)\n\n cache_dir = os.path.join(args.dir_path, \"cache-flow\")\n rootdir = os.path.join(args.dir_path, \"reconstructions\")\n\n print(\"***************************** \", scene_name)\n img_data = np.load(os.path.join(rootdir, \"images.npy\"))[:, ::-1, ...]\n disp_data = np.load(os.path.join(rootdir, \"disps.npy\")) + 1e-6\n intrinsics = np.load(os.path.join(rootdir, \"intrinsics.npy\"))\n poses = np.load(os.path.join(rootdir, \"poses.npy\"))\n mot_prob = np.load(os.path.join(rootdir, \"motion_prob.npy\"))\n\n flows = np.load(os.path.join(cache_dir, \"flows.npy\"), allow_pickle=True)\n flow_masks = np.load(os.path.join(cache_dir, \"flows_masks.npy\"), allow_pickle=True)\n flow_masks = np.float32(flow_masks)\n iijj = np.load(os.path.join(cache_dir, \"ii-jj.npy\"), allow_pickle=True)\n\n intrinsics = intrinsics[0]\n poses_th = torch.as_tensor(poses, device=\"cpu\").float().cuda()\n\n K = np.eye(3)\n K[0, 0] = intrinsics[0]\n K[1, 1] = intrinsics[1]\n K[0, 2] = intrinsics[2]\n K[1, 2] = intrinsics[3]\n\n img_data_pt = (\n torch.from_numpy(np.ascontiguousarray(img_data)).float().cuda() / 255.0\n )\n flows = torch.from_numpy(np.ascontiguousarray(flows)).float().cuda()\n flow_masks = (\n torch.from_numpy(np.ascontiguousarray(flow_masks)).float().cuda()\n ) # .unsqueeze(1)\n iijj = torch.from_numpy(np.ascontiguousarray(iijj)).float().cuda()\n ii = iijj[0, ...].long()\n jj = iijj[1, ...].long()\n K = torch.from_numpy(K).float().cuda()\n\n init_disp = torch.from_numpy(disp_data).float().cuda()\n disp_data = torch.from_numpy(disp_data).float().cuda()\n\n assert init_disp.shape == disp_data.shape\n\n init_disp = torch.nn.functional.interpolate(\n init_disp.unsqueeze(1),\n scale_factor=(RESIZE_FACTOR, RESIZE_FACTOR),\n mode=\"bilinear\",\n ).squeeze(1)\n disp_data = torch.nn.functional.interpolate(\n disp_data.unsqueeze(1),\n scale_factor=(RESIZE_FACTOR, RESIZE_FACTOR),\n mode=\"bilinear\",\n ).squeeze(1)\n\n fg_alpha = sobel_fg_alpha(init_disp[:, None, ...]) > 0.2\n fg_alpha = fg_alpha.squeeze(1).float() + 0.2\n\n cvd_prob = torch.nn.functional.interpolate(\n torch.from_numpy(mot_prob).unsqueeze(1).cuda(),\n scale_factor=(4, 4),\n mode=\"bilinear\",\n )\n cvd_prob[cvd_prob > 0.5] = 0.5\n cvd_prob = torch.clamp(cvd_prob, 1e-3, 1.0)\n\n # rescale intrinsic matrix to small resolution\n K_o = K.clone()\n K[0:2, ...] *= RESIZE_FACTOR\n K_inv = torch.linalg.inv(K)\n\n disp_data.requires_grad = False\n poses_th.requires_grad = False\n\n uncertainty = cvd_prob\n\n # First optimize scale and shift to align them\n log_scale_ = torch.log(torch.ones(init_disp.shape[0]).to(disp_data.device))\n shift_ = torch.zeros(init_disp.shape[0]).to(disp_data.device)\n log_scale_.requires_grad = True\n shift_.requires_grad = True\n uncertainty.requires_grad = True\n\n optim = torch.optim.Adam(\n [\n {\"params\": log_scale_, \"lr\": 1e-2},\n {\"params\": shift_, \"lr\": 1e-2},\n {\"params\": uncertainty, \"lr\": 1e-2},\n ]\n )\n\n compute_normals = []\n compute_normals.append(NormalGenerator(disp_data.shape[-2], disp_data.shape[-1]))\n init_disp = torch.clamp(init_disp, 1e-3, 1e3)\n\n for i in range(100):\n optim.zero_grad()\n cam_c2w = SE3(poses_th).inv().matrix()\n scale_ = torch.exp(log_scale_)\n\n loss = consistency_loss(\n cam_c2w,\n K,\n K_inv,\n torch.clamp(\n disp_data * scale_[..., None, None] + shift_[..., None, None],\n 1e-3,\n 1e3,\n ),\n init_disp,\n torch.clamp(uncertainty, 1e-4, 1e3),\n flows,\n flow_masks,\n ii,\n jj,\n compute_normals,\n fg_alpha,\n )\n\n loss.backward()\n uncertainty.grad = torch.nan_to_num(uncertainty.grad, nan=0.0)\n log_scale_.grad = torch.nan_to_num(log_scale_.grad, nan=0.0)\n shift_.grad = torch.nan_to_num(shift_.grad, nan=0.0)\n\n optim.step()\n print(\"step \", i, loss.item())\n\n # Then optimize depth and uncertainty\n disp_data = (\n disp_data * torch.exp(log_scale_)[..., None, None].detach()\n + shift_[..., None, None].detach()\n )\n init_disp = (\n init_disp * torch.exp(log_scale_)[..., None, None].detach()\n + shift_[..., None, None].detach()\n )\n init_disp = torch.clamp(init_disp, 1e-3, 1e3)\n\n disp_data.requires_grad = True\n uncertainty.requires_grad = True\n poses_th.requires_grad = False # True\n\n optim = torch.optim.Adam(\n [\n {\"params\": disp_data, \"lr\": 5e-3},\n {\"params\": uncertainty, \"lr\": 5e-3},\n ]\n )\n\n losses = []\n for i in range(400):\n optim.zero_grad()\n cam_c2w = SE3(poses_th).inv().matrix()\n loss = consistency_loss(\n cam_c2w,\n K,\n K_inv,\n torch.clamp(disp_data, 1e-3, 1e3),\n init_disp,\n torch.clamp(uncertainty, 1e-4, 1e3),\n flows,\n flow_masks,\n ii,\n jj,\n compute_normals,\n fg_alpha,\n w_ratio=1.0,\n w_flow=0.2,\n w_si=1,\n w_grad=args.w_grad,\n w_normal=args.w_normal,\n )\n\n loss.backward()\n disp_data.grad = torch.nan_to_num(disp_data.grad, nan=0.0)\n uncertainty.grad = torch.nan_to_num(uncertainty.grad, nan=0.0)\n\n optim.step()\n print(\"step \", i, loss.item())\n losses.append(loss)\n\n disp_data_opt = (\n torch.nn.functional.interpolate(\n disp_data.unsqueeze(1), scale_factor=(2, 2), mode=\"bilinear\"\n )\n .squeeze(1)\n .detach()\n .cpu()\n .numpy()\n )\n\n if args.only_depth:\n save_depth(\n os.path.join(args.dir_path, \"depth_opt.zip\"),\n disp_data_opt\n )\n else:\n np.savez(\n os.path.join(args.dir_path, \"sgd_cvd_hr.npz\"),\n images=np.uint8(img_data_pt.cpu().numpy().transpose(0, 2, 3, 1) * 255.0),\n depths=np.clip(np.float16(1.0 / disp_data_opt), 1e-3, 1e2),\n intrinsic=K_o.detach().cpu().numpy(),\n cam_c2w=cam_c2w.detach().cpu().numpy(),\n )\n" }, { "path": "camera_pose_annotation/cvd_opt/geometry_utils.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Geometry utils for MegaSaM.\"\"\"\n\n# pylint: disable=invalid-name\n\nimport kornia\nimport numpy as np\nimport torch\nfrom torch import jit\nfrom torch import nn\nfrom torch import Tensor # pylint: disable=g-importing-member\nimport torch.nn.functional as F\n\n\n@torch.jit.script\ndef to_homogeneous(input_tensor: Tensor, dim: int = 0) -> Tensor:\n \"\"\"Converts tensor to homogeneous coordinates by adding ones to the specified dimension.\"\"\"\n ones = torch.ones_like(input_tensor.select(dim, 0).unsqueeze(dim))\n output_bkn = torch.cat([input_tensor, ones], dim=dim)\n return output_bkn\n\n\nclass BackprojectDepth(nn.Module):\n \"\"\"Layer that projects points from 2D camera to 3D space.\n\n The 3D points are represented in homogeneous coordinates.\n \"\"\"\n\n def __init__(self, height: int, width: int):\n super().__init__()\n\n self.height = height\n self.width = width\n\n xx, yy = torch.meshgrid(\n torch.arange(self.width),\n torch.arange(self.height),\n indexing=\"xy\",\n )\n pix_coords_2hw = torch.stack((xx, yy), axis=0) + 0.5\n\n pix_coords_13N = (\n to_homogeneous(\n pix_coords_2hw,\n dim=0,\n )\n .flatten(1)\n .unsqueeze(0)\n )\n\n # make these tensors into buffers so they are put on the correct GPU\n # automatically\n self.register_buffer(\"pix_coords_13N\", pix_coords_13N)\n\n # @jit.script_method\n def forward(self, depth_b1hw: Tensor, invK_b44: Tensor) -> Tensor:\n \"\"\"Backprojects spatial points in 2D image space to world space using invK_b44 at the depths defined in depth_b1hw.\"\"\"\n cam_points_b3N = torch.matmul(\n invK_b44[:, :3, :3], self.pix_coords_13N.float().cuda()\n )\n cam_points_b3N = depth_b1hw.flatten(start_dim=2) * cam_points_b3N\n cam_points_b4N = to_homogeneous(cam_points_b3N, dim=1)\n return cam_points_b4N\n\n\nclass Project3D(jit.ScriptModule):\n \"\"\"Layer that projects 3D points into the 2D camera.\"\"\"\n\n def __init__(self, eps: float = 1e-8):\n super().__init__()\n\n self.register_buffer(\"eps\", torch.tensor(eps).view(1, 1, 1))\n\n @jit.script_method\n def forward(\n self, points_b4N: Tensor, K_b44: Tensor, cam_T_world_b44: Tensor\n ) -> Tensor:\n \"\"\"Projects spatial points in 3D world space to camera image space using the extrinsics matrix cam_T_world_b44 and intrinsics K_b44.\"\"\"\n P_b44 = K_b44 @ cam_T_world_b44\n\n cam_points_b3N = P_b44[:, :3] @ points_b4N\n\n # from Kornia and OpenCV:\n # https://kornia.readthedocs.io/en/latest/_modules/kornia/geometry/conversions.html#convert_points_from_homogeneous\n mask = torch.abs(cam_points_b3N[:, 2:]) > self.eps\n depth_b1N = cam_points_b3N[:, 2:] + self.eps\n scale = torch.where(\n mask, 1.0 / depth_b1N, torch.tensor(1.0, device=depth_b1N.device)\n )\n\n pix_coords_b2N = cam_points_b3N[:, :2] * scale\n\n return torch.cat([pix_coords_b2N, depth_b1N], dim=1)\n\n\nclass NormalGenerator(nn.Module):\n \"\"\"Estimates normals from depth maps.\"\"\"\n\n def __init__(\n self,\n height: int,\n width: int,\n smoothing_kernel_size: int = 5,\n smoothing_kernel_std: float = 2.0,\n ):\n \"\"\"Estimates normals from depth maps.\"\"\"\n super().__init__()\n self.height = height\n self.width = width\n\n self.backproject = BackprojectDepth(self.height, self.width)\n\n self.kernel_size = smoothing_kernel_size\n self.std = smoothing_kernel_std\n\n # @jit.script_method\n def forward(self, depth_b1hw: Tensor, invK_b44: Tensor) -> Tensor:\n \"\"\"Estimates a normal at each location in the depth map.\"\"\"\n\n # First smoothes incoming depth maps with a gaussian blur, backprojects\n # those depth points into world space (see BackprojectDepth), estimates\n # the spatial gradient at those points, and finally uses normalized cross\n # correlation to estimate a normal vector at each location.\n depth_smooth_b1hw = kornia.filters.gaussian_blur2d(\n depth_b1hw,\n (self.kernel_size, self.kernel_size),\n (self.std, self.std),\n )\n cam_points_b4N = self.backproject(depth_smooth_b1hw, invK_b44)\n cam_points_b3hw = cam_points_b4N[:, :3].view(-1, 3, self.height, self.width)\n\n gradients_b32hw = kornia.filters.spatial_gradient(cam_points_b3hw)\n\n return F.normalize(\n torch.cross(\n gradients_b32hw[:, :, 0],\n gradients_b32hw[:, :, 1],\n dim=1,\n ),\n dim=1,\n )\n\n\ndef get_camera_rays(\n world_T_cam_b44,\n world_points_b3N,\n in_camera_frame,\n cam_T_world_b44=None,\n eps=1e-4,\n):\n \"\"\"Computes camera rays for given camera data and points, optionally shifts rays to camera frame.\"\"\"\n del eps\n if in_camera_frame:\n batch_size = world_points_b3N.shape[0]\n num_points = world_points_b3N.shape[2]\n world_points_b4N = torch.cat(\n [\n world_points_b3N,\n torch.ones(batch_size, 1, num_points).to(world_points_b3N.device),\n ],\n 1,\n )\n camera_points_b3N = torch.matmul(\n cam_T_world_b44[:, :3, :4], world_points_b4N\n )\n rays_b3N = camera_points_b3N\n else:\n rays_b3N = world_points_b3N - world_T_cam_b44[:, 0:3, 3][:, :, None].expand(\n world_points_b3N.shape\n )\n\n rays_b3N = torch.nn.functional.normalize(rays_b3N, dim=1)\n\n return rays_b3N\n\n\ndef pose_distance(pose_b44):\n \"\"\"DVMVS frame pose distance.\"\"\"\n\n R = pose_b44[:, :3, :3]\n t = pose_b44[:, :3, 3]\n R_trace = R.diagonal(offset=0, dim1=-1, dim2=-2).sum(-1)\n R_measure = torch.sqrt(\n 2 * (1 - torch.minimum(torch.ones_like(R_trace) * 3.0, R_trace) / 3)\n )\n t_measure = torch.norm(t, dim=1)\n combined_measure = torch.sqrt(t_measure**2 + R_measure**2)\n\n return combined_measure, R_measure, t_measure\n\n\ndef qvec2rotmat(qvec):\n \"\"\"Quaternion to 3x3 rotation matrix.\"\"\"\n return np.array([\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\ndef rotx(t):\n \"\"\"3D Rotation about the x-axis.\"\"\"\n c = np.cos(t)\n s = np.sin(t)\n return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])\n\n\ndef roty(t):\n \"\"\"3D Rotation about the y-axis.\"\"\"\n c = np.cos(t)\n s = np.sin(t)\n return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])\n\n\ndef rotz(t):\n \"\"\"3D Rotation about the z-axis.\"\"\"\n c = np.cos(t)\n s = np.sin(t)\n return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])" }, { "path": "camera_pose_annotation/cvd_opt/inference_batch.py", "content": "\"\"\"\nBatch inference script for CVD (Camera View Depth) optimization.\nProcesses multiple video clips in parallel using multi-GPU setup.\n\"\"\"\n\nimport pandas as pd\nimport os\nimport argparse\nimport concurrent.futures\nfrom multiprocessing import Manager\nimport subprocess\nimport queue\nfrom tqdm import tqdm\n\n\ndef process_single_row(row, index, args, worker_id=0):\n \"\"\"Process a single video clip for CVD optimization.\"\"\"\n dir_path = os.path.join(args.dir_path, row[\"id\"])\n device_id = worker_id % args.gpu_num\n\n # Build command for CVD optimization with specific GPU\n cmd = (\n f\"CUDA_VISIBLE_DEVICES={args.gpu_id[device_id]} python camera_pose_annotation/cvd_opt/cvd_opt.py \"\n f\"--dir_path {dir_path} \"\n f\"--w_grad 2.0 --w_normal 5.0 \"\n )\n if args.only_depth:\n cmd += \"--only_depth \"\n process = subprocess.Popen(\n cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE\n )\n stdout, stderr = process.communicate()\n if process.returncode != 0:\n print(f\"Error optimizing CVD for {row['id']}: {stderr.decode()}\")\n\n\ndef worker(task_queue, args, worker_id, pbar):\n \"\"\"Worker function for parallel CVD optimization processing.\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n process_single_row(row, index, args, worker_id)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for CVD batch processing.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\"--dir_path\", type=str, default=\"./outputs\")\n parser.add_argument(\"--only_depth\", action=\"store_true\", help=\"Only save optimized depth\")\n parser.add_argument(\n \"--gpu_id\", type=str, default=\"0\", help=\"Comma-separated list of GPU IDs to use\"\n )\n parser.add_argument(\n \"--num_workers\",\n type=int,\n default=4,\n help=\"Number of workers for parallel processing\",\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n\n # Parse GPU configuration\n args.gpu_num = len(args.gpu_id.split(\",\"))\n args.gpu_id = [int(gpu) for gpu in args.gpu_id.split(\",\")]\n\n df = pd.read_csv(args.csv_path)\n\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(df.iterrows(), total=len(df), desc=\"Processing rows\"):\n process_single_row(row, index, args)\n else:\n # Parallel processing with multiple workers\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in df.iterrows():\n task_queue.put((index, row))\n with tqdm(total=len(df), desc=\"Processing rows\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for id in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, args, id, pbar))\n\n for future in concurrent.futures.as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/__init__.py", "content": "" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/__init__.py", "content": "" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/corr.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Correlation block for MegaSaM.\"\"\"\n\nimport torch\nimport torch.nn.functional as F\nfrom .utils.utils import bilinear_sampler\n\n# pylint: disable=g-import-not-at-top\ntry:\n import alt_cuda_corr\nexcept: # pylint: disable=bare-except\n # alt_cuda_corr is not compiled\n pass\n\n\nclass CorrBlock:\n \"\"\"Correlation block for MegaSaM.\"\"\"\n\n def __init__(self, fmap1, fmap2, num_levels=4, radius=4):\n self.num_levels = num_levels\n self.radius = radius\n self.corr_pyramid = []\n\n # all pairs correlation\n corr = CorrBlock.corr(fmap1, fmap2)\n\n batch, h1, w1, dim, h2, w2 = corr.shape\n corr = corr.reshape(batch * h1 * w1, dim, h2, w2)\n\n self.corr_pyramid.append(corr)\n for _ in range(self.num_levels - 1):\n corr = F.avg_pool2d(corr, 2, stride=2)\n self.corr_pyramid.append(corr)\n\n def __call__(self, coords):\n r = self.radius\n coords = coords.permute(0, 2, 3, 1)\n batch, h1, w1, _ = coords.shape\n\n out_pyramid = []\n for i in range(self.num_levels):\n corr = self.corr_pyramid[i]\n dx = torch.linspace(-r, r, 2 * r + 1)\n dy = torch.linspace(-r, r, 2 * r + 1)\n delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(coords.device)\n\n centroid_lvl = coords.reshape(batch * h1 * w1, 1, 1, 2) / 2**i\n delta_lvl = delta.view(1, 2 * r + 1, 2 * r + 1, 2)\n coords_lvl = centroid_lvl + delta_lvl\n\n corr = bilinear_sampler(corr, coords_lvl)\n corr = corr.view(batch, h1, w1, -1)\n out_pyramid.append(corr)\n\n out = torch.cat(out_pyramid, dim=-1)\n return out.permute(0, 3, 1, 2).contiguous().float()\n\n @classmethod\n def corr(cls, fmap1, fmap2):\n del cls\n batch, dim, ht, wd = fmap1.shape\n fmap1 = fmap1.view(batch, dim, ht * wd)\n fmap2 = fmap2.view(batch, dim, ht * wd)\n\n corr = torch.matmul(fmap1.transpose(1, 2), fmap2)\n corr = corr.view(batch, ht, wd, 1, ht, wd)\n return corr / torch.sqrt(torch.tensor(dim).float())\n\n\nclass AlternateCorrBlock:\n \"\"\"Correlation block for MegaSaM.\"\"\"\n\n def __init__(self, fmap1, fmap2, num_levels=4, radius=4):\n self.num_levels = num_levels\n self.radius = radius\n\n self.pyramid = [(fmap1, fmap2)]\n for _ in range(self.num_levels):\n fmap1 = F.avg_pool2d(fmap1, 2, stride=2)\n fmap2 = F.avg_pool2d(fmap2, 2, stride=2)\n self.pyramid.append((fmap1, fmap2))\n\n def __call__(self, coords):\n coords = coords.permute(0, 2, 3, 1)\n # pylint: disable=invalid-name\n B, H, W, _ = coords.shape\n dim = self.pyramid[0][0].shape[1]\n\n corr_list = []\n for i in range(self.num_levels):\n r = self.radius\n fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous()\n fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous()\n\n coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous()\n (corr,) = alt_cuda_corr.forward(fmap1_i, fmap2_i, coords_i, r)\n corr_list.append(corr.squeeze(1))\n\n corr = torch.stack(corr_list, dim=1)\n corr = corr.reshape(B, -1, H, W)\n return corr / torch.sqrt(torch.tensor(dim).float())\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/datasets.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Dataset classes for MegaSaM.\"\"\"\n\nimport glob\nimport os\nimport os.path as osp\nimport random\n\nimport numpy as np\nimport torch\nfrom torch.utils import data\nfrom utils import frame_utils\nfrom utils.augmentor import FlowAugmentor\nfrom utils.augmentor import SparseFlowAugmentor\n\n\nclass FlowDataset(data.Dataset):\n \"\"\"Base class for flow datasets.\"\"\"\n\n def __init__(self, aug_params=None, sparse=False):\n self.augmentor = None\n self.sparse = sparse\n if aug_params is not None:\n if sparse:\n self.augmentor = SparseFlowAugmentor(**aug_params)\n else:\n self.augmentor = FlowAugmentor(**aug_params)\n\n self.is_test = False\n self.init_seed = False\n self.flow_list = []\n self.image_list = []\n self.extra_info = []\n\n def __getitem__(self, index):\n\n if self.is_test:\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n img1 = np.array(img1).astype(np.uint8)[..., :3]\n img2 = np.array(img2).astype(np.uint8)[..., :3]\n img1 = torch.from_numpy(img1).permute(2, 0, 1).float()\n img2 = torch.from_numpy(img2).permute(2, 0, 1).float()\n return img1, img2, self.extra_info[index]\n\n if not self.init_seed:\n worker_info = torch.utils.data.get_worker_info()\n if worker_info is not None:\n torch.manual_seed(worker_info.id)\n np.random.seed(worker_info.id)\n random.seed(worker_info.id)\n self.init_seed = True\n\n index = index % len(self.image_list)\n valid = None\n if self.sparse:\n flow, valid = frame_utils.readFlowKITTI(self.flow_list[index])\n else:\n flow = frame_utils.read_gen(self.flow_list[index])\n\n img1 = frame_utils.read_gen(self.image_list[index][0])\n img2 = frame_utils.read_gen(self.image_list[index][1])\n\n flow = np.array(flow).astype(np.float32)\n img1 = np.array(img1).astype(np.uint8)\n img2 = np.array(img2).astype(np.uint8)\n\n # grayscale images\n if len(img1.shape) == 2:\n img1 = np.tile(img1[..., None], (1, 1, 3))\n img2 = np.tile(img2[..., None], (1, 1, 3))\n else:\n img1 = img1[..., :3]\n img2 = img2[..., :3]\n\n if self.augmentor is not None:\n if self.sparse:\n img1, img2, flow, valid = self.augmentor(img1, img2, flow, valid)\n else:\n img1, img2, flow = self.augmentor(img1, img2, flow)\n\n img1 = torch.from_numpy(img1).permute(2, 0, 1).float()\n img2 = torch.from_numpy(img2).permute(2, 0, 1).float()\n flow = torch.from_numpy(flow).permute(2, 0, 1).float()\n\n if valid is not None:\n valid = torch.from_numpy(valid)\n else:\n valid = (flow[0].abs() < 1000) & (flow[1].abs() < 1000)\n\n return img1, img2, flow, valid.float()\n\n def __rmul__(self, v):\n self.flow_list = v * self.flow_list\n self.image_list = v * self.image_list\n return self\n\n def __len__(self):\n return len(self.image_list)\n\n\nclass MpiSintel(FlowDataset):\n \"\"\"MpiSintel dataset.\"\"\"\n\n def __init__(\n self,\n aug_params=None,\n split='training',\n root='datasets/Sintel',\n dstype='clean',\n ):\n super(MpiSintel, self).__init__(aug_params)\n flow_root = osp.join(root, split, 'flow')\n image_root = osp.join(root, split, dstype)\n\n if split == 'test':\n self.is_test = True\n\n for scene in os.listdir(image_root):\n image_list = sorted(glob(osp.join(image_root, scene, '*.png')))\n for i in range(len(image_list) - 1):\n self.image_list += [[image_list[i], image_list[i + 1]]]\n self.extra_info += [(scene, i)] # scene and frame_id\n\n if split != 'test':\n self.flow_list += sorted(glob(osp.join(flow_root, scene, '*.flo')))\n\n\nclass FlyingChairs(FlowDataset):\n \"\"\"FlyingChairs dataset.\"\"\"\n\n def __init__(\n self,\n aug_params=None,\n split='train',\n root='datasets/FlyingChairs_release/data',\n ):\n super(FlyingChairs, self).__init__(aug_params)\n\n images = sorted(glob(osp.join(root, '*.ppm')))\n flows = sorted(glob(osp.join(root, '*.flo')))\n assert len(images) // 2 == len(flows)\n\n split_list = np.loadtxt('chairs_split.txt', dtype=np.int32)\n for i in range(len(flows)):\n exid = split_list[i]\n if (split == 'training' and exid == 1) or (\n split == 'validation' and exid == 2\n ):\n self.flow_list += [flows[i]]\n self.image_list += [[images[2 * i], images[2 * i + 1]]]\n\n\nclass FlyingThings3D(FlowDataset):\n \"\"\"FlyingThings3D dataset.\"\"\"\n\n def __init__(\n self,\n aug_params=None,\n root='datasets/FlyingThings3D',\n dstype='frames_cleanpass',\n ):\n super(FlyingThings3D, self).__init__(aug_params)\n\n for cam in ['left']:\n for direction in ['into_future', 'into_past']:\n image_dirs = sorted(glob(osp.join(root, dstype, 'TRAIN/*/*')))\n image_dirs = sorted([osp.join(f, cam) for f in image_dirs])\n\n flow_dirs = sorted(glob(osp.join(root, 'optical_flow/TRAIN/*/*')))\n flow_dirs = sorted([osp.join(f, direction, cam) for f in flow_dirs])\n\n for idir, fdir in zip(image_dirs, flow_dirs):\n images = sorted(glob(osp.join(idir, '*.png')))\n flows = sorted(glob(osp.join(fdir, '*.pfm')))\n for i in range(len(flows) - 1):\n if direction == 'into_future':\n self.image_list += [[images[i], images[i + 1]]]\n self.flow_list += [flows[i]]\n elif direction == 'into_past':\n self.image_list += [[images[i + 1], images[i]]]\n self.flow_list += [flows[i + 1]]\n\n\nclass KITTI(FlowDataset):\n \"\"\"KITTI dataset.\"\"\"\n\n def __init__(self, aug_params=None, split='training', root='datasets/KITTI'):\n super(KITTI, self).__init__(aug_params, sparse=True)\n if split == 'testing':\n self.is_test = True\n\n root = osp.join(root, split)\n images1 = sorted(glob(osp.join(root, 'image_2/*_10.png')))\n images2 = sorted(glob(osp.join(root, 'image_2/*_11.png')))\n\n for img1, img2 in zip(images1, images2):\n frame_id = img1.split('/')[-1]\n self.extra_info += [[frame_id]]\n self.image_list += [[img1, img2]]\n\n if split == 'training':\n self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png')))\n\n\nclass HD1K(FlowDataset):\n \"\"\"HD1K dataset.\"\"\"\n\n def __init__(self, aug_params=None, root='datasets/HD1k'):\n super(HD1K, self).__init__(aug_params, sparse=True)\n\n seq_ix = 0\n while 1:\n flows = sorted(\n glob(\n os.path.join(root, 'hd1k_flow_gt', 'flow_occ/%06d_*.png' % seq_ix)\n )\n )\n images = sorted(\n glob(os.path.join(root, 'hd1k_input', 'image_2/%06d_*.png' % seq_ix))\n )\n\n if not flows:\n break\n\n for i in range(len(flows) - 1):\n self.flow_list += [flows[i]]\n self.image_list += [[images[i], images[i + 1]]]\n\n seq_ix += 1\n\n\n# pylint: disable=invalid-name\ndef fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'):\n \"\"\"Create the data loader for the corresponding training set.\"\"\"\n\n if args.stage == 'chairs':\n aug_params = {\n 'crop_size': args.image_size,\n 'min_scale': -0.1,\n 'max_scale': 1.0,\n 'do_flip': True,\n }\n train_dataset = FlyingChairs(aug_params, split='training')\n\n elif args.stage == 'things':\n aug_params = {\n 'crop_size': args.image_size,\n 'min_scale': -0.4,\n 'max_scale': 0.8,\n 'do_flip': True,\n }\n clean_dataset = FlyingThings3D(aug_params, dstype='frames_cleanpass')\n final_dataset = FlyingThings3D(aug_params, dstype='frames_finalpass')\n train_dataset = clean_dataset + final_dataset\n\n elif args.stage == 'sintel':\n aug_params = {\n 'crop_size': args.image_size,\n 'min_scale': -0.2,\n 'max_scale': 0.6,\n 'do_flip': True,\n }\n things = FlyingThings3D(aug_params, dstype='frames_cleanpass')\n sintel_clean = MpiSintel(aug_params, split='training', dstype='clean')\n sintel_final = MpiSintel(aug_params, split='training', dstype='final')\n\n if TRAIN_DS == 'C+T+K+S+H':\n kitti = KITTI({\n 'crop_size': args.image_size,\n 'min_scale': -0.3,\n 'max_scale': 0.5,\n 'do_flip': True,\n })\n hd1k = HD1K({\n 'crop_size': args.image_size,\n 'min_scale': -0.5,\n 'max_scale': 0.2,\n 'do_flip': True,\n })\n train_dataset = (\n 100 * sintel_clean\n + 100 * sintel_final\n + 200 * kitti\n + 5 * hd1k\n + things\n )\n\n elif TRAIN_DS == 'C+T+K/S':\n train_dataset = 100 * sintel_clean + 100 * sintel_final + things\n else:\n raise ValueError('Unknown split: %s' % TRAIN_DS)\n\n elif args.stage == 'kitti':\n aug_params = {\n 'crop_size': args.image_size,\n 'min_scale': -0.2,\n 'max_scale': 0.4,\n 'do_flip': False,\n }\n train_dataset = KITTI(aug_params, split='training')\n else:\n raise ValueError('Unknown training set: %s' % args.stage)\n\n train_loader = data.DataLoader(\n train_dataset,\n batch_size=args.batch_size,\n pin_memory=False,\n shuffle=True,\n num_workers=4,\n drop_last=True,\n )\n\n print('Training with %d image pairs' % len(train_dataset))\n return train_loader\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/extractor.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Network layer classes for MegaSaM.\"\"\"\n\nimport torch\nfrom torch import nn\n\n\nclass ResidualBlock(nn.Module):\n \"\"\"Residual block for MegaSaM.\"\"\"\n\n def __init__(self, in_planes, planes, norm_fn='group', stride=1):\n super(ResidualBlock, self).__init__()\n\n self.conv1 = nn.Conv2d(\n in_planes, planes, kernel_size=3, padding=1, stride=stride\n )\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)\n self.relu = nn.ReLU(inplace=True)\n\n num_groups = planes // 8\n\n if norm_fn == 'group':\n self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)\n self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)\n if stride != 1:\n self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)\n\n elif norm_fn == 'batch':\n self.norm1 = nn.BatchNorm2d(planes)\n self.norm2 = nn.BatchNorm2d(planes)\n if stride != 1:\n self.norm3 = nn.BatchNorm2d(planes)\n\n elif norm_fn == 'instance':\n self.norm1 = nn.InstanceNorm2d(planes)\n self.norm2 = nn.InstanceNorm2d(planes)\n if stride != 1:\n self.norm3 = nn.InstanceNorm2d(planes)\n\n elif norm_fn == 'none':\n self.norm1 = nn.Sequential()\n self.norm2 = nn.Sequential()\n if stride != 1:\n self.norm3 = nn.Sequential()\n\n if stride == 1:\n self.downsample = None\n\n else:\n self.downsample = nn.Sequential(\n nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3\n )\n\n def forward(self, x):\n y = x\n y = self.relu(self.norm1(self.conv1(y)))\n y = self.relu(self.norm2(self.conv2(y)))\n\n if self.downsample is not None:\n x = self.downsample(x)\n\n return self.relu(x + y)\n\n\nclass BottleneckBlock(nn.Module):\n \"\"\"Bottleneck block for MegaSaM.\"\"\"\n\n def __init__(self, in_planes, planes, norm_fn='group', stride=1):\n super(BottleneckBlock, self).__init__()\n\n self.conv1 = nn.Conv2d(in_planes, planes // 4, kernel_size=1, padding=0)\n self.conv2 = nn.Conv2d(\n planes // 4, planes // 4, kernel_size=3, padding=1, stride=stride\n )\n self.conv3 = nn.Conv2d(planes // 4, planes, kernel_size=1, padding=0)\n self.relu = nn.ReLU(inplace=True)\n\n num_groups = planes // 8\n\n if norm_fn == 'group':\n self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4)\n self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4)\n self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)\n if stride != 1:\n self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)\n\n elif norm_fn == 'batch':\n self.norm1 = nn.BatchNorm2d(planes // 4)\n self.norm2 = nn.BatchNorm2d(planes // 4)\n self.norm3 = nn.BatchNorm2d(planes)\n if stride != 1:\n self.norm4 = nn.BatchNorm2d(planes)\n\n elif norm_fn == 'instance':\n self.norm1 = nn.InstanceNorm2d(planes // 4)\n self.norm2 = nn.InstanceNorm2d(planes // 4)\n self.norm3 = nn.InstanceNorm2d(planes)\n if stride != 1:\n self.norm4 = nn.InstanceNorm2d(planes)\n\n elif norm_fn == 'none':\n self.norm1 = nn.Sequential()\n self.norm2 = nn.Sequential()\n self.norm3 = nn.Sequential()\n if stride != 1:\n self.norm4 = nn.Sequential()\n\n if stride == 1:\n self.downsample = None\n\n else:\n self.downsample = nn.Sequential(\n nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4\n )\n\n def forward(self, x):\n y = x\n y = self.relu(self.norm1(self.conv1(y)))\n y = self.relu(self.norm2(self.conv2(y)))\n y = self.relu(self.norm3(self.conv3(y)))\n\n if self.downsample is not None:\n x = self.downsample(x)\n\n return self.relu(x + y)\n\n\nclass BasicEncoder(nn.Module):\n \"\"\"Basic encoder for MegaSaM.\"\"\"\n\n def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):\n super(BasicEncoder, self).__init__()\n self.norm_fn = norm_fn\n\n if self.norm_fn == 'group':\n self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)\n\n elif self.norm_fn == 'batch':\n self.norm1 = nn.BatchNorm2d(64)\n\n elif self.norm_fn == 'instance':\n self.norm1 = nn.InstanceNorm2d(64)\n\n elif self.norm_fn == 'none':\n self.norm1 = nn.Sequential()\n\n self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)\n self.relu1 = nn.ReLU(inplace=True)\n\n self.in_planes = 64\n self.layer1 = self._make_layer(64, stride=1)\n self.layer2 = self._make_layer(96, stride=2)\n self.layer3 = self._make_layer(128, stride=2)\n\n # output convolution\n self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)\n\n self.dropout = None\n if dropout > 0:\n self.dropout = nn.Dropout2d(p=dropout)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):\n if m.weight is not None:\n nn.init.constant_(m.weight, 1)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def _make_layer(self, dim, stride=1):\n layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)\n layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)\n layers = (layer1, layer2)\n\n self.in_planes = dim\n return nn.Sequential(*layers)\n\n def forward(self, x):\n\n # if input is list, combine batch dimension\n is_list = isinstance(x, tuple) or isinstance(x, list)\n if is_list:\n batch_dim = x[0].shape[0]\n x = torch.cat(x, dim=0)\n\n x = self.conv1(x)\n x = self.norm1(x)\n x = self.relu1(x)\n\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n\n x = self.conv2(x)\n\n if self.training and self.dropout is not None:\n x = self.dropout(x)\n\n if is_list:\n x = torch.split(x, [batch_dim, batch_dim], dim=0) # pylint: disable=undefined-variable\n\n return x\n\n\nclass SmallEncoder(nn.Module):\n \"\"\"Small encoder for MegaSaM.\"\"\"\n\n def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):\n super(SmallEncoder, self).__init__()\n self.norm_fn = norm_fn\n\n if self.norm_fn == 'group':\n self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32)\n\n elif self.norm_fn == 'batch':\n self.norm1 = nn.BatchNorm2d(32)\n\n elif self.norm_fn == 'instance':\n self.norm1 = nn.InstanceNorm2d(32)\n\n elif self.norm_fn == 'none':\n self.norm1 = nn.Sequential()\n\n self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3)\n self.relu1 = nn.ReLU(inplace=True)\n\n self.in_planes = 32\n self.layer1 = self._make_layer(32, stride=1)\n self.layer2 = self._make_layer(64, stride=2)\n self.layer3 = self._make_layer(96, stride=2)\n\n self.dropout = None\n if dropout > 0:\n self.dropout = nn.Dropout2d(p=dropout)\n\n self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):\n if m.weight is not None:\n nn.init.constant_(m.weight, 1)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def _make_layer(self, dim, stride=1):\n layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride)\n layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1)\n layers = (layer1, layer2)\n\n self.in_planes = dim\n return nn.Sequential(*layers)\n\n def forward(self, x):\n\n # if input is list, combine batch dimension\n is_list = isinstance(x, tuple) or isinstance(x, list)\n if is_list:\n batch_dim = x[0].shape[0]\n x = torch.cat(x, dim=0)\n\n x = self.conv1(x)\n x = self.norm1(x)\n x = self.relu1(x)\n\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.conv2(x)\n\n if self.training and self.dropout is not None:\n x = self.dropout(x)\n\n if is_list:\n x = torch.split(x, [batch_dim, batch_dim], dim=0) # pylint: disable=undefined-variable\n\n return x\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/raft.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"RAFT network for MegaSaM.\"\"\"\n\nfrom .corr import AlternateCorrBlock\nfrom .corr import CorrBlock\nfrom .extractor import BasicEncoder\nfrom .extractor import SmallEncoder\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom .update import BasicUpdateBlock\nfrom .update import SmallUpdateBlock\nfrom .utils.utils import coords_grid\nfrom .utils.utils import upflow8\n\ntry:\n autocast = torch.cuda.amp.autocast\nexcept: # pylint: disable=bare-except\n # dummy autocast for PyTorch < 1.6\n class autocast: # pylint: disable=invalid-name\n\n def __init__(self, enabled):\n pass\n\n def __enter__(self):\n pass\n\n def __exit__(self, *args):\n pass\n\n\nclass RAFT(nn.Module):\n \"\"\"RAFT network for MegaSaM.\"\"\"\n\n def __init__(self, args):\n super(RAFT, self).__init__()\n self.args = args\n self.mixed_precision = True\n if args.small:\n self.hidden_dim = hdim = 96\n self.context_dim = cdim = 64\n args.corr_levels = 4\n args.corr_radius = 3\n\n else:\n self.hidden_dim = hdim = 128\n self.context_dim = cdim = 128\n args.corr_levels = 4\n args.corr_radius = 4\n\n if 'dropout' not in self.args:\n self.args.dropout = 0\n\n if 'alternate_corr' not in self.args:\n self.args.alternate_corr = False\n\n # feature network, context network, and update block\n if args.small:\n self.fnet = SmallEncoder(\n output_dim=128, norm_fn='instance', dropout=args.dropout\n )\n self.cnet = SmallEncoder(\n output_dim=hdim + cdim, norm_fn='none', dropout=args.dropout\n )\n self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)\n\n else:\n self.fnet = BasicEncoder(\n output_dim=256, norm_fn='instance', dropout=args.dropout\n )\n self.cnet = BasicEncoder(\n output_dim=hdim + cdim, norm_fn='batch', dropout=args.dropout\n )\n self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)\n\n def freeze_bn(self):\n for m in self.modules():\n if isinstance(m, nn.BatchNorm2d):\n m.eval()\n\n def initialize_flow(self, img):\n \"\"\"Flow is represented as difference between two coordinate grids flow = coords1 - coords0.\"\"\"\n # pylint: disable=invalid-name\n N, _, H, W = img.shape\n coords0 = coords_grid(N, H // 8, W // 8).to(img.device)\n coords1 = coords_grid(N, H // 8, W // 8).to(img.device)\n\n # optical flow computed as difference: flow = coords1 - coords0\n return coords0, coords1\n\n def upsample_flow(self, flow, mask):\n \"\"\"Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination.\"\"\"\n # pylint: disable=invalid-name\n N, _, H, W = flow.shape\n mask = mask.view(N, 1, 9, 8, 8, H, W)\n mask = torch.softmax(mask, dim=2)\n\n up_flow = F.unfold(8 * flow, [3, 3], padding=1)\n up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)\n\n up_flow = torch.sum(mask * up_flow, dim=2)\n up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)\n return up_flow.reshape(N, 2, 8 * H, 8 * W)\n\n def forward(\n self,\n image1,\n image2,\n iters=12,\n flow_init=None,\n upsample=True,\n test_mode=False,\n ):\n \"\"\"Estimate optical flow between pair of frames.\"\"\"\n\n image1 = 2 * (image1 / 255.0) - 1.0\n image2 = 2 * (image2 / 255.0) - 1.0\n\n image1 = image1.contiguous()\n image2 = image2.contiguous()\n\n hdim = self.hidden_dim\n cdim = self.context_dim\n\n # run the feature network\n with autocast(enabled=self.mixed_precision):\n fmap1, fmap2 = self.fnet([image1, image2])\n\n fmap1 = fmap1.float()\n fmap2 = fmap2.float()\n if self.args.alternate_corr:\n corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)\n else:\n corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)\n\n # run the context network\n with autocast(enabled=self.mixed_precision):\n cnet = self.cnet(image1)\n net, inp = torch.split(cnet, [hdim, cdim], dim=1)\n net = torch.tanh(net)\n inp = torch.relu(inp)\n\n coords0, coords1 = self.initialize_flow(image1)\n\n if flow_init is not None:\n coords1 = coords1 + flow_init\n\n flow_predictions = []\n flow_up = None\n for _ in range(iters):\n coords1 = coords1.detach()\n corr = corr_fn(coords1) # index correlation volume\n\n flow = coords1 - coords0\n with autocast(enabled=self.mixed_precision):\n net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)\n\n # F(t+1) = F(t) + \\Delta(t)\n coords1 = coords1 + delta_flow\n\n # upsample predictions\n if up_mask is None:\n flow_up = upflow8(coords1 - coords0)\n else:\n flow_up = self.upsample_flow(coords1 - coords0, up_mask)\n\n flow_predictions.append(flow_up)\n\n if test_mode:\n if flow_up is None:\n raise ValueError('flow_up is None')\n return coords1 - coords0, flow_up, net\n\n return flow_predictions\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/update.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Update block for consistent video depth optimization.\"\"\"\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\n\n\nclass FlowHead(nn.Module):\n\n def __init__(self, input_dim=128, hidden_dim=256):\n super(FlowHead, self).__init__()\n self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)\n self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n return self.conv2(self.relu(self.conv1(x)))\n\n\nclass ConvGRU(nn.Module):\n \"\"\"GRU with convolution.\"\"\"\n\n def __init__(self, hidden_dim=128, input_dim=192 + 128):\n super(ConvGRU, self).__init__()\n self.convz = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)\n self.convr = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)\n self.convq = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)\n\n def forward(self, h, x):\n hx = torch.cat([h, x], dim=1)\n\n z = torch.sigmoid(self.convz(hx))\n r = torch.sigmoid(self.convr(hx))\n q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1)))\n\n h = (1 - z) * h + z * q\n return h\n\n\nclass SepConvGRU(nn.Module):\n \"\"\"GRU with separate convolution for horizontal and vertical directions.\"\"\"\n\n def __init__(self, hidden_dim=128, input_dim=192 + 128):\n super(SepConvGRU, self).__init__()\n self.convz1 = nn.Conv2d(\n hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)\n )\n self.convr1 = nn.Conv2d(\n hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)\n )\n self.convq1 = nn.Conv2d(\n hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)\n )\n\n self.convz2 = nn.Conv2d(\n hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)\n )\n self.convr2 = nn.Conv2d(\n hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)\n )\n self.convq2 = nn.Conv2d(\n hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)\n )\n\n def forward(self, h, x):\n # horizontal\n hx = torch.cat([h, x], dim=1)\n z = torch.sigmoid(self.convz1(hx))\n r = torch.sigmoid(self.convr1(hx))\n q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1)))\n h = (1 - z) * h + z * q\n\n # vertical\n hx = torch.cat([h, x], dim=1)\n z = torch.sigmoid(self.convz2(hx))\n r = torch.sigmoid(self.convr2(hx))\n q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1)))\n h = (1 - z) * h + z * q\n\n return h\n\n\nclass SmallMotionEncoder(nn.Module):\n \"\"\"Small motion encoder for MegaSaM.\"\"\"\n\n def __init__(self, args):\n super(SmallMotionEncoder, self).__init__()\n cor_planes = args.corr_levels * (2 * args.corr_radius + 1) ** 2\n self.convc1 = nn.Conv2d(cor_planes, 96, 1, padding=0)\n self.convf1 = nn.Conv2d(2, 64, 7, padding=3)\n self.convf2 = nn.Conv2d(64, 32, 3, padding=1)\n self.conv = nn.Conv2d(128, 80, 3, padding=1)\n\n def forward(self, flow, corr):\n cor = F.relu(self.convc1(corr))\n flo = F.relu(self.convf1(flow))\n flo = F.relu(self.convf2(flo))\n cor_flo = torch.cat([cor, flo], dim=1)\n out = F.relu(self.conv(cor_flo))\n return torch.cat([out, flow], dim=1)\n\n\nclass BasicMotionEncoder(nn.Module):\n \"\"\"Basic motion encoder for MegaSaM.\"\"\"\n\n def __init__(self, args):\n super(BasicMotionEncoder, self).__init__()\n cor_planes = args.corr_levels * (2 * args.corr_radius + 1) ** 2\n self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)\n self.convc2 = nn.Conv2d(256, 192, 3, padding=1)\n self.convf1 = nn.Conv2d(2, 128, 7, padding=3)\n self.convf2 = nn.Conv2d(128, 64, 3, padding=1)\n self.conv = nn.Conv2d(64 + 192, 128 - 2, 3, padding=1)\n\n def forward(self, flow, corr):\n cor = F.relu(self.convc1(corr))\n cor = F.relu(self.convc2(cor))\n flo = F.relu(self.convf1(flow))\n flo = F.relu(self.convf2(flo))\n\n cor_flo = torch.cat([cor, flo], dim=1)\n out = F.relu(self.conv(cor_flo))\n return torch.cat([out, flow], dim=1)\n\n\nclass SmallUpdateBlock(nn.Module):\n \"\"\"Small update block for MegaSaM.\"\"\"\n\n def __init__(self, args, hidden_dim=96):\n super(SmallUpdateBlock, self).__init__()\n self.encoder = SmallMotionEncoder(args)\n self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=82 + 64)\n self.flow_head = FlowHead(hidden_dim, hidden_dim=128)\n\n def forward(self, net, inp, corr, flow):\n motion_features = self.encoder(flow, corr)\n inp = torch.cat([inp, motion_features], dim=1)\n net = self.gru(net, inp)\n delta_flow = self.flow_head(net)\n\n return net, None, delta_flow\n\n\nclass BasicUpdateBlock(nn.Module):\n \"\"\"Basic update block for MegaSaM.\"\"\"\n\n def __init__(self, args, hidden_dim=128, input_dim=128):\n super(BasicUpdateBlock, self).__init__()\n self.args = args\n self.encoder = BasicMotionEncoder(args)\n self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128 + hidden_dim)\n self.flow_head = FlowHead(hidden_dim, hidden_dim=256)\n\n self.mask = nn.Sequential(\n nn.Conv2d(128, 256, 3, padding=1),\n nn.ReLU(inplace=True),\n nn.Conv2d(256, 64 * 9, 1, padding=0),\n )\n\n def forward(self, net, inp, corr, flow, upsample=True):\n motion_features = self.encoder(flow, corr)\n inp = torch.cat([inp, motion_features], dim=1)\n\n net = self.gru(net, inp)\n delta_flow = self.flow_head(net)\n\n # scale mask to balence gradients\n mask = 0.25 * self.mask(net)\n return net, mask, delta_flow\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/utils/__init__.py", "content": "" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/utils/augmentor.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Augmentation utils for MegaSaM.\"\"\"\n# pylint: disable=g-import-not-at-top\n# pylint: disable=g-importing-member\n\nimport cv2\nimport numpy as np\nfrom PIL import Image\n\ncv2.setNumThreads(0)\ncv2.ocl.setUseOpenCL(False)\n\nfrom torchvision.transforms import ColorJitter\n\n\nclass FlowAugmentor:\n \"\"\"Augmentation for flow for MegaSaM.\"\"\"\n\n def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True):\n\n # spatial augmentation params\n self.crop_size = crop_size\n self.min_scale = min_scale\n self.max_scale = max_scale\n self.spatial_aug_prob = 0.8\n self.stretch_prob = 0.8\n self.max_stretch = 0.2\n\n # flip augmentation params\n self.do_flip = do_flip\n self.h_flip_prob = 0.5\n self.v_flip_prob = 0.1\n\n # photometric augmentation params\n self.photo_aug = ColorJitter(\n brightness=0.4, contrast=0.4, saturation=0.4, hue=0.5 / 3.14\n )\n self.asymmetric_color_aug_prob = 0.2\n self.eraser_aug_prob = 0.5\n\n def color_transform(self, img1, img2):\n \"\"\"Photometric augmentation.\"\"\"\n\n # asymmetric\n if np.random.rand() < self.asymmetric_color_aug_prob:\n img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8)\n img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8)\n\n # symmetric\n else:\n image_stack = np.concatenate([img1, img2], axis=0)\n image_stack = np.array(\n self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8\n )\n img1, img2 = np.split(image_stack, 2, axis=0)\n\n return img1, img2\n\n def eraser_transform(self, img1, img2, bounds=[50, 100]): # pylint: disable=dangerous-default-value\n \"\"\"Occlusion augmentation.\"\"\"\n\n ht, wd = img1.shape[:2]\n if np.random.rand() < self.eraser_aug_prob:\n mean_color = np.mean(img2.reshape(-1, 3), axis=0)\n for _ in range(np.random.randint(1, 3)):\n x0 = np.random.randint(0, wd)\n y0 = np.random.randint(0, ht)\n dx = np.random.randint(bounds[0], bounds[1])\n dy = np.random.randint(bounds[0], bounds[1])\n img2[y0 : y0 + dy, x0 : x0 + dx, :] = mean_color\n\n return img1, img2\n\n def spatial_transform(self, img1, img2, flow):\n \"\"\"Spatial augmentation.\"\"\"\n\n # randomly sample scale\n ht, wd = img1.shape[:2]\n min_scale = np.maximum(\n (self.crop_size[0] + 8) / float(ht), (self.crop_size[1] + 8) / float(wd)\n )\n\n scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)\n scale_x = scale\n scale_y = scale\n if np.random.rand() < self.stretch_prob:\n scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)\n scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)\n\n scale_x = np.clip(scale_x, min_scale, None)\n scale_y = np.clip(scale_y, min_scale, None)\n\n if np.random.rand() < self.spatial_aug_prob:\n # rescale the images\n img1 = cv2.resize(\n img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR\n )\n img2 = cv2.resize(\n img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR\n )\n flow = cv2.resize(\n flow, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR\n )\n flow = flow * [scale_x, scale_y]\n\n if self.do_flip:\n if np.random.rand() < self.h_flip_prob: # h-flip\n img1 = img1[:, ::-1]\n img2 = img2[:, ::-1]\n flow = flow[:, ::-1] * [-1.0, 1.0]\n\n if np.random.rand() < self.v_flip_prob: # v-flip\n img1 = img1[::-1, :]\n img2 = img2[::-1, :]\n flow = flow[::-1, :] * [1.0, -1.0]\n\n y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])\n x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])\n\n img1 = img1[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]\n img2 = img2[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]\n flow = flow[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]\n\n return img1, img2, flow\n\n def __call__(self, img1, img2, flow):\n img1, img2 = self.color_transform(img1, img2)\n img1, img2 = self.eraser_transform(img1, img2)\n img1, img2, flow = self.spatial_transform(img1, img2, flow)\n\n img1 = np.ascontiguousarray(img1)\n img2 = np.ascontiguousarray(img2)\n flow = np.ascontiguousarray(flow)\n\n return img1, img2, flow\n\n\nclass SparseFlowAugmentor:\n \"\"\"Augmentation for sparse flow for MegaSaM.\"\"\"\n\n def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False):\n # spatial augmentation params\n self.crop_size = crop_size\n self.min_scale = min_scale\n self.max_scale = max_scale\n self.spatial_aug_prob = 0.8\n self.stretch_prob = 0.8\n self.max_stretch = 0.2\n\n # flip augmentation params\n self.do_flip = do_flip\n self.h_flip_prob = 0.5\n self.v_flip_prob = 0.1\n\n # photometric augmentation params\n self.photo_aug = ColorJitter(\n brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3 / 3.14\n )\n self.asymmetric_color_aug_prob = 0.2\n self.eraser_aug_prob = 0.5\n\n def color_transform(self, img1, img2):\n image_stack = np.concatenate([img1, img2], axis=0)\n image_stack = np.array(\n self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8\n )\n img1, img2 = np.split(image_stack, 2, axis=0)\n return img1, img2\n\n def eraser_transform(self, img1, img2):\n ht, wd = img1.shape[:2]\n if np.random.rand() < self.eraser_aug_prob:\n mean_color = np.mean(img2.reshape(-1, 3), axis=0)\n for _ in range(np.random.randint(1, 3)):\n x0 = np.random.randint(0, wd)\n y0 = np.random.randint(0, ht)\n dx = np.random.randint(50, 100)\n dy = np.random.randint(50, 100)\n img2[y0 : y0 + dy, x0 : x0 + dx, :] = mean_color\n\n return img1, img2\n\n def resize_sparse_flow_map(self, flow, valid, fx=1.0, fy=1.0):\n \"\"\"Resize sparse flow map.\"\"\"\n ht, wd = flow.shape[:2]\n coords = np.meshgrid(np.arange(wd), np.arange(ht))\n coords = np.stack(coords, axis=-1)\n\n coords = coords.reshape(-1, 2).astype(np.float32)\n flow = flow.reshape(-1, 2).astype(np.float32)\n valid = valid.reshape(-1).astype(np.float32)\n\n coords0 = coords[valid >= 1]\n flow0 = flow[valid >= 1]\n\n ht1 = int(round(ht * fy))\n wd1 = int(round(wd * fx))\n\n coords1 = coords0 * [fx, fy]\n flow1 = flow0 * [fx, fy]\n\n xx = np.round(coords1[:, 0]).astype(np.int32)\n yy = np.round(coords1[:, 1]).astype(np.int32)\n\n v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1)\n xx = xx[v]\n yy = yy[v]\n flow1 = flow1[v]\n\n flow_img = np.zeros([ht1, wd1, 2], dtype=np.float32)\n valid_img = np.zeros([ht1, wd1], dtype=np.int32)\n\n flow_img[yy, xx] = flow1\n valid_img[yy, xx] = 1\n\n return flow_img, valid_img\n\n def spatial_transform(self, img1, img2, flow, valid):\n \"\"\"Randomly sample scale and apply it to images and flow map.\"\"\"\n\n ht, wd = img1.shape[:2]\n min_scale = np.maximum(\n (self.crop_size[0] + 1) / float(ht), (self.crop_size[1] + 1) / float(wd)\n )\n\n scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)\n scale_x = np.clip(scale, min_scale, None)\n scale_y = np.clip(scale, min_scale, None)\n\n if np.random.rand() < self.spatial_aug_prob:\n # rescale the images\n img1 = cv2.resize(\n img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR\n )\n img2 = cv2.resize(\n img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR\n )\n flow, valid = self.resize_sparse_flow_map(\n flow, valid, fx=scale_x, fy=scale_y\n )\n\n if self.do_flip:\n if np.random.rand() < 0.5: # h-flip\n img1 = img1[:, ::-1]\n img2 = img2[:, ::-1]\n flow = flow[:, ::-1] * [-1.0, 1.0]\n valid = valid[:, ::-1]\n\n margin_y = 20\n margin_x = 50\n\n y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0] + margin_y)\n x0 = np.random.randint(\n -margin_x, img1.shape[1] - self.crop_size[1] + margin_x\n )\n\n y0 = np.clip(y0, 0, img1.shape[0] - self.crop_size[0])\n x0 = np.clip(x0, 0, img1.shape[1] - self.crop_size[1])\n\n img1 = img1[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]\n img2 = img2[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]\n flow = flow[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]\n valid = valid[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]\n return img1, img2, flow, valid\n\n def __call__(self, img1, img2, flow, valid):\n img1, img2 = self.color_transform(img1, img2)\n img1, img2 = self.eraser_transform(img1, img2)\n img1, img2, flow, valid = self.spatial_transform(img1, img2, flow, valid)\n\n img1 = np.ascontiguousarray(img1)\n img2 = np.ascontiguousarray(img2)\n flow = np.ascontiguousarray(flow)\n valid = np.ascontiguousarray(valid)\n\n return img1, img2, flow, valid\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/utils/flow_viz.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Flow visualization code.\n\nBased on https://github.com/tomrunia/OpticalFlow_Visualization\n\"\"\"\n\nimport numpy as np\n\n\ndef make_colorwheel():\n \"\"\"Generates a color wheel for optical flow visualization.\n\n Baker et al. \"A Database and Evaluation Methodology for Optical Flow\"\n (ICCV, 2007)\n URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf\n\n Code follows the original C++ source code of Daniel Scharstein.\n Code follows the the Matlab source code of Deqing Sun.\n\n Returns:\n np.ndarray: Color wheel\n \"\"\"\n\n # pylint: disable=invalid-name\n RY = 15\n YG = 6\n GC = 4\n CB = 11\n BM = 13\n MR = 6\n\n ncols = RY + YG + GC + CB + BM + MR\n colorwheel = np.zeros((ncols, 3))\n col = 0\n\n # RY\n colorwheel[0:RY, 0] = 255\n colorwheel[0:RY, 1] = np.floor(255 * np.arange(0, RY) / RY)\n col = col + RY\n # YG\n colorwheel[col : col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG) / YG)\n colorwheel[col : col + YG, 1] = 255\n col = col + YG\n # GC\n colorwheel[col : col + GC, 1] = 255\n colorwheel[col : col + GC, 2] = np.floor(255 * np.arange(0, GC) / GC)\n col = col + GC\n # CB\n colorwheel[col : col + CB, 1] = 255 - np.floor(255 * np.arange(CB) / CB)\n colorwheel[col : col + CB, 2] = 255\n col = col + CB\n # BM\n colorwheel[col : col + BM, 2] = 255\n colorwheel[col : col + BM, 0] = np.floor(255 * np.arange(0, BM) / BM)\n col = col + BM\n # MR\n colorwheel[col : col + MR, 2] = 255 - np.floor(255 * np.arange(MR) / MR)\n colorwheel[col : col + MR, 0] = 255\n return colorwheel\n\n\ndef flow_uv_to_colors(u, v, convert_to_bgr=False):\n \"\"\"Applies the flow color wheel to (possibly clipped) flow components u and v.\n\n According to the C++ source code of Daniel Scharstein\n According to the Matlab source code of Deqing Sun\n\n Args:\n u (np.ndarray): Input horizontal flow of shape [H,W]\n v (np.ndarray): Input vertical flow of shape [H,W]\n convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to\n False.\n\n Returns:\n np.ndarray: Flow visualization image of shape [H,W,3]\n \"\"\"\n flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)\n colorwheel = make_colorwheel() # shape [55x3]\n ncols = colorwheel.shape[0]\n rad = np.sqrt(np.square(u) + np.square(v))\n a = np.arctan2(-v, -u) / np.pi\n fk = (a + 1) / 2 * (ncols - 1)\n k0 = np.floor(fk).astype(np.int32)\n k1 = k0 + 1\n k1[k1 == ncols] = 0\n f = fk - k0\n for i in range(colorwheel.shape[1]):\n tmp = colorwheel[:, i]\n col0 = tmp[k0] / 255.0\n col1 = tmp[k1] / 255.0\n col = (1 - f) * col0 + f * col1\n idx = rad <= 1\n col[idx] = 1 - rad[idx] * (1 - col[idx])\n col[~idx] = col[~idx] * 0.75 # out of range\n # Note the 2-i => BGR instead of RGB\n ch_idx = 2 - i if convert_to_bgr else i\n flow_image[:, :, ch_idx] = np.floor(255 * col)\n return flow_image\n\n\ndef flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):\n \"\"\"Expects a two dimensional flow image of shape.\n\n Args:\n flow_uv (np.ndarray): Flow UV image of shape [H,W,2]\n clip_flow (float, optional): Clip maximum of flow values. Defaults to\n None.\n convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to\n False.\n\n Returns:\n np.ndarray: Flow visualization image of shape [H,W,3]\n \"\"\"\n assert flow_uv.ndim == 3, 'input flow must have three dimensions'\n assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'\n if clip_flow is not None:\n flow_uv = np.clip(flow_uv, 0, clip_flow)\n u = flow_uv[:, :, 0]\n v = flow_uv[:, :, 1]\n rad = np.sqrt(np.square(u) + np.square(v))\n rad_max = np.max(rad)\n epsilon = 1e-5\n u = u / (rad_max + epsilon)\n v = v / (rad_max + epsilon)\n return flow_uv_to_colors(u, v, convert_to_bgr)\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/utils/frame_utils.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Frame utils for MegaSaM.\"\"\"\n\n# pylint: disable=invalid-name\n# pylint: disable=g-doc-args\n# pylint: disable=broad-exception-raised\n\nimport os\nimport re\n\nimport cv2\nimport numpy as np\nfrom PIL import Image\n\ncv2.setNumThreads(0)\ncv2.ocl.setUseOpenCL(False)\n\nTAG_CHAR = np.array([202021.25], np.float32)\n\n\ndef readFlow(fn):\n \"\"\"Read .flo file in Middlebury format.\"\"\"\n # Code adapted from:\n # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy\n\n # WARNING: this will work on little-endian architectures (eg Intel x86) only!\n # print 'fn = %s'%(fn)\n with open(fn, 'rb') as f:\n magic = np.fromfile(f, np.float32, count=1)\n if 202021.25 != magic:\n print('Magic number incorrect. Invalid .flo file')\n return None\n else:\n w = np.fromfile(f, np.int32, count=1)\n h = np.fromfile(f, np.int32, count=1)\n # print 'Reading %d x %d flo file\\n' % (w, h)\n data = np.fromfile(f, np.float32, count=2 * int(w) * int(h))\n # Reshape data into 3D array (columns, rows, bands)\n # The reshape here is for visualization, the original code is (w,h,2)\n return np.resize(data, (int(h), int(w), 2))\n\n\ndef readPFM(file):\n \"\"\"Read PFM file.\"\"\"\n file = open(file, 'rb')\n\n header = file.readline().rstrip()\n if header == b'PF':\n color = True\n elif header == b'Pf':\n color = False\n else:\n raise Exception('Not a PFM file.')\n\n dim_match = re.match(rb'^(\\d+)\\s(\\d+)\\s$', file.readline())\n if dim_match:\n width, height = map(int, dim_match.groups())\n else:\n raise Exception('Malformed PFM header.')\n\n scale = float(file.readline().rstrip())\n if scale < 0: # little-endian\n endian = '<'\n else:\n endian = '>' # big-endian\n\n data = np.fromfile(file, endian + 'f')\n shape = (height, width, 3) if color else (height, width)\n\n data = np.reshape(data, shape)\n data = np.flipud(data)\n return data\n\n\ndef writeFlow(filename, uv, v=None):\n \"\"\"Write optical flow to file.\n\n If v is None, uv is assumed to contain both u and v channels,\n stacked in depth.\n Original code by Deqing Sun, adapted from Daniel Scharstein.\n \"\"\"\n nBands = 2\n\n if v is None:\n assert uv.ndim == 3\n assert uv.shape[2] == 2\n u = uv[:, :, 0]\n v = uv[:, :, 1]\n else:\n u = uv\n\n assert u.shape == v.shape\n height, width = u.shape\n f = open(filename, 'wb')\n # write the header\n f.write(TAG_CHAR)\n np.array(width).astype(np.int32).tofile(f)\n np.array(height).astype(np.int32).tofile(f)\n # arrange into matrix form\n tmp = np.zeros((height, width * nBands))\n tmp[:, np.arange(width) * 2] = u\n tmp[:, np.arange(width) * 2 + 1] = v\n tmp.astype(np.float32).tofile(f)\n f.close()\n\n\ndef readFlowKITTI(filename):\n flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR)\n flow = flow[:, :, ::-1].astype(np.float32)\n flow, valid = flow[:, :, :2], flow[:, :, 2]\n flow = (flow - 2**15) / 64.0\n return flow, valid\n\n\ndef readDispKITTI(filename):\n disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0\n valid = disp > 0.0\n flow = np.stack([-disp, np.zeros_like(disp)], -1)\n return flow, valid\n\n\ndef writeFlowKITTI(filename, uv):\n uv = 64.0 * uv + 2**15\n valid = np.ones([uv.shape[0], uv.shape[1], 1])\n uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)\n cv2.imwrite(filename, uv[..., ::-1])\n\n\ndef read_gen(file_name, pil=False):\n \"\"\"Read image or flow file.\"\"\"\n del pil\n ext = os.path.splitext(file_name)[-1]\n if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':\n return Image.open(file_name)\n elif ext == '.bin' or ext == '.raw':\n return np.load(file_name)\n elif ext == '.flo':\n return readFlow(file_name).astype(np.float32) # pylint: disable=attribute-error\n elif ext == '.pfm':\n flow = readPFM(file_name).astype(np.float32)\n if len(flow.shape) == 2:\n return flow\n else:\n return flow[:, :, :-1]\n return []\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/core/utils/utils.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Utility functions for MegaSaM.\"\"\"\n\n# pylint: disable=invalid-name\n\nimport numpy as np\nfrom scipy import interpolate\nimport torch\nimport torch.nn.functional as F\n\n\nclass InputPadder:\n \"\"\"Pads images such that dimensions are divisible by 8.\"\"\"\n\n def __init__(self, dims, mode='sintel'):\n self.ht, self.wd = dims[-2:]\n pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8\n pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8\n if mode == 'sintel':\n self._pad = [\n pad_wd // 2,\n pad_wd - pad_wd // 2,\n pad_ht // 2,\n pad_ht - pad_ht // 2,\n ]\n else:\n self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht]\n\n def pad(self, *inputs):\n return [F.pad(x, self._pad, mode='replicate') for x in inputs]\n\n def unpad(self, x):\n ht, wd = x.shape[-2:]\n c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]]\n return x[..., c[0] : c[1], c[2] : c[3]]\n\n\ndef forward_interpolate(flow):\n \"\"\"Interpolate flow map to match the original image size.\"\"\"\n flow = flow.detach().cpu().numpy()\n dx, dy = flow[0], flow[1]\n\n ht, wd = dx.shape\n x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))\n\n x1 = x0 + dx\n y1 = y0 + dy\n\n x1 = x1.reshape(-1)\n y1 = y1.reshape(-1)\n dx = dx.reshape(-1)\n dy = dy.reshape(-1)\n\n valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)\n x1 = x1[valid]\n y1 = y1[valid]\n dx = dx[valid]\n dy = dy[valid]\n\n flow_x = interpolate.griddata(\n (x1, y1), dx, (x0, y0), method='nearest', fill_value=0\n )\n\n flow_y = interpolate.griddata(\n (x1, y1), dy, (x0, y0), method='nearest', fill_value=0\n )\n\n flow = np.stack([flow_x, flow_y], axis=0)\n return torch.from_numpy(flow).float()\n\n\ndef bilinear_sampler(img, coords, mode='bilinear', mask=False):\n \"\"\"Wrapper for grid_sample, uses pixel coordinates.\"\"\"\n del mode\n H, W = img.shape[-2:]\n xgrid, ygrid = coords.split([1, 1], dim=-1)\n xgrid = 2 * xgrid / (W - 1) - 1\n ygrid = 2 * ygrid / (H - 1) - 1\n\n grid = torch.cat([xgrid, ygrid], dim=-1)\n img = F.grid_sample(img, grid, align_corners=True)\n\n if mask:\n mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)\n return img, mask.float()\n\n return img\n\n\ndef coords_grid(batch, ht, wd):\n coords = torch.meshgrid(torch.arange(ht), torch.arange(wd))\n coords = torch.stack(coords[::-1], dim=0).float()\n return coords[None].repeat(batch, 1, 1, 1)\n\n\ndef upflow8(flow, mode='bilinear'):\n new_size = (8 * flow.shape[2], 8 * flow.shape[3])\n return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/inference_batch.py", "content": "\"\"\"\nBatch inference script for optical flow preprocessing using RAFT model.\nProcesses multiple video clips in parallel to generate optical flow data for CVD optimization.\n\"\"\"\n\nimport pandas as pd\nimport os\nimport argparse\nimport concurrent.futures\nfrom multiprocessing import Manager\nimport subprocess\nimport queue\nfrom tqdm import tqdm\n\n\ndef process_single_row(row, index, args, worker_id=0):\n \"\"\"Process a single video clip for optical flow generation.\"\"\"\n dir_path = os.path.join(args.dir_path, row[\"id\"])\n device_id = worker_id % args.gpu_num\n\n # Build command for optical flow preprocessing with RAFT model\n cmd = (\n f\"CUDA_VISIBLE_DEVICES={args.gpu_id[device_id]} python camera_pose_annotation/cvd_opt/preprocess/preprocess_flow.py \"\n f\"--dir_path {dir_path} \"\n f\"--model {args.checkpoints_path}/raft-things.pth \"\n f\"--mixed_precision\"\n )\n process = subprocess.Popen(\n cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE\n )\n stdout, stderr = process.communicate()\n if process.returncode != 0:\n print(f\"Error generating optical flow for {row['id']}: {stderr.decode()}\")\n\n\ndef worker(task_queue, args, worker_id, pbar):\n \"\"\"Worker function for parallel optical flow preprocessing.\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n process_single_row(row, index, args, worker_id)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for optical flow preprocessing.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\"--dir_path\", type=str, default=\"./outputs\")\n parser.add_argument(\"--checkpoints_path\", type=str, default=\"./checkpoints\")\n parser.add_argument(\n \"--gpu_id\", type=str, default=\"0\", help=\"Comma-separated list of GPU IDs to use\"\n )\n parser.add_argument(\n \"--num_workers\",\n type=int,\n default=4,\n help=\"Number of workers for parallel processing\",\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n\n # Parse GPU configuration\n args.gpu_num = len(args.gpu_id.split(\",\"))\n args.gpu_id = [int(gpu) for gpu in args.gpu_id.split(\",\")]\n\n df = pd.read_csv(args.csv_path)\n\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(df.iterrows(), total=len(df), desc=\"Processing rows\"):\n process_single_row(row, index, args)\n else:\n # Parallel processing with multiple workers\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in df.iterrows():\n task_queue.put((index, row))\n with tqdm(total=len(df), desc=\"Processing rows\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for id in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, args, id, pbar))\n\n for future in concurrent.futures.as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/cvd_opt/preprocess/preprocess_flow.py", "content": "# Copyright 2025 DeepMind Technologies Limited\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\n\"\"\"Preprocess flow for MegaSaM.\"\"\"\n\nimport cv2\nimport tqdm\nimport argparse\nfrom pathlib import Path # pylint: disable=g-importing-member\nfrom core.utils.utils import InputPadder\nfrom core.raft import RAFT\nimport glob\nimport os\nimport sys\n\nimport numpy as np\nimport torch\n\n\ndef warp_flow(img, flow):\n h, w = flow.shape[:2]\n flow_new = flow.copy()\n flow_new[:, :, 0] += np.arange(w)\n flow_new[:, :, 1] += np.arange(h)[:, np.newaxis]\n\n res = cv2.remap(\n img, flow_new, None, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT\n )\n return res\n\n\ndef resize_flow(flow, img_h, img_w):\n # flow = np.load(flow_path)\n flow_h, flow_w = flow.shape[0], flow.shape[1]\n flow[:, :, 0] *= float(img_w) / float(flow_w)\n flow[:, :, 1] *= float(img_h) / float(flow_h)\n flow = cv2.resize(flow, (img_w, img_h), cv2.INTER_LINEAR)\n\n return flow\n\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--model\", default=\"raft-things.pth\", help=\"restore checkpoint\")\n parser.add_argument(\"--small\", action=\"store_true\", help=\"use small model\")\n parser.add_argument(\"--dir_path\", help=\"dataset for evaluation\")\n parser.add_argument(\n \"--num_heads\",\n default=1,\n type=int,\n help=\"number of heads in attention and aggregation\",\n )\n parser.add_argument(\n \"--position_only\",\n default=False,\n action=\"store_true\",\n help=\"only use position-wise attention\",\n )\n parser.add_argument(\n \"--position_and_content\",\n default=False,\n action=\"store_true\",\n help=\"use position and content-wise attention\",\n )\n parser.add_argument(\n \"--mixed_precision\", action=\"store_true\", help=\"use mixed precision\"\n )\n\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n\n model = torch.nn.DataParallel(RAFT(args))\n model.load_state_dict(torch.load(args.model))\n flow_model = model.module\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n flow_model.to(device).eval()\n\n img_path = os.path.join(args.dir_path, \"img\")\n\n image_list = sorted(glob.glob(os.path.join(img_path, \"*.png\"))) # [::stride]\n image_list += sorted(glob.glob(os.path.join(img_path, \"*.jpg\"))) # [::stride]\n img_data = []\n\n for t, (image_file) in tqdm.tqdm(enumerate(image_list)):\n image = cv2.imread(image_file)[..., ::-1] # rgb\n h0, w0, _ = image.shape\n h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))\n w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))\n image = cv2.resize(image, (w1, h1))\n image = image[: h1 - h1 % 8, : w1 - w1 % 8].transpose(2, 0, 1)\n img_data.append(image)\n\n img_data = np.array(img_data)\n\n flows_low = []\n\n flows_high = []\n flow_masks_high = []\n\n flow_init = None\n flows_arr_low_bwd = {}\n flows_arr_low_fwd = {}\n\n ii = []\n jj = []\n flows_arr_up = []\n masks_arr_up = []\n\n for step in [1, 2, 4, 8, 15]:\n flows_arr_low = []\n for i in tqdm.tqdm(range(max(0, -step), img_data.shape[0] - max(0, step))):\n image1 = (\n torch.as_tensor(np.ascontiguousarray(img_data[i : i + 1]))\n .float()\n .cuda()\n )\n image2 = (\n torch.as_tensor(np.ascontiguousarray(img_data[i + step : i + step + 1]))\n .float()\n .cuda()\n )\n\n ii.append(i)\n jj.append(i + step)\n\n with torch.no_grad():\n padder = InputPadder(image1.shape)\n image1, image2 = padder.pad(image1, image2)\n if np.abs(step) > 1:\n flow_init = np.stack(\n [flows_arr_low_fwd[i], flows_arr_low_bwd[i + step]], axis=0\n )\n flow_init = (\n torch.as_tensor(np.ascontiguousarray(flow_init))\n .float()\n .cuda()\n .permute(0, 3, 1, 2)\n )\n else:\n flow_init = None\n\n flow_low, flow_up, _ = flow_model(\n torch.cat([image1, image2], dim=0),\n torch.cat([image2, image1], dim=0),\n iters=22,\n test_mode=True,\n flow_init=flow_init,\n )\n\n flow_low_fwd = flow_low[0].cpu().numpy().transpose(1, 2, 0)\n flow_low_bwd = flow_low[1].cpu().numpy().transpose(1, 2, 0)\n\n flow_up_fwd = resize_flow(\n flow_up[0].cpu().numpy().transpose(1, 2, 0),\n flow_up.shape[-2] // 2,\n flow_up.shape[-1] // 2,\n )\n flow_up_bwd = resize_flow(\n flow_up[1].cpu().numpy().transpose(1, 2, 0),\n flow_up.shape[-2] // 2,\n flow_up.shape[-1] // 2,\n )\n\n bwd2fwd_flow = warp_flow(flow_up_bwd, flow_up_fwd)\n fwd_lr_error = np.linalg.norm(flow_up_fwd + bwd2fwd_flow, axis=-1)\n fwd_mask_up = fwd_lr_error < 1.0\n\n # flows_arr_low.append(flow_low_fwd)\n flows_arr_low_bwd[i + step] = flow_low_bwd\n flows_arr_low_fwd[i] = flow_low_fwd\n\n # masks_arr_low.append(fwd_mask_low)\n flows_arr_up.append(flow_up_fwd)\n masks_arr_up.append(fwd_mask_up)\n\n iijj = np.stack((ii, jj), axis=0)\n flows_high = np.array(flows_arr_up).transpose(0, 3, 1, 2)\n flow_masks_high = np.array(masks_arr_up)[:, None, ...]\n\n output_path = os.path.join(args.dir_path, \"cache-flow\")\n if not os.path.exists(output_path):\n os.makedirs(output_path)\n\n np.save(os.path.join(output_path, \"flows.npy\"), np.float16(flows_high))\n np.save(os.path.join(output_path, \"flows_masks.npy\"), flow_masks_high)\n np.save(os.path.join(output_path, \"ii-jj.npy\"), iijj)\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/__init__.py", "content": "" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n#\n# This source code is licensed under the Apache License, Version 2.0\n# found in the LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py\n\nfrom functools import partial\nimport math\nimport logging\nfrom typing import Sequence, Tuple, Union, Callable\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.checkpoint\nfrom torch.nn.init import trunc_normal_\n\nfrom .dinov2_layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block\n\n\nlogger = logging.getLogger(\"dinov2\")\n\n\ndef named_apply(fn: Callable, module: nn.Module, name=\"\", depth_first=True, include_root=False) -> nn.Module:\n if not depth_first and include_root:\n fn(module=module, name=name)\n for child_name, child_module in module.named_children():\n child_name = \".\".join((name, child_name)) if name else child_name\n named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)\n if depth_first and include_root:\n fn(module=module, name=name)\n return module\n\n\nclass BlockChunk(nn.ModuleList):\n def forward(self, x):\n for b in self:\n x = b(x)\n return x\n\n\nclass DinoVisionTransformer(nn.Module):\n def __init__(\n self,\n img_size=224,\n patch_size=16,\n in_chans=3,\n embed_dim=768,\n depth=12,\n num_heads=12,\n mlp_ratio=4.0,\n qkv_bias=True,\n ffn_bias=True,\n proj_bias=True,\n drop_path_rate=0.0,\n drop_path_uniform=False,\n init_values=None, # for layerscale: None or 0 => no layerscale\n embed_layer=PatchEmbed,\n act_layer=nn.GELU,\n block_fn=Block,\n ffn_layer=\"mlp\",\n block_chunks=1,\n num_register_tokens=0,\n interpolate_antialias=False,\n interpolate_offset=0.1,\n ):\n \"\"\"\n Args:\n img_size (int, tuple): input image size\n patch_size (int, tuple): patch size\n in_chans (int): number of input channels\n embed_dim (int): embedding dimension\n depth (int): depth of transformer\n num_heads (int): number of attention heads\n mlp_ratio (int): ratio of mlp hidden dim to embedding dim\n qkv_bias (bool): enable bias for qkv if True\n proj_bias (bool): enable bias for proj in attn if True\n ffn_bias (bool): enable bias for ffn if True\n drop_path_rate (float): stochastic depth rate\n drop_path_uniform (bool): apply uniform drop rate across blocks\n weight_init (str): weight init scheme\n init_values (float): layer-scale init values\n embed_layer (nn.Module): patch embedding layer\n act_layer (nn.Module): MLP activation layer\n block_fn (nn.Module): transformer block class\n ffn_layer (str): \"mlp\", \"swiglu\", \"swiglufused\" or \"identity\"\n block_chunks: (int) split block sequence into block_chunks units for FSDP wrap\n num_register_tokens: (int) number of extra cls tokens (so-called \"registers\")\n interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings\n interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings\n \"\"\"\n super().__init__()\n norm_layer = partial(nn.LayerNorm, eps=1e-6)\n\n self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models\n self.num_tokens = 1\n self.n_blocks = depth\n self.num_heads = num_heads\n self.patch_size = patch_size\n self.num_register_tokens = num_register_tokens\n self.interpolate_antialias = interpolate_antialias\n self.interpolate_offset = interpolate_offset\n\n self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)\n num_patches = self.patch_embed.num_patches\n\n self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))\n assert num_register_tokens >= 0\n self.register_tokens = (\n nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None\n )\n\n if drop_path_uniform is True:\n dpr = [drop_path_rate] * depth\n else:\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule\n\n if ffn_layer == \"mlp\":\n logger.info(\"using MLP layer as FFN\")\n ffn_layer = Mlp\n elif ffn_layer == \"swiglufused\" or ffn_layer == \"swiglu\":\n logger.info(\"using SwiGLU layer as FFN\")\n ffn_layer = SwiGLUFFNFused\n elif ffn_layer == \"identity\":\n logger.info(\"using Identity layer as FFN\")\n\n def f(*args, **kwargs):\n return nn.Identity()\n\n ffn_layer = f\n else:\n raise NotImplementedError\n\n blocks_list = [\n block_fn(\n dim=embed_dim,\n num_heads=num_heads,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias,\n proj_bias=proj_bias,\n ffn_bias=ffn_bias,\n drop_path=dpr[i],\n norm_layer=norm_layer,\n act_layer=act_layer,\n ffn_layer=ffn_layer,\n init_values=init_values,\n )\n for i in range(depth)\n ]\n if block_chunks > 0:\n self.chunked_blocks = True\n chunked_blocks = []\n chunksize = depth // block_chunks\n for i in range(0, depth, chunksize):\n # this is to keep the block index consistent if we chunk the block list\n chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])\n self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])\n else:\n self.chunked_blocks = False\n self.blocks = nn.ModuleList(blocks_list)\n\n self.norm = norm_layer(embed_dim)\n self.head = nn.Identity()\n\n self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))\n\n self.init_weights()\n\n def init_weights(self):\n trunc_normal_(self.pos_embed, std=0.02)\n nn.init.normal_(self.cls_token, std=1e-6)\n if self.register_tokens is not None:\n nn.init.normal_(self.register_tokens, std=1e-6)\n named_apply(init_weights_vit_timm, self)\n\n def interpolate_pos_encoding(self, x, w, h):\n previous_dtype = x.dtype\n npatch = x.shape[1] - 1\n N = self.pos_embed.shape[1] - 1\n if npatch == N and w == h:\n return self.pos_embed\n pos_embed = self.pos_embed.float()\n class_pos_embed = pos_embed[:, 0]\n patch_pos_embed = pos_embed[:, 1:]\n dim = x.shape[-1]\n w0 = w // self.patch_size\n h0 = h // self.patch_size\n # we add a small number to avoid floating point error in the interpolation\n # see discussion at https://github.com/facebookresearch/dino/issues/8\n # DINOv2 with register modify the interpolate_offset from 0.1 to 0.0\n w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset\n # w0, h0 = w0 + 0.1, h0 + 0.1\n \n sqrt_N = math.sqrt(N)\n sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N\n patch_pos_embed = nn.functional.interpolate(\n patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),\n scale_factor=(sx, sy),\n # (int(w0), int(h0)), # to solve the upsampling shape issue\n mode=\"bicubic\",\n antialias=self.interpolate_antialias\n )\n \n assert int(w0) == patch_pos_embed.shape[-2]\n assert int(h0) == patch_pos_embed.shape[-1]\n patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)\n return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)\n\n def prepare_tokens_with_masks(self, x, masks=None):\n B, nc, w, h = x.shape\n x = self.patch_embed(x)\n if masks is not None:\n x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)\n\n x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)\n x = x + self.interpolate_pos_encoding(x, w, h)\n\n if self.register_tokens is not None:\n x = torch.cat(\n (\n x[:, :1],\n self.register_tokens.expand(x.shape[0], -1, -1),\n x[:, 1:],\n ),\n dim=1,\n )\n\n return x\n\n def forward_features_list(self, x_list, masks_list):\n x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]\n for blk in self.blocks:\n x = blk(x)\n\n all_x = x\n output = []\n for x, masks in zip(all_x, masks_list):\n x_norm = self.norm(x)\n output.append(\n {\n \"x_norm_clstoken\": x_norm[:, 0],\n \"x_norm_regtokens\": x_norm[:, 1 : self.num_register_tokens + 1],\n \"x_norm_patchtokens\": x_norm[:, self.num_register_tokens + 1 :],\n \"x_prenorm\": x,\n \"masks\": masks,\n }\n )\n return output\n\n def forward_features(self, x, masks=None):\n if isinstance(x, list):\n return self.forward_features_list(x, masks)\n\n x = self.prepare_tokens_with_masks(x, masks)\n\n for blk in self.blocks:\n x = blk(x)\n\n x_norm = self.norm(x)\n return {\n \"x_norm_clstoken\": x_norm[:, 0],\n \"x_norm_regtokens\": x_norm[:, 1 : self.num_register_tokens + 1],\n \"x_norm_patchtokens\": x_norm[:, self.num_register_tokens + 1 :],\n \"x_prenorm\": x,\n \"masks\": masks,\n }\n\n def _get_intermediate_layers_not_chunked(self, x, n=1):\n x = self.prepare_tokens_with_masks(x)\n # If n is an int, take the n last blocks. If it's a list, take them\n output, total_block_len = [], len(self.blocks)\n blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n\n for i, blk in enumerate(self.blocks):\n x = blk(x)\n if i in blocks_to_take:\n output.append(x)\n assert len(output) == len(blocks_to_take), f\"only {len(output)} / {len(blocks_to_take)} blocks found\"\n return output\n\n def _get_intermediate_layers_chunked(self, x, n=1):\n x = self.prepare_tokens_with_masks(x)\n output, i, total_block_len = [], 0, len(self.blocks[-1])\n # If n is an int, take the n last blocks. If it's a list, take them\n blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n\n for block_chunk in self.blocks:\n for blk in block_chunk[i:]: # Passing the nn.Identity()\n x = blk(x)\n if i in blocks_to_take:\n output.append(x)\n i += 1\n assert len(output) == len(blocks_to_take), f\"only {len(output)} / {len(blocks_to_take)} blocks found\"\n return output\n\n def get_intermediate_layers(\n self,\n x: torch.Tensor,\n n: Union[int, Sequence] = 1, # Layers or n last layers to take\n reshape: bool = False,\n return_class_token: bool = False,\n norm=True\n ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:\n if self.chunked_blocks:\n outputs = self._get_intermediate_layers_chunked(x, n)\n else:\n outputs = self._get_intermediate_layers_not_chunked(x, n)\n if norm:\n outputs = [self.norm(out) for out in outputs]\n class_tokens = [out[:, 0] for out in outputs]\n outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]\n if reshape:\n B, _, w, h = x.shape\n outputs = [\n out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()\n for out in outputs\n ]\n if return_class_token:\n return tuple(zip(outputs, class_tokens))\n return tuple(outputs)\n\n def forward(self, *args, is_training=False, **kwargs):\n ret = self.forward_features(*args, **kwargs)\n if is_training:\n return ret\n else:\n return self.head(ret[\"x_norm_clstoken\"])\n\n\ndef init_weights_vit_timm(module: nn.Module, name: str = \"\"):\n \"\"\"ViT weight initialization, original timm impl (for reproducibility)\"\"\"\n if isinstance(module, nn.Linear):\n trunc_normal_(module.weight, std=0.02)\n if module.bias is not None:\n nn.init.zeros_(module.bias)\n\n\ndef vit_small(patch_size=16, num_register_tokens=0, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=384,\n depth=12,\n num_heads=6,\n mlp_ratio=4,\n block_fn=partial(Block, attn_class=MemEffAttention),\n num_register_tokens=num_register_tokens,\n **kwargs,\n )\n return model\n\n\ndef vit_base(patch_size=16, num_register_tokens=0, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=768,\n depth=12,\n num_heads=12,\n mlp_ratio=4,\n block_fn=partial(Block, attn_class=MemEffAttention),\n num_register_tokens=num_register_tokens,\n **kwargs,\n )\n return model\n\n\ndef vit_large(patch_size=16, num_register_tokens=0, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=1024,\n depth=24,\n num_heads=16,\n mlp_ratio=4,\n block_fn=partial(Block, attn_class=MemEffAttention),\n num_register_tokens=num_register_tokens,\n **kwargs,\n )\n return model\n\n\ndef vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):\n \"\"\"\n Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64\n \"\"\"\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=1536,\n depth=40,\n num_heads=24,\n mlp_ratio=4,\n block_fn=partial(Block, attn_class=MemEffAttention),\n num_register_tokens=num_register_tokens,\n **kwargs,\n )\n return model\n\n\ndef DINOv2(model_name):\n model_zoo = {\n \"vits\": vit_small, \n \"vitb\": vit_base, \n \"vitl\": vit_large, \n \"vitg\": vit_giant2\n }\n \n return model_zoo[model_name](\n img_size=518,\n patch_size=14,\n init_values=1.0,\n ffn_layer=\"mlp\" if model_name != \"vitg\" else \"swiglufused\",\n block_chunks=0,\n num_register_tokens=0,\n interpolate_antialias=False,\n interpolate_offset=0.1\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom .mlp import Mlp\nfrom .patch_embed import PatchEmbed\nfrom .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused\nfrom .block import NestedTensorBlock\nfrom .attention import MemEffAttention\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/attention.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py\n\nimport logging\n\nfrom torch import Tensor\nfrom torch import nn\n\n\nlogger = logging.getLogger(\"dinov2\")\n\n\ntry:\n from xformers.ops import memory_efficient_attention, unbind, fmha\n\n XFORMERS_AVAILABLE = True\nexcept ImportError:\n logger.warning(\"xFormers not available\")\n XFORMERS_AVAILABLE = False\n\n\nclass Attention(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int = 8,\n qkv_bias: bool = False,\n proj_bias: bool = True,\n attn_drop: float = 0.0,\n proj_drop: float = 0.0,\n ) -> None:\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.scale = head_dim**-0.5\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(dim, dim, bias=proj_bias)\n self.proj_drop = nn.Dropout(proj_drop)\n\n def forward(self, x: Tensor) -> Tensor:\n B, N, C = x.shape\n qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n\n q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]\n attn = q @ k.transpose(-2, -1)\n\n attn = attn.softmax(dim=-1)\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n\n\nclass MemEffAttention(Attention):\n def forward(self, x: Tensor, attn_bias=None) -> Tensor:\n if not XFORMERS_AVAILABLE:\n assert attn_bias is None, \"xFormers is required for nested tensors usage\"\n return super().forward(x)\n\n B, N, C = x.shape\n qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)\n\n q, k, v = unbind(qkv, 2)\n\n x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)\n x = x.reshape([B, N, C])\n\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n\n " }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/block.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py\n\nimport logging\nfrom typing import Callable, List, Any, Tuple, Dict\n\nimport torch\nfrom torch import nn, Tensor\n\nfrom .attention import Attention, MemEffAttention\nfrom .drop_path import DropPath\nfrom .layer_scale import LayerScale\nfrom .mlp import Mlp\n\n\nlogger = logging.getLogger(\"dinov2\")\n\n\ntry:\n from xformers.ops import fmha\n from xformers.ops import scaled_index_add, index_select_cat\n\n XFORMERS_AVAILABLE = True\nexcept ImportError:\n logger.warning(\"xFormers not available\")\n XFORMERS_AVAILABLE = False\n\n\nclass Block(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int,\n mlp_ratio: float = 4.0,\n qkv_bias: bool = False,\n proj_bias: bool = True,\n ffn_bias: bool = True,\n drop: float = 0.0,\n attn_drop: float = 0.0,\n init_values=None,\n drop_path: float = 0.0,\n act_layer: Callable[..., nn.Module] = nn.GELU,\n norm_layer: Callable[..., nn.Module] = nn.LayerNorm,\n attn_class: Callable[..., nn.Module] = Attention,\n ffn_layer: Callable[..., nn.Module] = Mlp,\n ) -> None:\n super().__init__()\n # print(f\"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}\")\n self.norm1 = norm_layer(dim)\n self.attn = attn_class(\n dim,\n num_heads=num_heads,\n qkv_bias=qkv_bias,\n proj_bias=proj_bias,\n attn_drop=attn_drop,\n proj_drop=drop,\n )\n self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()\n self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = ffn_layer(\n in_features=dim,\n hidden_features=mlp_hidden_dim,\n act_layer=act_layer,\n drop=drop,\n bias=ffn_bias,\n )\n self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()\n self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n self.sample_drop_ratio = drop_path\n\n def forward(self, x: Tensor) -> Tensor:\n def attn_residual_func(x: Tensor) -> Tensor:\n return self.ls1(self.attn(self.norm1(x)))\n\n def ffn_residual_func(x: Tensor) -> Tensor:\n return self.ls2(self.mlp(self.norm2(x)))\n\n if self.training and self.sample_drop_ratio > 0.1:\n # the overhead is compensated only for a drop path rate larger than 0.1\n x = drop_add_residual_stochastic_depth(\n x,\n residual_func=attn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n )\n x = drop_add_residual_stochastic_depth(\n x,\n residual_func=ffn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n )\n elif self.training and self.sample_drop_ratio > 0.0:\n x = x + self.drop_path1(attn_residual_func(x))\n x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2\n else:\n x = x + attn_residual_func(x)\n x = x + ffn_residual_func(x)\n return x\n\n\ndef drop_add_residual_stochastic_depth(\n x: Tensor,\n residual_func: Callable[[Tensor], Tensor],\n sample_drop_ratio: float = 0.0,\n) -> Tensor:\n # 1) extract subset using permutation\n b, n, d = x.shape\n sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)\n brange = (torch.randperm(b, device=x.device))[:sample_subset_size]\n x_subset = x[brange]\n\n # 2) apply residual_func to get residual\n residual = residual_func(x_subset)\n\n x_flat = x.flatten(1)\n residual = residual.flatten(1)\n\n residual_scale_factor = b / sample_subset_size\n\n # 3) add the residual\n x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)\n return x_plus_residual.view_as(x)\n\n\ndef get_branges_scales(x, sample_drop_ratio=0.0):\n b, n, d = x.shape\n sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)\n brange = (torch.randperm(b, device=x.device))[:sample_subset_size]\n residual_scale_factor = b / sample_subset_size\n return brange, residual_scale_factor\n\n\ndef add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):\n if scaling_vector is None:\n x_flat = x.flatten(1)\n residual = residual.flatten(1)\n x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)\n else:\n x_plus_residual = scaled_index_add(\n x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor\n )\n return x_plus_residual\n\n\nattn_bias_cache: Dict[Tuple, Any] = {}\n\n\ndef get_attn_bias_and_cat(x_list, branges=None):\n \"\"\"\n this will perform the index select, cat the tensors, and provide the attn_bias from cache\n \"\"\"\n batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]\n all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))\n if all_shapes not in attn_bias_cache.keys():\n seqlens = []\n for b, x in zip(batch_sizes, x_list):\n for _ in range(b):\n seqlens.append(x.shape[1])\n attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)\n attn_bias._batch_sizes = batch_sizes\n attn_bias_cache[all_shapes] = attn_bias\n\n if branges is not None:\n cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])\n else:\n tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)\n cat_tensors = torch.cat(tensors_bs1, dim=1)\n\n return attn_bias_cache[all_shapes], cat_tensors\n\n\ndef drop_add_residual_stochastic_depth_list(\n x_list: List[Tensor],\n residual_func: Callable[[Tensor, Any], Tensor],\n sample_drop_ratio: float = 0.0,\n scaling_vector=None,\n) -> Tensor:\n # 1) generate random set of indices for dropping samples in the batch\n branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]\n branges = [s[0] for s in branges_scales]\n residual_scale_factors = [s[1] for s in branges_scales]\n\n # 2) get attention bias and index+concat the tensors\n attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)\n\n # 3) apply residual_func to get residual, and split the result\n residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore\n\n outputs = []\n for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):\n outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))\n return outputs\n\n\nclass NestedTensorBlock(Block):\n def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:\n \"\"\"\n x_list contains a list of tensors to nest together and run\n \"\"\"\n assert isinstance(self.attn, MemEffAttention)\n\n if self.training and self.sample_drop_ratio > 0.0:\n\n def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.attn(self.norm1(x), attn_bias=attn_bias)\n\n def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.mlp(self.norm2(x))\n\n x_list = drop_add_residual_stochastic_depth_list(\n x_list,\n residual_func=attn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,\n )\n x_list = drop_add_residual_stochastic_depth_list(\n x_list,\n residual_func=ffn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,\n )\n return x_list\n else:\n\n def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))\n\n def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:\n return self.ls2(self.mlp(self.norm2(x)))\n\n attn_bias, x = get_attn_bias_and_cat(x_list)\n x = x + attn_residual_func(x, attn_bias=attn_bias)\n x = x + ffn_residual_func(x)\n return attn_bias.split(x)\n\n def forward(self, x_or_x_list):\n if isinstance(x_or_x_list, Tensor):\n return super().forward(x_or_x_list)\n elif isinstance(x_or_x_list, list):\n assert XFORMERS_AVAILABLE, \"Please install xFormers for nested tensors usage\"\n return self.forward_nested(x_or_x_list)\n else:\n raise AssertionError\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/drop_path.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py\n\n\nfrom torch import nn\n\n\ndef drop_path(x, drop_prob: float = 0.0, training: bool = False):\n if drop_prob == 0.0 or not training:\n return x\n keep_prob = 1 - drop_prob\n shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets\n random_tensor = x.new_empty(shape).bernoulli_(keep_prob)\n if keep_prob > 0.0:\n random_tensor.div_(keep_prob)\n output = x * random_tensor\n return output\n\n\nclass DropPath(nn.Module):\n \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\"\"\"\n\n def __init__(self, drop_prob=None):\n super(DropPath, self).__init__()\n self.drop_prob = drop_prob\n\n def forward(self, x):\n return drop_path(x, self.drop_prob, self.training)\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/layer_scale.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110\n\nfrom typing import Union\n\nimport torch\nfrom torch import Tensor\nfrom torch import nn\n\n\nclass LayerScale(nn.Module):\n def __init__(\n self,\n dim: int,\n init_values: Union[float, Tensor] = 1e-5,\n inplace: bool = False,\n ) -> None:\n super().__init__()\n self.inplace = inplace\n self.gamma = nn.Parameter(init_values * torch.ones(dim))\n\n def forward(self, x: Tensor) -> Tensor:\n return x.mul_(self.gamma) if self.inplace else x * self.gamma\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/mlp.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py\n\n\nfrom typing import Callable, Optional\n\nfrom torch import Tensor, nn\n\n\nclass Mlp(nn.Module):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = nn.GELU,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x: Tensor) -> Tensor:\n x = self.fc1(x)\n x = self.act(x)\n x = self.drop(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/patch_embed.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py\n\nfrom typing import Callable, Optional, Tuple, Union\n\nfrom torch import Tensor\nimport torch.nn as nn\n\n\ndef make_2tuple(x):\n if isinstance(x, tuple):\n assert len(x) == 2\n return x\n\n assert isinstance(x, int)\n return (x, x)\n\n\nclass PatchEmbed(nn.Module):\n \"\"\"\n 2D image to patch embedding: (B,C,H,W) -> (B,N,D)\n\n Args:\n img_size: Image size.\n patch_size: Patch token size.\n in_chans: Number of input image channels.\n embed_dim: Number of linear projection output channels.\n norm_layer: Normalization layer.\n \"\"\"\n\n def __init__(\n self,\n img_size: Union[int, Tuple[int, int]] = 224,\n patch_size: Union[int, Tuple[int, int]] = 16,\n in_chans: int = 3,\n embed_dim: int = 768,\n norm_layer: Optional[Callable] = None,\n flatten_embedding: bool = True,\n ) -> None:\n super().__init__()\n\n image_HW = make_2tuple(img_size)\n patch_HW = make_2tuple(patch_size)\n patch_grid_size = (\n image_HW[0] // patch_HW[0],\n image_HW[1] // patch_HW[1],\n )\n\n self.img_size = image_HW\n self.patch_size = patch_HW\n self.patches_resolution = patch_grid_size\n self.num_patches = patch_grid_size[0] * patch_grid_size[1]\n\n self.in_chans = in_chans\n self.embed_dim = embed_dim\n\n self.flatten_embedding = flatten_embedding\n\n self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)\n self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()\n\n def forward(self, x: Tensor) -> Tensor:\n _, _, H, W = x.shape\n patch_H, patch_W = self.patch_size\n\n assert H % patch_H == 0, f\"Input image height {H} is not a multiple of patch height {patch_H}\"\n assert W % patch_W == 0, f\"Input image width {W} is not a multiple of patch width: {patch_W}\"\n\n x = self.proj(x) # B C H W\n H, W = x.size(2), x.size(3)\n x = x.flatten(2).transpose(1, 2) # B HW C\n x = self.norm(x)\n if not self.flatten_embedding:\n x = x.reshape(-1, H, W, self.embed_dim) # B H W C\n return x\n\n def flops(self) -> float:\n Ho, Wo = self.patches_resolution\n flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])\n if self.norm is not None:\n flops += Ho * Wo * self.embed_dim\n return flops\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dinov2_layers/swiglu_ffn.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Callable, Optional\n\nfrom torch import Tensor, nn\nimport torch.nn.functional as F\n\n\nclass SwiGLUFFN(nn.Module):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = None,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)\n self.w3 = nn.Linear(hidden_features, out_features, bias=bias)\n\n def forward(self, x: Tensor) -> Tensor:\n x12 = self.w12(x)\n x1, x2 = x12.chunk(2, dim=-1)\n hidden = F.silu(x1) * x2\n return self.w3(hidden)\n\n\ntry:\n from xformers.ops import SwiGLU\n\n XFORMERS_AVAILABLE = True\nexcept ImportError:\n SwiGLU = SwiGLUFFN\n XFORMERS_AVAILABLE = False\n\n\nclass SwiGLUFFNFused(SwiGLU):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = None,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8\n super().__init__(\n in_features=in_features,\n hidden_features=hidden_features,\n out_features=out_features,\n bias=bias,\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/dpt.py", "content": "import cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torchvision.transforms import Compose\n\nfrom .dinov2 import DINOv2\nfrom .util.blocks import FeatureFusionBlock, _make_scratch\nfrom .util.transform import Resize, NormalizeImage, PrepareForNet\n\n\ndef _make_fusion_block(features, use_bn, size=None):\n return FeatureFusionBlock(\n features,\n nn.ReLU(False),\n deconv=False,\n bn=use_bn,\n expand=False,\n align_corners=True,\n size=size,\n )\n\n\nclass ConvBlock(nn.Module):\n def __init__(self, in_feature, out_feature):\n super().__init__()\n \n self.conv_block = nn.Sequential(\n nn.Conv2d(in_feature, out_feature, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(out_feature),\n nn.ReLU(True)\n )\n \n def forward(self, x):\n return self.conv_block(x)\n\n\nclass DPTHead(nn.Module):\n def __init__(\n self, \n in_channels, \n features=256, \n use_bn=False, \n out_channels=[256, 512, 1024, 1024], \n use_clstoken=False\n ):\n super(DPTHead, self).__init__()\n \n self.use_clstoken = use_clstoken\n \n self.projects = nn.ModuleList([\n nn.Conv2d(\n in_channels=in_channels,\n out_channels=out_channel,\n kernel_size=1,\n stride=1,\n padding=0,\n ) for out_channel in out_channels\n ])\n \n self.resize_layers = nn.ModuleList([\n nn.ConvTranspose2d(\n in_channels=out_channels[0],\n out_channels=out_channels[0],\n kernel_size=4,\n stride=4,\n padding=0),\n nn.ConvTranspose2d(\n in_channels=out_channels[1],\n out_channels=out_channels[1],\n kernel_size=2,\n stride=2,\n padding=0),\n nn.Identity(),\n nn.Conv2d(\n in_channels=out_channels[3],\n out_channels=out_channels[3],\n kernel_size=3,\n stride=2,\n padding=1)\n ])\n \n if use_clstoken:\n self.readout_projects = nn.ModuleList()\n for _ in range(len(self.projects)):\n self.readout_projects.append(\n nn.Sequential(\n nn.Linear(2 * in_channels, in_channels),\n nn.GELU()))\n \n self.scratch = _make_scratch(\n out_channels,\n features,\n groups=1,\n expand=False,\n )\n \n self.scratch.stem_transpose = None\n \n self.scratch.refinenet1 = _make_fusion_block(features, use_bn)\n self.scratch.refinenet2 = _make_fusion_block(features, use_bn)\n self.scratch.refinenet3 = _make_fusion_block(features, use_bn)\n self.scratch.refinenet4 = _make_fusion_block(features, use_bn)\n \n head_features_1 = features\n head_features_2 = 32\n \n self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)\n self.scratch.output_conv2 = nn.Sequential(\n nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),\n nn.ReLU(True),\n nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),\n nn.ReLU(True),\n nn.Identity(),\n )\n \n def forward(self, out_features, patch_h, patch_w):\n out = []\n for i, x in enumerate(out_features):\n if self.use_clstoken:\n x, cls_token = x[0], x[1]\n readout = cls_token.unsqueeze(1).expand_as(x)\n x = self.readout_projects[i](torch.cat((x, readout), -1))\n else:\n x = x[0]\n \n x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))\n \n x = self.projects[i](x)\n x = self.resize_layers[i](x)\n \n out.append(x)\n \n layer_1, layer_2, layer_3, layer_4 = out\n \n layer_1_rn = self.scratch.layer1_rn(layer_1)\n layer_2_rn = self.scratch.layer2_rn(layer_2)\n layer_3_rn = self.scratch.layer3_rn(layer_3)\n layer_4_rn = self.scratch.layer4_rn(layer_4)\n \n path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:]) \n path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])\n path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])\n path_1 = self.scratch.refinenet1(path_2, layer_1_rn)\n \n out = self.scratch.output_conv1(path_1)\n out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode=\"bilinear\", align_corners=True)\n out = self.scratch.output_conv2(out)\n \n return out\n\n\nclass DepthAnythingV2(nn.Module):\n def __init__(\n self, \n encoder='vitl', \n features=256, \n out_channels=[256, 512, 1024, 1024], \n use_bn=False, \n use_clstoken=False\n ):\n super(DepthAnythingV2, self).__init__()\n \n self.intermediate_layer_idx = {\n 'vits': [2, 5, 8, 11],\n 'vitb': [2, 5, 8, 11], \n 'vitl': [4, 11, 17, 23], \n 'vitg': [9, 19, 29, 39]\n }\n \n self.encoder = encoder\n self.pretrained = DINOv2(model_name=encoder)\n \n self.depth_head = DPTHead(self.pretrained.embed_dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken)\n \n def forward(self, x):\n patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14\n \n features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], return_class_token=True)\n \n depth = self.depth_head(features, patch_h, patch_w)\n depth = F.relu(depth)\n \n return depth.squeeze(1)\n \n @torch.no_grad()\n def infer_image(self, raw_image, input_size=518):\n image, (h, w) = self.image2tensor(raw_image, input_size)\n \n depth = self.forward(image)\n \n depth = F.interpolate(depth[:, None], (h, w), mode=\"bilinear\", align_corners=True)[0, 0]\n \n return depth.cpu().numpy()\n \n def image2tensor(self, raw_image, input_size=518): \n transform = Compose([\n Resize(\n width=input_size,\n height=input_size,\n resize_target=False,\n keep_aspect_ratio=True,\n ensure_multiple_of=14,\n resize_method='lower_bound',\n image_interpolation_method=cv2.INTER_CUBIC,\n ),\n NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n PrepareForNet(),\n ])\n \n h, w = raw_image.shape[:2]\n \n image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0\n \n image = transform({'image': image})['image']\n image = torch.from_numpy(image).unsqueeze(0)\n \n DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'\n image = image.to(DEVICE)\n \n return image, (h, w)\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/util/blocks.py", "content": "import torch.nn as nn\n\n\ndef _make_scratch(in_shape, out_shape, groups=1, expand=False):\n scratch = nn.Module()\n\n out_shape1 = out_shape\n out_shape2 = out_shape\n out_shape3 = out_shape\n if len(in_shape) >= 4:\n out_shape4 = out_shape\n\n if expand:\n out_shape1 = out_shape\n out_shape2 = out_shape * 2\n out_shape3 = out_shape * 4\n if len(in_shape) >= 4:\n out_shape4 = out_shape * 8\n\n scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n if len(in_shape) >= 4:\n scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n\n return scratch\n\n\nclass ResidualConvUnit(nn.Module):\n \"\"\"Residual convolution module.\n \"\"\"\n\n def __init__(self, features, activation, bn):\n \"\"\"Init.\n\n Args:\n features (int): number of features\n \"\"\"\n super().__init__()\n\n self.bn = bn\n\n self.groups=1\n\n self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)\n \n self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)\n\n if self.bn == True:\n self.bn1 = nn.BatchNorm2d(features)\n self.bn2 = nn.BatchNorm2d(features)\n\n self.activation = activation\n\n self.skip_add = nn.quantized.FloatFunctional()\n\n def forward(self, x):\n \"\"\"Forward pass.\n\n Args:\n x (tensor): input\n\n Returns:\n tensor: output\n \"\"\"\n \n out = self.activation(x)\n out = self.conv1(out)\n if self.bn == True:\n out = self.bn1(out)\n \n out = self.activation(out)\n out = self.conv2(out)\n if self.bn == True:\n out = self.bn2(out)\n\n if self.groups > 1:\n out = self.conv_merge(out)\n\n return self.skip_add.add(out, x)\n\n\nclass FeatureFusionBlock(nn.Module):\n \"\"\"Feature fusion block.\n \"\"\"\n\n def __init__(\n self, \n features, \n activation, \n deconv=False, \n bn=False, \n expand=False, \n align_corners=True,\n size=None\n ):\n \"\"\"Init.\n \n Args:\n features (int): number of features\n \"\"\"\n super(FeatureFusionBlock, self).__init__()\n\n self.deconv = deconv\n self.align_corners = align_corners\n\n self.groups=1\n\n self.expand = expand\n out_features = features\n if self.expand == True:\n out_features = features // 2\n \n self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)\n\n self.resConfUnit1 = ResidualConvUnit(features, activation, bn)\n self.resConfUnit2 = ResidualConvUnit(features, activation, bn)\n \n self.skip_add = nn.quantized.FloatFunctional()\n\n self.size=size\n\n def forward(self, *xs, size=None):\n \"\"\"Forward pass.\n\n Returns:\n tensor: output\n \"\"\"\n output = xs[0]\n\n if len(xs) == 2:\n res = self.resConfUnit1(xs[1])\n output = self.skip_add.add(output, res)\n\n output = self.resConfUnit2(output)\n\n if (size is None) and (self.size is None):\n modifier = {\"scale_factor\": 2}\n elif size is None:\n modifier = {\"size\": self.size}\n else:\n modifier = {\"size\": size}\n\n output = nn.functional.interpolate(output, **modifier, mode=\"bilinear\", align_corners=self.align_corners)\n \n output = self.out_conv(output)\n\n return output\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/depth_anything_v2/util/transform.py", "content": "import numpy as np\nimport cv2\n\n\nclass Resize(object):\n \"\"\"Resize sample to given size (width, height).\n \"\"\"\n\n def __init__(\n self,\n width,\n height,\n resize_target=True,\n keep_aspect_ratio=False,\n ensure_multiple_of=1,\n resize_method=\"lower_bound\",\n image_interpolation_method=cv2.INTER_AREA,\n ):\n \"\"\"Init.\n\n Args:\n width (int): desired output width\n height (int): desired output height\n resize_target (bool, optional):\n True: Resize the full sample (image, mask, target).\n False: Resize image only.\n Defaults to True.\n keep_aspect_ratio (bool, optional):\n True: Keep the aspect ratio of the input sample.\n Output sample might not have the given width and height, and\n resize behaviour depends on the parameter 'resize_method'.\n Defaults to False.\n ensure_multiple_of (int, optional):\n Output width and height is constrained to be multiple of this parameter.\n Defaults to 1.\n resize_method (str, optional):\n \"lower_bound\": Output will be at least as large as the given size.\n \"upper_bound\": Output will be at max as large as the given size. (Output size might be smaller than given size.)\n \"minimal\": Scale as least as possible. (Output size might be smaller than given size.)\n Defaults to \"lower_bound\".\n \"\"\"\n self.__width = width\n self.__height = height\n\n self.__resize_target = resize_target\n self.__keep_aspect_ratio = keep_aspect_ratio\n self.__multiple_of = ensure_multiple_of\n self.__resize_method = resize_method\n self.__image_interpolation_method = image_interpolation_method\n\n def constrain_to_multiple_of(self, x, min_val=0, max_val=None):\n y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)\n\n if max_val is not None and y > max_val:\n y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)\n\n if y < min_val:\n y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)\n\n return y\n\n def get_size(self, width, height):\n # determine new height and width\n scale_height = self.__height / height\n scale_width = self.__width / width\n\n if self.__keep_aspect_ratio:\n if self.__resize_method == \"lower_bound\":\n # scale such that output size is lower bound\n if scale_width > scale_height:\n # fit width\n scale_height = scale_width\n else:\n # fit height\n scale_width = scale_height\n elif self.__resize_method == \"upper_bound\":\n # scale such that output size is upper bound\n if scale_width < scale_height:\n # fit width\n scale_height = scale_width\n else:\n # fit height\n scale_width = scale_height\n elif self.__resize_method == \"minimal\":\n # scale as least as possbile\n if abs(1 - scale_width) < abs(1 - scale_height):\n # fit width\n scale_height = scale_width\n else:\n # fit height\n scale_width = scale_height\n else:\n raise ValueError(f\"resize_method {self.__resize_method} not implemented\")\n\n if self.__resize_method == \"lower_bound\":\n new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)\n new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)\n elif self.__resize_method == \"upper_bound\":\n new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)\n new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)\n elif self.__resize_method == \"minimal\":\n new_height = self.constrain_to_multiple_of(scale_height * height)\n new_width = self.constrain_to_multiple_of(scale_width * width)\n else:\n raise ValueError(f\"resize_method {self.__resize_method} not implemented\")\n\n return (new_width, new_height)\n\n def __call__(self, sample):\n width, height = self.get_size(sample[\"image\"].shape[1], sample[\"image\"].shape[0])\n \n # resize sample\n sample[\"image\"] = cv2.resize(sample[\"image\"], (width, height), interpolation=self.__image_interpolation_method)\n\n if self.__resize_target:\n if \"depth\" in sample:\n sample[\"depth\"] = cv2.resize(sample[\"depth\"], (width, height), interpolation=cv2.INTER_NEAREST)\n \n if \"mask\" in sample:\n sample[\"mask\"] = cv2.resize(sample[\"mask\"].astype(np.float32), (width, height), interpolation=cv2.INTER_NEAREST)\n \n return sample\n\n\nclass NormalizeImage(object):\n \"\"\"Normlize image by given mean and std.\n \"\"\"\n\n def __init__(self, mean, std):\n self.__mean = mean\n self.__std = std\n\n def __call__(self, sample):\n sample[\"image\"] = (sample[\"image\"] - self.__mean) / self.__std\n\n return sample\n\n\nclass PrepareForNet(object):\n \"\"\"Prepare sample for usage as network input.\n \"\"\"\n\n def __init__(self):\n pass\n\n def __call__(self, sample):\n image = np.transpose(sample[\"image\"], (2, 0, 1))\n sample[\"image\"] = np.ascontiguousarray(image).astype(np.float32)\n\n if \"depth\" in sample:\n depth = sample[\"depth\"].astype(np.float32)\n sample[\"depth\"] = np.ascontiguousarray(depth)\n \n if \"mask\" in sample:\n sample[\"mask\"] = sample[\"mask\"].astype(np.float32)\n sample[\"mask\"] = np.ascontiguousarray(sample[\"mask\"])\n \n return sample" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/inference.py", "content": "\"\"\"\nSingle-threaded inference script for Depth-Anything V2 model.\nProcesses images in a directory to generate depth maps sequentially.\n\"\"\"\n\nimport argparse\nimport cv2\nimport glob\nimport numpy as np\nimport os\nimport torch\n\nfrom depth_anything_v2.dpt import DepthAnythingV2\n\n# Model configuration for different encoder variants\nmodel_configs = {\n \"vits\": {\"encoder\": \"vits\", \"features\": 64, \"out_channels\": [48, 96, 192, 384]},\n \"vitb\": {\"encoder\": \"vitb\", \"features\": 128, \"out_channels\": [96, 192, 384, 768]},\n \"vitl\": {\n \"encoder\": \"vitl\",\n \"features\": 256,\n \"out_channels\": [256, 512, 1024, 1024],\n },\n \"vitg\": {\n \"encoder\": \"vitg\",\n \"features\": 384,\n \"out_channels\": [1536, 1536, 1536, 1536],\n },\n}\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for depth estimation.\"\"\"\n parser = argparse.ArgumentParser(description=\"Depth Anything V2\")\n parser.add_argument(\"--input-size\", type=int, default=518)\n parser.add_argument(\"--dir_path\", type=str, default=\"./vis_depth\")\n parser.add_argument(\n \"--encoder\", type=str, default=\"vitl\", choices=[\"vits\", \"vitb\", \"vitl\", \"vitg\"]\n )\n parser.add_argument(\n \"--load-from\",\n type=str,\n default=\"checkpoints/Depth-Anything/depth_anything_v2_vitl.pth\",\n )\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n\n # Auto-detect best available device\n DEVICE = (\n \"cuda\"\n if torch.cuda.is_available()\n else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n )\n print(f\"Using device: {DEVICE}\")\n\n # Initialize Depth-Anything V2 model\n depth_anything = DepthAnythingV2(**model_configs[args.encoder])\n depth_anything.load_state_dict(torch.load(args.load_from, map_location=\"cpu\"))\n depth_anything = depth_anything.to(DEVICE).eval()\n\n # Setup input and output paths\n img_path = os.path.join(args.dir_path, \"img\")\n out_path = os.path.join(args.dir_path, \"depth-anything\")\n\n if not os.path.exists(out_path):\n os.makedirs(out_path)\n\n # Collect all image files\n img_list = sorted(glob.glob(os.path.join(img_path, \"*.jpg\")))\n img_list += sorted(glob.glob(os.path.join(img_path, \"*.png\")))\n\n # Process each image sequentially\n for k, img in enumerate(img_list):\n print(f\"Progress {k+1}/{len(img_list)}: {img}\")\n\n # Load and process image\n raw_image = cv2.imread(img)\n\n # Generate depth map\n depth = depth_anything.infer_image(raw_image, args.input_size)\n\n # Save depth map as numpy array\n output_path = os.path.join(\n out_path, os.path.splitext(os.path.basename(img))[0] + \".npy\"\n )\n np.save(output_path, depth)\n" }, { "path": "camera_pose_annotation/depth_estimation/Depth-Anything/inference_batch.py", "content": "\"\"\"\nDistributed batch inference script for Depth-Anything V2 model.\nProcesses video frames to generate depth maps using distributed computing.\n\"\"\"\n\nimport argparse\nfrom datetime import timedelta\nimport cv2\nimport glob\nimport numpy as np\nimport pandas as pd\nimport os\nimport torch\nimport torch.distributed as dist\nfrom torch.utils.data import Dataset, DataLoader, DistributedSampler\nfrom torchvision.transforms import Compose\nimport torch.nn.functional as F\nfrom torchvision.transforms import ToTensor\nfrom tqdm import tqdm\n\nfrom depth_anything_v2.util.transform import Resize, NormalizeImage, PrepareForNet\nfrom depth_anything_v2.dpt import DepthAnythingV2\n\n# Model configuration for different encoder variants\nmodel_configs = {\n \"vits\": {\"encoder\": \"vits\", \"features\": 64, \"out_channels\": [48, 96, 192, 384]},\n \"vitb\": {\"encoder\": \"vitb\", \"features\": 128, \"out_channels\": [96, 192, 384, 768]},\n \"vitl\": {\n \"encoder\": \"vitl\",\n \"features\": 256,\n \"out_channels\": [256, 512, 1024, 1024],\n },\n \"vitg\": {\n \"encoder\": \"vitg\",\n \"features\": 384,\n \"out_channels\": [1536, 1536, 1536, 1536],\n },\n}\n\n\nclass ImageDataset(Dataset):\n \"\"\"Dataset for loading and preprocessing images for depth estimation.\"\"\"\n\n def __init__(self, img_list, input_size):\n self.img_list = img_list\n self.input_size = input_size\n self.transform = Compose(\n [\n Resize(\n width=input_size,\n height=input_size,\n resize_target=False,\n keep_aspect_ratio=True,\n ensure_multiple_of=14,\n resize_method=\"lower_bound\",\n image_interpolation_method=cv2.INTER_CUBIC,\n ),\n NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n PrepareForNet(),\n ]\n )\n\n def __len__(self):\n return len(self.img_list)\n\n def image2tensor(self, raw_image):\n \"\"\"Convert raw image to tensor format for model input.\"\"\"\n h, w = raw_image.shape[:2]\n\n image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0\n\n image = self.transform({\"image\": image})[\"image\"]\n image = torch.from_numpy(image)\n\n return image, (h, w)\n\n def __getitem__(self, idx):\n \"\"\"Load and preprocess a single image with error handling.\"\"\"\n\n def inner_func(idx):\n img_path = self.img_list[idx]\n raw_image = cv2.imread(img_path)\n\n image, (original_h, original_w) = self.image2tensor(raw_image)\n\n data = {\n \"image\": image,\n \"path\": img_path,\n \"original_size\": (original_h, original_w),\n }\n return data\n\n while True:\n try:\n return inner_func(idx)\n except Exception as e:\n print(f\"e: [{e}], path: {self.img_list[idx]}, try to get next idx\")\n idx += 1\n if idx >= len(self.img_list):\n raise StopIteration\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for depth estimation.\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Depth Anything V2 Distributed Inference\"\n )\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\"--input-size\", type=int, default=518)\n parser.add_argument(\"--output_dir\", type=str, default=\"./output\")\n parser.add_argument(\n \"--encoder\", type=str, default=\"vitl\", choices=[\"vits\", \"vitb\", \"vitl\", \"vitg\"]\n )\n parser.add_argument(\"--checkpoints_path\", type=str, default=\"./checkpoints\")\n parser.add_argument(\"--bs\", type=int, default=8, help=\"Batch size for inference\")\n parser.add_argument(\n \"--num_workers\", type=int, default=4, help=\"Number of data loading workers\"\n )\n return parser.parse_args()\n\n\ndef collate_fn(batch):\n \"\"\"Custom collate function for batching data.\"\"\"\n return_batch = {}\n for key in batch[0].keys():\n if key == \"image\":\n return_batch[key] = torch.stack([item[key] for item in batch], dim=0)\n else:\n return_batch[key] = [item[key] for item in batch]\n return return_batch\n\n\ndef main():\n args = parse_args()\n\n # Initialize distributed environment\n dist.init_process_group(backend=\"nccl\", timeout=timedelta(hours=24))\n local_rank = dist.get_rank()\n torch.cuda.set_device(local_rank)\n DEVICE = f\"cuda:{local_rank}\"\n\n # Load data list from CSV\n df = pd.read_csv(args.csv_path)\n\n img_list = []\n for index, row in tqdm(\n df.iterrows(), total=len(df), desc=\"Loading images\", disable=(local_rank != 0)\n ):\n img_dir = os.path.join(args.output_dir, row[\"id\"], \"img\")\n if not os.path.exists(img_dir):\n print(f\"Image directory not found: {img_dir}\")\n continue\n img_list += sorted(glob.glob(os.path.join(img_dir, \"*.jpg\")))\n img_list += sorted(glob.glob(os.path.join(img_dir, \"*.png\")))\n\n # Create dataset and distributed sampler\n dataset = ImageDataset(img_list, args.input_size)\n sampler = DistributedSampler(\n dataset,\n num_replicas=dist.get_world_size(),\n rank=local_rank,\n shuffle=False,\n drop_last=False,\n )\n dataloader = DataLoader(\n dataset,\n batch_size=args.bs,\n sampler=sampler,\n num_workers=args.num_workers,\n pin_memory=True,\n collate_fn=collate_fn,\n )\n\n # Initialize Depth-Anything V2 model\n depth_anything = DepthAnythingV2(**model_configs[args.encoder])\n load_from = os.path.join(\n args.checkpoints_path, f\"Depth-Anything/depth_anything_v2_{args.encoder}.pth\"\n )\n depth_anything.load_state_dict(torch.load(load_from, map_location=\"cpu\"))\n depth_anything = depth_anything.to(DEVICE).eval()\n\n # Run inference and save depth maps\n with torch.no_grad():\n for batch in tqdm(\n dataloader, desc=\"Depth inference\", disable=(local_rank != 0)\n ):\n images = batch[\"image\"].to(DEVICE)\n original_sizes = batch[\"original_size\"]\n paths = batch[\"path\"]\n\n # Forward pass through depth model\n depth = depth_anything(images)\n\n # Upsample to original image size\n original_h, original_w = original_sizes[0]\n depth = F.interpolate(\n depth[:, None],\n size=(original_h, original_w),\n mode=\"bilinear\",\n align_corners=False,\n )\n\n # Save depth maps as numpy arrays\n for i in range(depth.shape[0]):\n depth_i = depth[i, 0].cpu().numpy()\n img_path = paths[i]\n output_filename = (\n os.path.splitext(os.path.basename(img_path))[0] + \".npy\"\n )\n output_dir = os.path.join(\n os.path.dirname(os.path.dirname(img_path)), \"depth-anything\"\n )\n os.makedirs(output_dir, exist_ok=True)\n output_path = os.path.join(output_dir, output_filename)\n np.save(output_path, depth_i)\n\n dist.destroy_process_group()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/__init__.py", "content": "" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/inference.py", "content": "\"\"\"\nSingle-threaded inference script for UniDepth V2 model.\nProcesses images in a directory to generate depth maps and camera parameters sequentially.\n\"\"\"\n\nimport argparse\nimport glob\nimport os\n\nimport cv2\nimport numpy as np\nfrom PIL import Image\nimport torch\nfrom unidepth.models import UniDepthV2\n\n# Maximum dimension for image resizing\nLONG_DIM = 640\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for UniDepth inference.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--dir_path\", type=str, default=\"./vis_depth\")\n parser.add_argument(\"--load-from\", type=str, default=\"checkpoints/UniDepth\")\n\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n\n # Initialize UniDepth V2 model\n model = UniDepthV2.from_pretrained(args.load_from)\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n model = model.to(device)\n\n # Setup input and output paths\n img_path = os.path.join(args.dir_path, \"img\")\n out_path = os.path.join(args.dir_path, \"unidepth\")\n\n if not os.path.exists(out_path):\n os.makedirs(out_path)\n\n # Collect all image files\n img_list = sorted(glob.glob(os.path.join(img_path, \"*.jpg\")))\n img_list += sorted(glob.glob(os.path.join(img_path, \"*.png\")))\n\n fovs = []\n # Process each image sequentially\n for img_path in img_list:\n # Load and preprocess image\n rgb = np.array(Image.open(img_path))[..., :3]\n\n # Calculate target size maintaining aspect ratio\n if rgb.shape[1] > rgb.shape[0]:\n final_w, final_h = LONG_DIM, int(\n round(LONG_DIM * rgb.shape[0] / rgb.shape[1])\n )\n else:\n final_w, final_h = (\n int(round(LONG_DIM * rgb.shape[1] / rgb.shape[0])),\n LONG_DIM,\n )\n rgb = cv2.resize(rgb, (final_w, final_h), cv2.INTER_AREA)\n\n # Convert to tensor format\n rgb_torch = torch.from_numpy(rgb).permute(2, 0, 1)\n\n # Predict depth and intrinsics\n predictions = model.infer(rgb_torch)\n\n # Calculate FOV (horizontal field of view) from predicted intrinsics\n fov_ = np.rad2deg(\n 2\n * np.arctan(\n predictions[\"depth\"].shape[-1]\n / (2 * predictions[\"intrinsics\"][0, 0, 0].cpu().numpy())\n )\n )\n depth = predictions[\"depth\"][0, 0].cpu().numpy()\n print(fov_)\n fovs.append(fov_)\n\n # Save depth map and FOV\n np.savez(\n os.path.join(out_path, img_path.split(\"/\")[-1][:-4] + \".npz\"),\n depth=np.float32(depth),\n fov=fov_,\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/inference_batch.py", "content": "\"\"\"\nDistributed batch inference script for UniDepth V2 model.\nProcesses video frames to generate depth maps and camera intrinsics using distributed computing.\n\"\"\"\n\nimport argparse\nfrom datetime import timedelta\nimport glob\nimport os\nimport cv2\nimport pandas as pd\nimport numpy as np\nfrom PIL import Image\nimport torch\nimport torch.distributed as dist\nfrom torch.utils.data import Dataset, DataLoader, DistributedSampler\nfrom tqdm import tqdm\nfrom unidepth.models import UniDepthV2\n\n\nclass ImageDataset(Dataset):\n \"\"\"Dataset for loading and preprocessing images for UniDepth inference.\"\"\"\n\n def __init__(self, img_list, input_size):\n self.img_list = img_list\n self.input_size = input_size\n\n def __len__(self):\n return len(self.img_list)\n\n def __getitem__(self, idx):\n \"\"\"Load and preprocess a single image with error handling.\"\"\"\n\n def inner_func(idx):\n img_path = self.img_list[idx]\n rgb = np.array(Image.open(img_path))[..., :3]\n\n h, w = rgb.shape[:2]\n\n # Calculate target size maintaining aspect ratio\n if w > h:\n final_w, final_h = self.input_size, int(round(self.input_size * h / w))\n else:\n final_w, final_h = int(round(self.input_size * w / h)), self.input_size\n\n rgb_resized = cv2.resize(rgb, (final_w, final_h), cv2.INTER_AREA)\n rgb_torch = (\n torch.from_numpy(rgb_resized).permute(2, 0, 1).float()\n ) # Convert to CHW format\n\n return {\n \"image\": rgb_torch,\n \"path\": img_path,\n }\n\n while True:\n try:\n return inner_func(idx)\n except Exception as e:\n print(f\"e: [{e}], path: {self.img_list[idx]}, try to get next idx\")\n idx = (idx + 1) % len(self.img_list)\n if idx >= len(self.img_list):\n raise StopIteration\n\n\ndef collate_fn(batch):\n \"\"\"Custom collate function for batching data.\"\"\"\n return_batch = {}\n for key in batch[0].keys():\n if key == \"image\":\n return_batch[key] = torch.stack([item[key] for item in batch], dim=0)\n else:\n return_batch[key] = [item[key] for item in batch]\n return return_batch\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for UniDepth inference.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\"--output_dir\", type=str, default=\"./output\")\n parser.add_argument(\"--checkpoints_path\", type=str, default=\"./checkpoints\")\n parser.add_argument(\n \"--input_size\", type=int, default=640, help=\"Input size for the model\"\n )\n parser.add_argument(\"--bs\", type=int, default=8, help=\"Inference batch size\")\n parser.add_argument(\n \"--num_workers\", type=int, default=4, help=\"Data loading workers\"\n )\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n\n # Initialize distributed environment\n dist.init_process_group(backend=\"nccl\", timeout=timedelta(hours=24))\n local_rank = dist.get_rank()\n torch.cuda.set_device(local_rank)\n DEVICE = f\"cuda:{local_rank}\"\n\n # Load data list from CSV\n df = pd.read_csv(args.csv_path)\n\n img_list = []\n for index, row in tqdm(\n df.iterrows(), total=len(df), desc=\"Loading images\", disable=local_rank != 0\n ):\n img_dir = os.path.join(args.output_dir, row[\"id\"], \"img\")\n if not os.path.exists(img_dir):\n print(f\"Image directory not found: {img_dir}\")\n continue\n img_list += sorted(glob.glob(os.path.join(img_dir, \"*.jpg\")))\n img_list += sorted(glob.glob(os.path.join(img_dir, \"*.png\")))\n\n # Create dataset and distributed sampler\n dataset = ImageDataset(img_list, args.input_size)\n sampler = DistributedSampler(\n dataset,\n num_replicas=dist.get_world_size(),\n rank=local_rank,\n shuffle=False,\n drop_last=False,\n )\n dataloader = DataLoader(\n dataset,\n batch_size=args.bs,\n sampler=sampler,\n num_workers=args.num_workers,\n pin_memory=True,\n collate_fn=collate_fn,\n )\n\n # Initialize UniDepth V2 model\n load_from = os.path.join(args.checkpoints_path, \"UniDepth\")\n model = UniDepthV2.from_pretrained(load_from)\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n model = model.to(device).eval()\n\n # Run inference and save results\n with torch.no_grad():\n for batch in tqdm(\n dataloader, desc=\"Processing batches\", disable=(local_rank != 0)\n ):\n images = batch[\"image\"].to(device)\n paths = batch[\"path\"]\n\n # Model inference\n predictions = model.infer(images)\n\n # Process results for each sample\n for i in range(len(paths)):\n depth = predictions[\"depth\"][i, 0].cpu().numpy() # [H, W]\n intrinsics = predictions[\"intrinsics\"][i].cpu().numpy()\n focal_length = intrinsics[\n 0, 0\n ] # Assume principal point at center, take fx\n w = depth.shape[-1] # Width\n\n # Calculate FOV (horizontal field of view)\n fov = np.rad2deg(2 * np.arctan(w / (2 * focal_length)))\n\n # Save results\n img_path = paths[i]\n output_filename = (\n os.path.splitext(os.path.basename(img_path))[0] + \".npz\"\n )\n output_dir = os.path.join(\n os.path.dirname(os.path.dirname(img_path)), \"unidepth\"\n )\n os.makedirs(output_dir, exist_ok=True)\n output_path = os.path.join(output_dir, output_filename)\n np.savez(output_path, depth=np.float32(depth), fov=fov)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/_2d3ds.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.pipelines import Compose, PanoCrop, PanoRoll\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass _2D3DS(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 512.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"2D3DS.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"cam2w\", \"camera_params\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n self.resizer = Compose(\n [PanoCrop(), PanoRoll(test_mode=test_mode), self.resizer]\n )\n\n def preprocess(self, results):\n self.resizer.ctx = None\n if self.test_mode:\n for i, seq in enumerate(results[\"sequence_fields\"]):\n results[seq][\"points\"] = results[seq][\"camera\"].reconstruct(\n results[seq][\"depth\"]\n )\n results[seq][\"depth\"] = results[seq][\"points\"][:, -1:]\n results[seq][\"gt_fields\"].add(\"points\")\n return super().preprocess(results)\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/_4dor.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass _4DOR(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n default_fps = 10\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"4DOR.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"si\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/__init__.py", "content": "from ._2d3ds import _2D3DS\nfrom ._4dor import _4DOR\nfrom .a2d2 import A2D2\nfrom .adt import ADT\nfrom .aimotive import aiMotive\nfrom .argoverse import Argoverse\nfrom .argoverse2 import Argoverse2\nfrom .arkit import ARKit\nfrom .ase import ASE\nfrom .base_dataset import BaseDataset\nfrom .bdd import BDD\nfrom .bedlam import BEDLAM\nfrom .behave import Behave\nfrom .blendedmvg import BlendedMVG\nfrom .cityscape import Cityscape\nfrom .ddad import DDAD\nfrom .deep360 import Deep360\nfrom .dense import DENSE\nfrom .diode import DiodeIndoor, DiodeIndoor_F\nfrom .dl3dv import DL3DV\nfrom .driving_stereo import DrivingStereo\nfrom .dtu_rmvd import DTURMVD\nfrom .dummy import Dummy\nfrom .dynamic_replica import DynReplica\nfrom .eden import EDEN\nfrom .eth3d import ETH3D, ETH3D_F\nfrom .eth3d_rmvd import ETH3DRMVD\nfrom .facedepth import FaceDepth\nfrom .flsea import FLSea\nfrom .futurehouse import FutureHouse\nfrom .gibson import Gibson\nfrom .hammer import HAMMER\nfrom .hm3d import HM3D\nfrom .hoi4d import HOI4D\nfrom .hypersim import HyperSim\nfrom .ibims import IBims, IBims_F\nfrom .image_dataset import ImageDataset\nfrom .ken_burns import KenBurns\nfrom .kitti import KITTI, KITTIBenchmark\nfrom .kitti360 import KITTI360\nfrom .kitti_multi import KITTIMulti\nfrom .kitti_rmvd import KITTIRMVD\nfrom .lyft import Lyft\nfrom .mapillary import Mapillary\nfrom .matrix_city import MatrixCity\nfrom .matterport3d import Matterport3D\nfrom .megadepth import MegaDepth\nfrom .megadepth_s import MegaDepthS\nfrom .midair import MidAir\nfrom .mip import MIP\nfrom .ms2 import MS2\nfrom .mvimgnet import MVImgNet\nfrom .mvsynth import MVSynth\nfrom .nerds360 import NeRDS360\nfrom .niantic_mapfree import NianticMapFree\nfrom .nuscenes import Nuscenes\nfrom .nyuv2 import NYUv2Depth\nfrom .point_odyssey import PointOdyssey\nfrom .proteus import Proteus\nfrom .samplers import DistributedSamplerNoDuplicate\nfrom .scannet import ScanNet\nfrom .scannetpp import ScanNetpp, ScanNetpp_F\nfrom .sequence_dataset import SequenceDataset\nfrom .sintel import Sintel\nfrom .sunrgbd import SUNRGBD\nfrom .synscapes import Synscapes\nfrom .tartanair import TartanAir\nfrom .taskonomy import Taskonomy\nfrom .tat_rmvd import TATRMVD\nfrom .theo import Theo\nfrom .unrealstereo4k import UnrealStereo4K\nfrom .urbansyn import UrbanSyn\nfrom .utils import ConcatDataset, collate_fn, get_weights\nfrom .vkitti import VKITTI\nfrom .void import VOID\nfrom .waymo import Waymo\nfrom .wildrgbd import WildRGBD\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/a2d2.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass A2D2(ImageDataset):\n min_depth = 0.01\n max_depth = 120.0\n depth_scale = 256.0\n train_split = \"train_clean.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"a2d2.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(\n intrinsics[os.path.join(*image_filename.split(\"/\")[:2])]\n ).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n # if not self.test_mode:\n # dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/adt.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass ADT(SequenceDataset):\n min_depth = 0.01\n max_depth = 20.0\n depth_scale = 1000.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"ADT.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields, # if not test_mode else [*decode_fields, \"points\"],\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def preprocess(self, results):\n self.resizer.ctx = None\n for i, seq in enumerate(results[\"sequence_fields\"]):\n # Create a mask where the distance from the center is less than H/2\n H, W = results[seq][\"image\"].shape[-2:]\n x = torch.linspace(-W / 2 - 0.5, W / 2 + 0.5, W)\n y = torch.linspace(-H / 2 - 0.5, H / 2 + 0.5, H)\n xv, yv = torch.meshgrid(x, y, indexing=\"xy\")\n distance_from_center = torch.sqrt(xv**2 + yv**2).reshape(1, 1, H, W)\n results[seq][\"validity_mask\"] = distance_from_center < (H / 2) + 20\n results[seq][\"depth_mask\"] = results[seq][\"validity_mask\"].clone()\n results[seq][\"mask_fields\"].add(\"depth_mask\")\n results[seq][\"mask_fields\"].add(\"validity_mask\")\n return super().preprocess(results)\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/aimotive.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass aiMotive(SequenceDataset):\n min_depth = 0.01\n max_depth = 100.0\n depth_scale = 256.0\n default_fps = 10\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"aiMotive.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/argoverse.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass Argoverse(ImageDataset):\n min_depth = 0.05\n max_depth = 120.0\n depth_scale = 256.0\n test_split = \"argo_val.txt\"\n train_split = \"argo_train.txt\"\n intrisics_file = \"argo_intrinsics.json\"\n hdf5_paths = [\"argoverse11.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/argoverse2.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Argoverse2(SequenceDataset):\n min_depth = 0.05\n max_depth = 120.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences_clean.json\"\n hdf5_paths = [f\"AV2_viz.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/arkit.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass ARKit(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"Training.txt\"\n train_split = \"Training.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"ARKitS.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/ase.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass ASE(SequenceDataset):\n min_depth = 0.01\n max_depth = 20.0\n depth_scale = 1000.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"ASE.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def preprocess(self, results):\n self.resizer.ctx = None\n for i, seq in enumerate(results[\"sequence_fields\"]):\n # Create a mask where the distance from the center is less than H/2\n H, W = results[seq][\"image\"].shape[-2:]\n x = torch.linspace(-W / 2 - 0.5, W / 2 + 0.5, W)\n y = torch.linspace(-H / 2 - 0.5, H / 2 + 0.5, H)\n xv, yv = torch.meshgrid(x, y, indexing=\"xy\")\n distance_from_center = torch.sqrt(xv**2 + yv**2).reshape(1, 1, H, W)\n results[seq][\"validity_mask\"] = distance_from_center < (H / 2) + 20\n results[seq][\"mask_fields\"].add(\"validity_mask\")\n\n return super().preprocess(results)\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/base_dataset.py", "content": "import os\nfrom abc import abstractmethod\nfrom copy import deepcopy\nfrom math import ceil, log\nfrom typing import Any, Dict, Tuple\n\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset\n\nimport unidepth.datasets.pipelines as pipelines\nfrom unidepth.utils import (eval_3d, eval_depth, identity, is_main_process,\n recursive_index, sync_tensor_across_gpus)\nfrom unidepth.utils.constants import (IMAGENET_DATASET_MEAN,\n IMAGENET_DATASET_STD,\n OPENAI_DATASET_MEAN, OPENAI_DATASET_STD)\n\n\nclass BaseDataset(Dataset):\n min_depth = 0.01\n max_depth = 1000.0\n\n def __init__(\n self,\n image_shape: Tuple[int, int],\n split_file: str,\n test_mode: bool,\n benchmark: bool,\n normalize: bool,\n augmentations_db: Dict[str, Any],\n resize_method: str,\n mini: float,\n num_copies: int = 1,\n **kwargs,\n ) -> None:\n super().__init__()\n assert normalize in [None, \"imagenet\", \"openai\"]\n\n self.split_file = split_file\n self.test_mode = test_mode\n self.data_root = os.environ[\"DATAROOT\"]\n self.image_shape = image_shape\n self.resize_method = resize_method\n self.mini = mini\n self.num_frames = 1\n self.num_copies = num_copies\n self.metrics_store = {}\n self.metrics_count = {}\n\n if normalize == \"imagenet\":\n self.normalization_stats = {\n \"mean\": torch.tensor(IMAGENET_DATASET_MEAN),\n \"std\": torch.tensor(IMAGENET_DATASET_STD),\n }\n elif normalize == \"openai\":\n self.normalization_stats = {\n \"mean\": torch.tensor(OPENAI_DATASET_MEAN),\n \"std\": torch.tensor(OPENAI_DATASET_STD),\n }\n else:\n self.normalization_stats = {\n \"mean\": torch.tensor([0.0, 0.0, 0.0]),\n \"std\": torch.tensor([1.0, 1.0, 1.0]),\n }\n\n for k, v in augmentations_db.items():\n setattr(self, k, v)\n if not self.test_mode:\n self._augmentation_space()\n\n self.masker = pipelines.AnnotationMask(\n min_value=0.0,\n max_value=self.max_depth if test_mode else None,\n custom_fn=identity,\n )\n self.filler = pipelines.RandomFiller(noise_pad=True)\n\n shape_mult = self.shape_constraints[\"shape_mult\"]\n self.image_shape = [\n ceil(self.image_shape[0] / shape_mult) * shape_mult,\n ceil(self.image_shape[1] / shape_mult) * shape_mult,\n ]\n self.resizer = pipelines.ContextCrop(\n image_shape=self.image_shape,\n train_ctx_range=(1.0 / self.random_scale, 1.0 * self.random_scale),\n test_min_ctx=self.test_context,\n keep_original=test_mode,\n shape_constraints=self.shape_constraints,\n )\n\n self.collecter = pipelines.Collect(\n keys=[\"image_fields\", \"mask_fields\", \"gt_fields\", \"camera_fields\"]\n )\n\n def __len__(self):\n return len(self.dataset)\n\n def pack_batch(self, results):\n results[\"paddings\"] = [\n results[x][\"paddings\"][0] for x in results[\"sequence_fields\"]\n ]\n for fields_name in [\n \"image_fields\",\n \"gt_fields\",\n \"mask_fields\",\n \"camera_fields\",\n ]:\n fields = results.get(fields_name)\n packed = {\n field: torch.cat(\n [results[seq][field] for seq in results[\"sequence_fields\"]]\n )\n for field in fields\n }\n results.update(packed)\n return results\n\n def unpack_batch(self, results):\n for fields_name in [\n \"image_fields\",\n \"gt_fields\",\n \"mask_fields\",\n \"camera_fields\",\n ]:\n fields = results.get(fields_name)\n unpacked = {\n field: {\n seq: results[field][idx : idx + 1]\n for idx, seq in enumerate(results[\"sequence_fields\"])\n }\n for field in fields\n }\n results.update(unpacked)\n return results\n\n def _augmentation_space(self):\n self.augmentations_dict = {\n \"Flip\": pipelines.RandomFlip(prob=self.flip_p),\n \"Jitter\": pipelines.RandomColorJitter(\n (-self.random_jitter, self.random_jitter), prob=self.jitter_p\n ),\n \"Gamma\": pipelines.RandomGamma(\n (-self.random_gamma, self.random_gamma), prob=self.gamma_p\n ),\n \"Blur\": pipelines.GaussianBlur(\n kernel_size=13, sigma=(0.1, self.random_blur), prob=self.blur_p\n ),\n \"Grayscale\": pipelines.RandomGrayscale(prob=self.grayscale_p),\n }\n\n def augment(self, results):\n for name, aug in self.augmentations_dict.items():\n results = aug(results)\n return results\n\n def prepare_depth_eval(self, inputs, preds):\n new_preds = {}\n keyframe_idx = getattr(self, \"keyframe_idx\", None)\n slice_idx = slice(\n keyframe_idx, keyframe_idx + 1 if keyframe_idx is not None else None\n )\n new_gts = inputs[\"depth\"][slice_idx]\n new_masks = inputs[\"depth_mask\"][slice_idx].bool()\n for key, val in preds.items():\n if \"depth\" in key:\n new_preds[key] = val[slice_idx]\n return new_gts, new_preds, new_masks\n\n def prepare_points_eval(self, inputs, preds):\n new_preds = {}\n new_gts = inputs[\"points\"]\n new_masks = inputs[\"depth_mask\"].bool()\n if \"points_mask\" in inputs:\n new_masks = inputs[\"points_mask\"].bool()\n for key, val in preds.items():\n if \"points\" in key:\n new_preds[key] = val\n return new_gts, new_preds, new_masks\n\n def add_points(self, inputs):\n inputs[\"points\"] = inputs.get(\"camera_original\", inputs[\"camera\"]).reconstruct(\n inputs[\"depth\"]\n )\n return inputs\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def accumulate_metrics(\n self,\n inputs,\n preds,\n keyframe_idx=None,\n metrics=[\"depth\", \"points\", \"flow_fwd\", \"pairwise\"],\n ):\n if \"depth\" in inputs and \"points\" not in inputs:\n inputs = self.add_points(inputs)\n\n available_metrics = []\n for metric in metrics:\n metric_in_gt = any((metric in k for k in inputs.keys()))\n metric_in_pred = any((metric in k for k in preds.keys()))\n if metric_in_gt and metric_in_pred:\n available_metrics.append(metric)\n\n if keyframe_idx is not None:\n inputs = recursive_index(inputs, slice(keyframe_idx, keyframe_idx + 1))\n preds = recursive_index(preds, slice(keyframe_idx, keyframe_idx + 1))\n\n if \"depth\" in available_metrics:\n depth_gt, depth_pred, depth_masks = self.prepare_depth_eval(inputs, preds)\n self.accumulate_metrics_depth(depth_gt, depth_pred, depth_masks)\n\n if \"points\" in available_metrics:\n points_gt, points_pred, points_masks = self.prepare_points_eval(\n inputs, preds\n )\n self.accumulate_metrics_3d(points_gt, points_pred, points_masks)\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def accumulate_metrics_depth(self, gts, preds, masks):\n for eval_type, pred in preds.items():\n log_name = eval_type.replace(\"depth\", \"\").strip(\"-\").strip(\"_\")\n if log_name not in self.metrics_store:\n self.metrics_store[log_name] = {}\n current_count = self.metrics_count.get(\n log_name, torch.tensor([], device=gts.device)\n )\n new_count = masks.view(gts.shape[0], -1).sum(dim=-1)\n self.metrics_count[log_name] = torch.cat([current_count, new_count])\n for k, v in eval_depth(gts, pred, masks, max_depth=self.max_depth).items():\n current_metric = self.metrics_store[log_name].get(\n k, torch.tensor([], device=gts.device)\n )\n self.metrics_store[log_name][k] = torch.cat([current_metric, v])\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def accumulate_metrics_3d(self, gts, preds, masks):\n thresholds = torch.linspace(\n log(self.min_depth),\n log(self.max_depth / 20),\n steps=100,\n device=gts.device,\n ).exp()\n for eval_type, pred in preds.items():\n log_name = eval_type.replace(\"points\", \"\").strip(\"-\").strip(\"_\")\n if log_name not in self.metrics_store:\n self.metrics_store[log_name] = {}\n current_count = self.metrics_count.get(\n log_name, torch.tensor([], device=gts.device)\n )\n new_count = masks.view(gts.shape[0], -1).sum(dim=-1)\n self.metrics_count[log_name] = torch.cat([current_count, new_count])\n for k, v in eval_3d(gts, pred, masks, thresholds=thresholds).items():\n current_metric = self.metrics_store[log_name].get(\n k, torch.tensor([], device=gts.device)\n )\n self.metrics_store[log_name][k] = torch.cat([current_metric, v])\n\n def get_evaluation(self, metrics=None):\n metric_vals = {}\n for eval_type in metrics if metrics is not None else self.metrics_store.keys():\n assert self.metrics_store[eval_type]\n cnts = sync_tensor_across_gpus(self.metrics_count[eval_type])\n for name, val in self.metrics_store[eval_type].items():\n # vals_r = (sync_tensor_across_gpus(val) * cnts / cnts.sum()).sum()\n vals_r = sync_tensor_across_gpus(val).mean()\n metric_vals[f\"{eval_type}_{name}\".strip(\"_\")] = np.round(\n vals_r.cpu().item(), 5\n )\n self.metrics_store[eval_type] = {}\n self.metrics_count = {}\n return metric_vals\n\n def replicate(self, results):\n for i in range(1, self.num_copies):\n results[(0, i)] = {k: deepcopy(v) for k, v in results[(0, 0)].items()}\n results[\"sequence_fields\"].append((0, i))\n return results\n\n def log_load_dataset(self):\n if is_main_process():\n info = f\"Loaded {self.__class__.__name__} with {len(self)} images.\"\n print(info)\n\n def pre_pipeline(self, results):\n results[\"image_fields\"] = results.get(\"image_fields\", set())\n results[\"gt_fields\"] = results.get(\"gt_fields\", set())\n results[\"mask_fields\"] = results.get(\"mask_fields\", set())\n results[\"sequence_fields\"] = results.get(\"sequence_fields\", set())\n results[\"camera_fields\"] = results.get(\"camera_fields\", set())\n results[\"dataset_name\"] = (\n [self.__class__.__name__] * self.num_frames * self.num_copies\n )\n results[\"depth_scale\"] = [self.depth_scale] * self.num_frames * self.num_copies\n results[\"si\"] = [False] * self.num_frames * self.num_copies\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n results[\"valid_camera\"] = [True] * self.num_frames * self.num_copies\n results[\"valid_pose\"] = [True] * self.num_frames * self.num_copies\n return results\n\n def eval_mask(self, valid_mask):\n return valid_mask\n\n def chunk(self, dataset, chunk_dim=1, pct=1.0):\n subsampled_datasets = [\n x\n for i in range(0, len(dataset), int(1 / pct * chunk_dim))\n for x in dataset[i : i + chunk_dim]\n ]\n return subsampled_datasets\n\n @abstractmethod\n def preprocess(self, results):\n raise NotImplementedError\n\n @abstractmethod\n def postprocess(self, results):\n raise NotImplementedError\n\n @abstractmethod\n def get_mapper(self):\n raise NotImplementedError\n\n @abstractmethod\n def get_intrinsics(self, idx, image_name):\n raise NotImplementedError\n\n @abstractmethod\n def get_extrinsics(self, idx, image_name):\n raise NotImplementedError\n\n @abstractmethod\n def load_dataset(self):\n raise NotImplementedError\n\n @abstractmethod\n def get_single_item(self, idx, sample=None, mapper=None):\n raise NotImplementedError\n\n @abstractmethod\n def __getitem__(self, idx):\n raise NotImplementedError\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/bdd.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass BDD(ImageDataset):\n min_depth = 0.01\n max_depth = 70.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train_clean.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"BDD.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(\n intrinsics[os.path.join(*image_filename.split(\"/\")[:2])]\n ).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n h5file.close()\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n if self.test_mode and not self.benchmark:\n dataset = self.chunk(dataset, chunk_dim=1, pct=0.1)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"si\"] = [True] * self.num_copies\n results[\"valid_camera\"] = [False] * self.num_copies\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [2] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/bedlam.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass BEDLAM(SequenceDataset):\n min_depth = 0.01\n max_depth = 256.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"val.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"BEDLAM.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/behave.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Behave(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n default_fps = 10\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"Behave.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"si\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/blendedmvg.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass BlendedMVG(SequenceDataset):\n min_depth = 0.01\n max_depth = 5000.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences_clean.json\"\n hdf5_paths = [\"BlendedMVG_.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"si\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/cityscape.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass Cityscape(ImageDataset):\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"cityscape.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"quality\"] = [2] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/ddad.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass DDAD(ImageDataset):\n min_depth = 0.05\n max_depth = 120.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [f\"ddad/ddad_{i}.hdf5\" for i in range(8)]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\").strip(\"\\n\")\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename, chunk_idx = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val, chunk_idx]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n \"K\": 2,\n \"chunk_idx\": 3,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/deep360.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.pipelines import Compose, PanoCrop, PanoRoll\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Deep360(SequenceDataset):\n min_depth = 0.1\n max_depth = 1000.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"Deep360.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"cam2w\", \"camera_params\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n self.resizer = Compose(\n [PanoCrop(), PanoRoll(test_mode=test_mode), self.resizer]\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/dense.py", "content": "import os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass DENSE(ImageDataset):\n CAM_INTRINSIC = {\n \"ALL\": torch.tensor(\n [\n [1177.8614, 0.0, 474.319027],\n [0.0, 1177.8614, 224.275919],\n [0.0, 0.0, 1.0],\n ]\n )\n }\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 255.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n hdf5_paths = [\"DENSE.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.intrisics = {}\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n sample = [image_filename, depth_filename]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_intrinsics(self, idx, image_name):\n return self.CAM_INTRINSIC[\"ALL\"].clone()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/diml.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass DIML(ImageDataset):\n min_depth = 0.01\n max_depth = 100.0\n depth_scale = 256.0\n test_split = \"test.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"DIML.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.intrisics = {}\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(\n intrinsics[image_filename.split(\"/\")[0]]\n ).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n h5file.close()\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"quality\"] = [2] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/diode.py", "content": "import os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass DiodeIndoor(ImageDataset):\n CAM_INTRINSIC = {\n \"ALL\": torch.tensor([[886.81, 0, 512], [0, 927.06, 384], [0, 0, 1]])\n }\n min_depth = 0.01\n max_depth = 25.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n hdf5_paths = [\"DiodeIndoor.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n # load annotations\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n sample = [\n image_filename,\n depth_filename,\n ]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_intrinsics(self, *args, **kwargs):\n return self.CAM_INTRINSIC[\"ALL\"].clone()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n\n\nclass DiodeIndoor_F(SequenceDataset):\n min_depth = 0.01\n max_depth = 25.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"DiodeIndoor-F.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, float],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=(\n decode_fields if not test_mode else [*decode_fields, \"points\"]\n ),\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n\n\nclass DiodeOutdoor(ImageDataset):\n CAM_INTRINSIC = {\n \"ALL\": torch.tensor([[886.81, 0, 512], [0, 927.06, 384], [0, 0, 1]])\n }\n min_depth = 0.1\n max_depth = 80.0\n log_mean = 0\n log_std = 1\n test_split = \"diode_outdoor_val.txt\"\n train_split = \"diode_outdoor_train.txt\"\n hdf5_paths = [\"diode.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n depth_scale=256,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.depth_scale = depth_scale\n\n self.masker = AnnotationMask(\n min_value=self.min_depth,\n max_value=self.max_depth if test_mode else None,\n custom_fn=self.eval_mask if test_mode else lambda x, *args, **kwargs: x,\n )\n # load annotations\n self.load_dataset()\n\n def load_dataset(self):\n self.h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_path),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(self.h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1]\n dataset = {\"depth_filename\": [], \"image_filename\": []}\n for line in txt_string.split(\"\\n\"):\n depth_filename = line.strip().split(\" \")[1]\n img_name = line.strip().split(\" \")[0]\n image_filename = img_name\n dataset[\"depth_filename\"].append(depth_filename)\n dataset[\"image_filename\"].append(image_filename)\n\n self.dataset = pl.from_dict(dataset)\n\n if not self.test_mode and self.mini:\n self.dataset = self.dataset[::2]\n\n\nclass Diode(ImageDataset):\n CAM_INTRINSIC = {\n \"ALL\": torch.tensor([[886.81, 0, 512], [0, 927.06, 384], [0, 0, 1]])\n }\n log_mean = 0\n log_std = 1\n min_depth = 0.6\n max_depth = 80.0\n test_split = \"diode_val.txt\"\n train_split = \"diode_train.txt\"\n hdf5_paths = [\"diode.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n depth_scale=256,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.depth_scale = depth_scale\n\n self.masker = AnnotationMask(\n min_value=self.min_depth,\n max_value=self.max_depth if test_mode else None,\n custom_fn=self.eval_mask if test_mode else lambda x, *args, **kwargs: x,\n )\n # load annotations\n self.load_dataset()\n\n def load_dataset(self):\n self.h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_path),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(self.h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1]\n dataset = {\"depth_filename\": [], \"image_filename\": []}\n for line in txt_string.split(\"\\n\"):\n depth_filename = line.strip().split(\" \")[1]\n image_filename = line.strip().split(\" \")[0]\n dataset[\"depth_filename\"].append(depth_filename)\n dataset[\"image_filename\"].append(image_filename)\n\n self.dataset = pl.from_dict(dataset)\n\n if not self.test_mode and self.mini:\n self.dataset = self.dataset[::2]\n\n def get_intrinsics(self, *args, **kwargs):\n return self.CAM_INTRINSIC[\"ALL\"].clone()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/dl3dv.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass DL3DV(SequenceDataset):\n min_depth = 0.001\n max_depth = 250.0\n depth_scale = 512.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"DL3DVcv.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/driving_stereo.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass DrivingStereo(ImageDataset):\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n test_split = \"drivingstereo_val.txt\"\n train_split = \"drivingstereo_train.txt\"\n intrisics_file = \"drivingstereo_intrinsics.json\"\n hdf5_paths = [\"DrivingStereo.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n h5file.close()\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n if self.test_mode and not self.benchmark:\n dataset = self.chunk(dataset, chunk_dim=1, pct=1.0)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/dtu_rmvd.py", "content": "import json\nimport os\nfrom typing import Any\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.pipelines import AnnotationMask, KittiCrop\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.datasets.utils import DatasetFromList\nfrom unidepth.utils import identity\n\n\nclass DTURMVD(SequenceDataset):\n min_depth = 0.05\n max_depth = 3.0\n depth_scale = 1000.0\n default_fps = 6\n test_split = \"test.txt\"\n train_split = \"test.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"dtu_rmvd.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/dummy.py", "content": "import numpy as np\nimport torch\nfrom torch.utils.data import Dataset\n\n\nclass Dummy(Dataset):\n train_split = None\n test_split = None\n\n def __init__(self, *args, **kwargs):\n super().__init__()\n self.dataset = np.arange(1_000_000)\n\n def get_single_item(self, idx):\n # results = {}\n # results[\"cam2w\"] = torch.eye(4).unsqueeze(0)\n # results[\"K\"] = torch.eye(3).unsqueeze(0)\n # results[\"image\"] = torch.zeros(1, 3, 1024, 1024).to(torch.uint8)\n # results[\"depth\"] = torch.zeros(1, 1, 1024, 1024).to(torch.float32)\n return {\n \"x\": {(0, 0): torch.rand(1, 3, 1024, 1024, dtype=torch.float32)},\n \"img_metas\": {\"val\": torch.rand(1, 1024, dtype=torch.float32)},\n }\n\n def __getitem__(self, idx):\n if isinstance(idx, (list, tuple)):\n results = [self.get_single_item(i) for i in idx]\n else:\n results = self.get_single_item(idx)\n return results\n\n def __len__(self):\n return len(self.dataset)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/dynamic_replica.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass DynReplica(SequenceDataset):\n min_depth = 0.01\n max_depth = 20.0\n default_fps = 30.0\n depth_scale = 512.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences_clean.json\"\n hdf5_paths = [\"DynReplica.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/eden.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass EDEN(SequenceDataset):\n min_depth = 0.1\n max_depth = 100.0\n depth_scale = 256.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"EDEN.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/eth3d.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass ETH3D(ImageDataset):\n min_depth = 0.01\n max_depth = 50.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"ETH3D.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n\n\nclass ETH3D_F(SequenceDataset):\n min_depth = 0.05\n max_depth = 60.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"ETH3D-F.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, float],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=(\n decode_fields if not test_mode else [*decode_fields, \"points\"]\n ),\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/eth3d_rmvd.py", "content": "import json\nimport os\nfrom typing import Any\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.pipelines import AnnotationMask, KittiCrop\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.datasets.utils import DatasetFromList\nfrom unidepth.utils import identity\n\n\nclass ETH3DRMVD(SequenceDataset):\n min_depth = 0.01\n max_depth = 50.0\n depth_scale = 1000.0\n default_fps = 6\n test_split = \"test.txt\"\n train_split = \"test.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"eth3d_rmvd.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/facedepth.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass FaceDepth(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n default_fps = 10\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"FaceDepth.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/flsea.py", "content": "import os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass FLSea(ImageDataset):\n CAM_INTRINSIC = {\n \"canyons\": torch.tensor(\n [\n [1175.3913431656817, 0.0, 466.2595428966926],\n [0.0, 1174.2805075232263, 271.2116633091501],\n [0.0, 0.0, 1.0],\n ]\n ),\n \"red_sea\": torch.tensor(\n [\n [1296.666758476217, 0.0, 501.50386149846],\n [0.0, 1300.831316354508, 276.161712082695],\n [0.0, 0.0, 1.0],\n ]\n ),\n }\n min_depth = 0.05\n max_depth = 20.0\n depth_scale = 1000.0\n train_split = \"train.txt\"\n hdf5_paths = [\"FLSea.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=False,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n sample = [image_filename, depth_filename]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n if self.test_mode and not self.benchmark:\n dataset = self.chunk(dataset, chunk_dim=1, pct=0.33)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_intrinsics(self, idx, image_name):\n return self.CAM_INTRINSIC[image_name.split(\"/\")[0]][:, :3].clone()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"quality\"] = [2] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/futurehouse.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.pipelines import Compose, PanoCrop, PanoRoll\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass FutureHouse(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"FutureHouse.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"cam2w\", \"camera_params\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n self.resizer = Compose(\n [PanoCrop(), PanoRoll(test_mode=test_mode), self.resizer]\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/gibson.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.pipelines import Compose, PanoCrop, PanoRoll\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Gibson(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"Gibson.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"cam2w\", \"camera_params\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n self.resizer = Compose(\n [PanoCrop(), PanoRoll(test_mode=test_mode), self.resizer]\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/hammer.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass HAMMER(ImageDataset):\n min_depth = 0.005\n max_depth = 10.0\n depth_scale = 1000.0\n train_split = \"test.txt\"\n test_split = \"test.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"hammer.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/hm3d.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass HM3D(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"val.txt\"\n train_split = \"full.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"HM3D.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/hoi4d.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass HOI4D(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n default_fps = 5\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"HOI4D.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/hrwsi.py", "content": "import os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass HRWSI(ImageDataset):\n min_depth = 0.01\n max_depth = 1000.0\n depth_scale = 50.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n hdf5_paths = [\"HRWSI.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.split_file), \"w\") as f:\n # f.write(txt_string)\n\n dataset = []\n\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n sample = [\n image_filename,\n depth_filename,\n ]\n dataset.append(sample)\n h5file.close()\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"ssi\"] = [True]\n results[\"valid_camera\"] = [False]\n return results\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/hypersim.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass HyperSim(ImageDataset):\n min_depth = 0.01\n max_depth = 50.0\n depth_scale = 1000.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [f\"hypersim/hypersim_{i}.hdf5\" for i in range(8)]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\").strip(\"\\n\")\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.split_file), \"w\") as f:\n # f.write(txt_string)\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.intrisics_file), \"w\") as f:\n # json.dump(intrinsics, f)\n\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename, chunk_idx = line.strip().split(\" \")\n intrinsics_val = torch.tensor(\n intrinsics[os.path.join(*image_filename.split(\"/\")[:2])]\n ).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val, chunk_idx]\n dataset.append(sample)\n h5file.close()\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n if self.test_mode and not self.benchmark: # corresponds to 712 images\n dataset = self.chunk(dataset, chunk_dim=1, pct=0.1)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n \"K\": 2,\n \"chunk_idx\": 3,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"synthetic\"] = [True] * self.num_copies\n results[\"quality\"] = [0] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/ibims.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass IBims(ImageDataset):\n min_depth = 0.005\n max_depth = 25.0\n depth_scale = 1000.0\n train_split = \"ibims_val.txt\"\n test_split = \"ibims_val.txt\"\n intrisics_file = \"ibims_intrinsics.json\"\n hdf5_paths = [\"ibims.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n\n\nclass IBims_F(SequenceDataset):\n min_depth = 0.01\n max_depth = 25.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"IBims-F.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, float],\n resize_method: str,\n mini: float,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=(\n decode_fields if not test_mode else [*decode_fields, \"points\"]\n ),\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/image_dataset.py", "content": "import io\nimport os\nfrom time import time\nfrom typing import Any, Dict, List, Tuple\n\nimport numpy as np\nimport tables\nimport torch\nimport torchvision\nimport torchvision.transforms.v2.functional as TF\nfrom PIL import Image\n\nfrom unidepth.datasets.base_dataset import BaseDataset\nfrom unidepth.utils import is_main_process\nfrom unidepth.utils.camera import BatchCamera, Pinhole\n\n\"\"\"\nAwful class for legacy reasons, we assume only pinhole cameras\nAnd we \"fake\" sequences by setting sequence_fields to [(0, 0)] and cam2w as eye(4)\n\"\"\"\n\n\nclass ImageDataset(BaseDataset):\n def __init__(\n self,\n image_shape: Tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: Dict[str, Any],\n resize_method: str,\n mini: float,\n benchmark: bool = False,\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.mapper = self.get_mapper()\n\n def get_single_item(self, idx, sample=None, mapper=None):\n sample = self.dataset[idx] if sample is None else sample\n mapper = self.mapper if mapper is None else mapper\n\n results = {\n (0, 0): dict(\n gt_fields=set(),\n image_fields=set(),\n mask_fields=set(),\n camera_fields=set(),\n )\n }\n results = self.pre_pipeline(results)\n results[\"sequence_fields\"] = [(0, 0)]\n\n chunk_idx = (\n int(sample[self.mapper[\"chunk_idx\"]]) if \"chunk_idx\" in self.mapper else 0\n )\n h5_path = os.path.join(self.data_root, self.hdf5_paths[chunk_idx])\n with tables.File(\n h5_path,\n mode=\"r\",\n libver=\"latest\",\n swmr=True,\n ) as h5file_chunk:\n for key_mapper, idx_mapper in mapper.items():\n if \"image\" not in key_mapper and \"depth\" not in key_mapper:\n continue\n value = sample[idx_mapper]\n results[(0, 0)][key_mapper] = value\n name = key_mapper.replace(\"_filename\", \"\")\n value_root = \"/\" + value\n\n if \"image\" in key_mapper:\n results[(0, 0)][\"filename\"] = value\n file = h5file_chunk.get_node(value_root).read()\n image = (\n torchvision.io.decode_image(torch.from_numpy(file))\n .to(torch.uint8)\n .squeeze()\n )\n results[(0, 0)][\"image_fields\"].add(name)\n results[(0, 0)][f\"image_ori_shape\"] = image.shape[-2:]\n results[(0, 0)][name] = image[None, ...]\n\n # collect camera information for the given image\n name = name.replace(\"image_\", \"\")\n results[(0, 0)][\"camera_fields\"].update({\"camera\", \"cam2w\"})\n K = self.get_intrinsics(idx, value)\n if K is None:\n K = torch.eye(3)\n K[0, 0] = K[1, 1] = 0.7 * self.image_shape[1]\n K[0, 2] = 0.5 * self.image_shape[1]\n K[1, 2] = 0.5 * self.image_shape[0]\n\n camera = Pinhole(K=K[None, ...].clone())\n results[(0, 0)][\"camera\"] = BatchCamera.from_camera(camera)\n results[(0, 0)][\"cam2w\"] = self.get_extrinsics(idx, value)[\n None, ...\n ]\n\n elif \"depth\" in key_mapper:\n # start = time()\n file = h5file_chunk.get_node(value_root).read()\n depth = Image.open(io.BytesIO(file))\n depth = TF.pil_to_tensor(depth).squeeze().to(torch.float32)\n if depth.ndim == 3:\n depth = depth[2] + depth[1] * 255 + depth[0] * 255 * 255\n\n results[(0, 0)][\"gt_fields\"].add(name)\n results[(0, 0)][f\"depth_ori_shape\"] = depth.shape\n\n depth = (\n depth.view(1, 1, *depth.shape).contiguous() / self.depth_scale\n )\n results[(0, 0)][name] = depth\n\n results = self.preprocess(results)\n if not self.test_mode:\n results = self.augment(results)\n results = self.postprocess(results)\n return results\n\n def preprocess(self, results):\n results = self.replicate(results)\n for i, seq in enumerate(results[\"sequence_fields\"]):\n self.resizer.ctx = None\n results[seq] = self.resizer(results[seq])\n num_pts = torch.count_nonzero(results[seq][\"depth\"] > 0)\n if num_pts < 50:\n raise IndexError(f\"Too few points in depth map ({num_pts})\")\n\n for key in results[seq].get(\"image_fields\", [\"image\"]):\n results[seq][key] = results[seq][key].to(torch.float32) / 255\n\n # update fields common in sequence\n for key in [\"image_fields\", \"gt_fields\", \"mask_fields\", \"camera_fields\"]:\n if key in results[(0, 0)]:\n results[key] = results[(0, 0)][key]\n results = self.pack_batch(results)\n return results\n\n def postprocess(self, results):\n # normalize after because color aug requires [0,255]?\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.normalize(results[key], **self.normalization_stats)\n results = self.filler(results)\n results = self.unpack_batch(results)\n results = self.masker(results)\n results = self.collecter(results)\n return results\n\n def __getitem__(self, idx):\n try:\n if isinstance(idx, (list, tuple)):\n results = [self.get_single_item(i) for i in idx]\n else:\n results = self.get_single_item(idx)\n except Exception as e:\n print(f\"Error loading sequence {idx} for {self.__class__.__name__}: {e}\")\n idx = np.random.randint(0, len(self.dataset))\n results = self[idx]\n return results\n\n def get_intrinsics(self, idx, image_name):\n idx_sample = self.mapper.get(\"K\", 1000)\n sample = self.dataset[idx]\n if idx_sample >= len(sample):\n return None\n return sample[idx_sample]\n\n def get_extrinsics(self, idx, image_name):\n idx_sample = self.mapper.get(\"cam2w\", 1000)\n sample = self.dataset[idx]\n if idx_sample >= len(sample):\n return torch.eye(4)\n return sample[idx_sample]\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n \"K\": 2,\n }\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/ken_burns.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass KenBurns(ImageDataset):\n min_depth = 0.05\n max_depth = 50.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [f\"3dkenburns/3DKenBurns_{i}.hdf5\" for i in range(8)]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\").strip(\"\\n\")\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.split_file), \"w\") as f:\n # f.write(txt_string)\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.intrisics_file), \"w\") as f:\n # json.dump(intrinsics, f)\n\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename, chunk_idx = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val, chunk_idx]\n dataset.append(sample)\n h5file.close()\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n if self.test_mode and not self.benchmark: # corresponds to 500 images\n dataset = self.chunk(dataset, chunk_dim=1, pct=0.25)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n \"K\": 2,\n \"chunk_idx\": 3,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"synthetic\"] = [True] * self.num_copies\n results[\"quality\"] = [0] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/kitti.py", "content": "import os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.pipelines import AnnotationMask, KittiCrop\nfrom unidepth.datasets.utils import DatasetFromList\nfrom unidepth.utils import identity\n\n\nclass KITTI(ImageDataset):\n CAM_INTRINSIC = {\n \"2011_09_26\": torch.tensor(\n [\n [7.215377e02, 0.000000e00, 6.095593e02, 4.485728e01],\n [0.000000e00, 7.215377e02, 1.728540e02, 2.163791e-01],\n [0.000000e00, 0.000000e00, 1.000000e00, 2.745884e-03],\n ]\n ),\n \"2011_09_28\": torch.tensor(\n [\n [7.070493e02, 0.000000e00, 6.040814e02, 4.575831e01],\n [0.000000e00, 7.070493e02, 1.805066e02, -3.454157e-01],\n [0.000000e00, 0.000000e00, 1.000000e00, 4.981016e-03],\n ]\n ),\n \"2011_09_29\": torch.tensor(\n [\n [7.183351e02, 0.000000e00, 6.003891e02, 4.450382e01],\n [0.000000e00, 7.183351e02, 1.815122e02, -5.951107e-01],\n [0.000000e00, 0.000000e00, 1.000000e00, 2.616315e-03],\n ]\n ),\n \"2011_09_30\": torch.tensor(\n [\n [7.070912e02, 0.000000e00, 6.018873e02, 4.688783e01],\n [0.000000e00, 7.070912e02, 1.831104e02, 1.178601e-01],\n [0.000000e00, 0.000000e00, 1.000000e00, 6.203223e-03],\n ]\n ),\n \"2011_10_03\": torch.tensor(\n [\n [7.188560e02, 0.000000e00, 6.071928e02, 4.538225e01],\n [0.000000e00, 7.188560e02, 1.852157e02, -1.130887e-01],\n [0.000000e00, 0.000000e00, 1.000000e00, 3.779761e-03],\n ]\n ),\n }\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n log_mean = 2.5462\n log_std = 0.5871\n test_split = \"kitti_eigen_test.txt\"\n train_split = \"kitti_eigen_train.txt\"\n test_split_benchmark = \"kitti_test.txt\"\n hdf5_paths = [\"kitti.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.masker = AnnotationMask(\n min_value=0.0,\n max_value=self.max_depth if test_mode else None,\n custom_fn=self.eval_mask if test_mode else lambda x, *args, **kwargs: x,\n )\n self.test_mode = test_mode\n self.crop = crop\n self.cropper_base = KittiCrop(crop_size=(352, 1216))\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename = line.strip().split(\" \")[0]\n depth_filename = line.strip().split(\" \")[1]\n if depth_filename == \"None\":\n self.invalid_depth_num += 1\n continue\n sample = [\n image_filename,\n depth_filename,\n ]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_intrinsics(self, idx, image_name):\n return self.CAM_INTRINSIC[image_name.split(\"/\")[0]][:, :3].clone()\n\n def preprocess(self, results):\n results = self.replicate(results)\n for i, seq in enumerate(results[\"sequence_fields\"]):\n self.resizer.ctx = None\n results[seq] = self.cropper_base(results[seq])\n results[seq] = self.resizer(results[seq])\n num_pts = torch.count_nonzero(results[seq][\"depth\"] > 0)\n if num_pts < 50:\n raise IndexError(f\"Too few points in depth map ({num_pts})\")\n\n for key in results[seq].get(\"image_fields\", [\"image\"]):\n results[seq][key] = results[seq][key].to(torch.float32) / 255\n\n # update fields common in sequence\n for key in [\"image_fields\", \"gt_fields\", \"mask_fields\", \"camera_fields\"]:\n if key in results[(0, 0)]:\n results[key] = results[(0, 0)][key]\n results = self.pack_batch(results)\n return results\n\n def eval_mask(self, valid_mask, info={}):\n \"\"\"Do grag_crop or eigen_crop for testing\"\"\"\n mask_height, mask_width = valid_mask.shape[-2:]\n eval_mask = torch.zeros_like(valid_mask)\n if \"garg\" in self.crop:\n eval_mask[\n ...,\n int(0.40810811 * mask_height) : int(0.99189189 * mask_height),\n int(0.03594771 * mask_width) : int(0.96405229 * mask_width),\n ] = 1\n elif \"eigen\" in self.crop:\n eval_mask[\n ...,\n int(0.3324324 * mask_height) : int(0.91351351 * mask_height),\n int(0.03594771 * mask_width) : int(0.96405229 * mask_width),\n ] = 1\n return torch.logical_and(valid_mask, eval_mask)\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n\n\nimport json\n\n\nclass KITTIBenchmark(ImageDataset):\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n test_split = \"test_split.txt\"\n train_split = \"val_split.txt\"\n intrinsics_file = \"intrinsics.json\"\n hdf5_paths = [\"kitti_benchmark.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=True,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n\n self.masker = AnnotationMask(\n min_value=self.min_depth,\n max_value=self.max_depth if test_mode else None,\n custom_fn=lambda x, *args, **kwargs: x,\n )\n self.collecter = Collect(keys=[\"image_fields\", \"mask_fields\", \"gt_fields\"])\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_path),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(self.h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrinsics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics = torch.tensor(\n intrinsics[os.path.join(*image_filename.split(\"/\")[:2])]\n ).squeeze()[:, :3]\n sample = {\n \"image_filename\": image_filename,\n \"depth_filename\": depth_filename,\n \"K\": intrinsics,\n }\n dataset.append(sample)\n\n self.dataset = DatasetFromList(dataset)\n\n self.log_load_dataset()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/kitti360.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass KITTI360(SequenceDataset):\n min_depth = 0.01\n max_depth = 80.0\n depth_scale = 256.0\n train_split = \"train.txt\"\n test_split = \"val_split.txt\"\n sequences_file = \"sequences_split.json\"\n hdf5_paths = [f\"KITTI360.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=(\n decode_fields if not test_mode else [*decode_fields, \"points\"]\n ),\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def preprocess(self, results):\n self.resizer.ctx = None\n for i, seq in enumerate(results[\"sequence_fields\"]):\n # Create a mask where the distance from the center is less than H/2\n H, W = results[seq][\"image\"].shape[-2:]\n x = torch.linspace(-W / 2, W / 2, W)\n y = torch.linspace(-H / 2, H / 2, H)\n xv, yv = torch.meshgrid(x, y, indexing=\"xy\")\n distance_from_center = torch.sqrt(xv**2 + yv**2).reshape(1, 1, H, W)\n results[seq][\"validity_mask\"] = distance_from_center < (H / 2)\n return super().preprocess(results)\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/kitti_multi.py", "content": "import json\nimport os\nfrom typing import Any\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.pipelines import AnnotationMask, KittiCrop\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.datasets.utils import DatasetFromList\nfrom unidepth.utils import identity\n\n\nclass KITTIMulti(SequenceDataset):\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n default_fps = 10.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"KITTI_sequence.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n\n self.cropper_base = KittiCrop(crop_size=(352, 1216))\n\n self.masker = AnnotationMask(\n min_value=0.0,\n max_value=self.max_depth if test_mode else None,\n custom_fn=self.eval_mask if test_mode else identity,\n )\n self.eval_last = True\n\n def __len__(self):\n if self.test_mode:\n return 64 # FIXME: Hardcoded for now\n return len(self.dataset)\n\n def preprocess(self, results):\n self.resizer.ctx = None\n for i, seq in enumerate(results[\"sequence_fields\"]):\n results[seq] = self.cropper_base(results[seq])\n results[seq] = self.resizer(results[seq])\n for key in results[seq].get(\"image_fields\", [\"image\"]):\n results[seq][key] = results[seq][key].to(torch.float32) / 255\n results.update({k: v for k, v in results[(0, 0)].items() if \"fields\" in k})\n\n results = self.pack_batch(results)\n return results\n\n def eval_mask(self, valid_mask, info={}):\n \"\"\"Do grag_crop or eigen_crop for testing\"\"\"\n mask_height, mask_width = valid_mask.shape[-2:]\n eval_mask = torch.zeros_like(valid_mask)\n if \"garg\" in self.crop:\n eval_mask[\n ...,\n int(0.40810811 * mask_height) : int(0.99189189 * mask_height),\n int(0.03594771 * mask_width) : int(0.96405229 * mask_width),\n ] = 1\n elif \"eigen\" in self.crop:\n eval_mask[\n ...,\n int(0.3324324 * mask_height) : int(0.91351351 * mask_height),\n int(0.03594771 * mask_width) : int(0.96405229 * mask_width),\n ] = 1\n else:\n return valid_mask\n return torch.logical_and(valid_mask, eval_mask)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/kitti_rmvd.py", "content": "import json\nimport os\nfrom typing import Any\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.pipelines import AnnotationMask, Compose, KittiCrop\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.utils import identity\n\n\nclass KITTIRMVD(SequenceDataset):\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n default_fps = 10\n test_split = \"test.txt\"\n train_split = \"test.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"kitti_rmvd.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n self.crop = crop\n self.resizer = Compose([KittiCrop(crop_size=(352, 1216)), self.resizer])\n\n def eval_mask(self, valid_mask, info={}):\n \"\"\"Do grag_crop or eigen_crop for testing\"\"\"\n mask_height, mask_width = valid_mask.shape[-2:]\n eval_mask = torch.zeros_like(valid_mask)\n if \"garg\" in self.crop:\n eval_mask[\n ...,\n int(0.40810811 * mask_height) : int(0.99189189 * mask_height),\n int(0.03594771 * mask_width) : int(0.96405229 * mask_width),\n ] = 1\n elif \"eigen\" in self.crop:\n eval_mask[\n ...,\n int(0.3324324 * mask_height) : int(0.91351351 * mask_height),\n int(0.03594771 * mask_width) : int(0.96405229 * mask_width),\n ] = 1\n else:\n return valid_mask\n return torch.logical_and(valid_mask, eval_mask)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/lyft.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass Lyft(ImageDataset):\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n test_split = \"test.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"Lyft2.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.split_file), \"w\") as f:\n # f.write(txt_string)\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.intrisics_file), \"w\") as f:\n # json.dump(intrinsics, f)\n\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [\n image_filename,\n depth_filename,\n intrinsics_val,\n ]\n dataset.append(sample)\n h5file.close()\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False]\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/mapillary.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass Mapillary(ImageDataset):\n min_depth = 0.01\n max_depth = 70.0\n depth_scale = 256.0\n test_split = \"mapillary_val.txt\"\n train_split = \"mapillary_train_clean.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"Mapillary.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n h5file.close()\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n if self.test_mode and not self.benchmark:\n dataset = self.chunk(dataset, chunk_dim=1, pct=0.05)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"si\"] = [True] * self.num_copies\n results[\"valid_camera\"] = [False] * self.num_copies\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [2] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/matrix_city.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass MatrixCity(SequenceDataset):\n min_depth = 0.01\n max_depth = 200.0\n depth_scale = 1000.0\n test_split = \"test.txt\"\n train_split = \"train_full.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"MatrixCity.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/matterport3d.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.pipelines import Compose, PanoCrop, PanoRoll\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Matterport3D(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"Matterport3D.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"cam2w\", \"camera_params\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n self.resizer = Compose(\n [PanoCrop(), PanoRoll(test_mode=test_mode), self.resizer]\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/megadepth.py", "content": "import os\n\nimport h5py\nimport numpy as np\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass MegaDepth(ImageDataset):\n min_depth = 0.01\n max_depth = 1000.0\n depth_scale = 50.0\n test_split = \"test.txt\"\n train_split = \"train.txt\"\n hdf5_paths = [\"MegaDepth.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.split_file), \"w\") as f:\n # f.write(txt_string)\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n sample = [\n image_filename,\n depth_filename,\n ]\n dataset.append(sample)\n h5file.close()\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n else:\n dataset = self.chunk(dataset, chunk_dim=1, pct=0.5)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"ssi\"] = [True]\n results[\"valid_camera\"] = [False]\n results[\"dense\"] = [False]\n return results\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/megadepth_s.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass MegaDepthS(SequenceDataset):\n min_depth = 0.001\n max_depth = 10000.0\n depth_scale = 512.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences_filter_clean.json\"\n hdf5_paths = [\"MegaDepthS.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"intrinsics\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/midair.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass MidAir(SequenceDataset):\n min_depth = 0.1\n max_depth = 1000.0\n depth_scale = 1000.0\n default_fps = 6\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"MidAir.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/mip.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass MIP(SequenceDataset):\n min_depth = 0.01\n max_depth = 100.0\n depth_scale = 1000.0\n default_fps = 10\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"MIP.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/ms2.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass MS2(SequenceDataset):\n min_depth = 0.01\n max_depth = 100.0\n depth_scale = 256.0\n default_fps = 5\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"MS2.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/mvimgnet.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\nINVALID_SEQUENCES = [\n \"1/000121f2-0\",\n \"15/1600ae56-0\",\n \"26/000000f3-0\",\n \"33/1d00e677-0\",\n \"43/22008925-0\",\n \"49/000147db-0\",\n \"51/23002a43-0\",\n \"51/23000916-0\",\n \"108/000133ae-0\",\n \"129/000037f2-0\",\n \"141/17012545-0\",\n \"141/1700f3de-0\",\n \"152/1b00e061-0\",\n \"154/1d00decb-0\",\n \"154/1d017c1c-0\",\n \"154/1d0019a5-0\",\n \"154/1d00334d-0\",\n \"154/1d012ed6-0\",\n \"154/1d016b8a-0\",\n \"154/1d016cc1-0\",\n \"154/1d008d5f-0\",\n \"159/000157f9-0\",\n \"159/00000b96-0\",\n \"159/000075c0-0\",\n \"159/0000445c-0\",\n \"159/000056a0-0\",\n \"159/00010c68-0\",\n \"159/0000573b-0\",\n \"159/00002698-0\",\n \"159/00008fca-0\",\n \"159/00009ef8-0\",\n \"159/00015f05-0\",\n \"159/0000c6df-0\",\n \"159/0000ee59-0\",\n \"163/290159d2-0\",\n \"163/29016c7c-0\",\n \"163/2900239c-0\",\n \"163/29002f7b-0\",\n \"163/29014b05-0\",\n \"163/29000196-0\",\n \"163/2901750f-0\",\n \"164/1b0145cf-0\",\n \"164/1b00eb1d-0\",\n \"164/1b00c28b-0\",\n \"164/1b0110d0-0\",\n \"164/1b00dd20-0\",\n \"165/2600e15a-0\",\n \"165/26008444-0\",\n \"165/260145c5-0\",\n \"165/26003a0c-0\",\n \"165/260106ba-0\",\n \"165/26001548-0\",\n \"167/2a0092b0-0\",\n \"167/2a014dbe-0\",\n \"167/2a003ce6-0\",\n \"169/1800c645-0\",\n \"171/2500014d-0\",\n \"176/1d0021c2-0\",\n \"176/1d014abf-0\",\n \"176/1d00e714-0\",\n \"176/1d0159cb-0\",\n \"176/1e016629-0\",\n \"178/000102b8-0\",\n \"191/23008fdb-0\",\n \"191/2300187f-0\",\n \"191/2300ae68-0\",\n \"191/230076dd-0\",\n \"191/24007d7e-0\",\n \"192/000107b5-0\",\n \"195/1f012359-0\",\n \"195/1f00f751-0\",\n \"195/1f011331-0\",\n \"195/1e00d999-0\",\n \"196/1c01304e-0\",\n \"198/1a00e02f-0\",\n \"198/050084ac-0\",\n \"198/1a0075fa-0\",\n \"199/1e001742-0\",\n \"199/1e00116a-0\",\n \"199/1e011d00-0\",\n \"199/1e018040-0\",\n \"199/1e001107-0\",\n]\n\n\nclass MVImgNet(SequenceDataset):\n min_depth = 0.005\n max_depth = 10.0\n # weird scale issue, should be 1000, but avg depth is ~10meters...\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"MVImgNet.hdf5\"]\n invalid_sequences = INVALID_SEQUENCES\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"intrinsics\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/mvsynth.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass MVSynth(SequenceDataset):\n min_depth = 0.1\n max_depth = 1000.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"MVSynth.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/nerds360.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass NeRDS360(SequenceDataset):\n min_depth = 0.01\n max_depth = 1000.0\n depth_scale = 1000.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"NeRDS360.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/niantic_mapfree.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass NianticMapFree(SequenceDataset):\n min_depth = 0.1\n max_depth = 250.0\n depth_scale = 512.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"NianticMapFree.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/nuscenes.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass Nuscenes(ImageDataset):\n min_depth = 0.05\n max_depth = 80.0\n depth_scale = 256.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n intrisics_file = \"intrinsics.json\"\n # hdf5_paths = [\"Nuscenes2.hdf5\"]\n hdf5_paths = [f\"nuscenes/nuscenes_{i}.hdf5\" for i in range(8)]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\").strip(\"\\n\")\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename, chunk_idx = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val, chunk_idx]\n dataset.append(sample)\n h5file.close()\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=6, pct=self.mini)\n if self.test_mode and not self.benchmark:\n dataset = self.chunk(dataset, chunk_dim=6, pct=0.1)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n \"K\": 2,\n \"chunk_idx\": 3,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_copies\n results[\"quality\"] = [1] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/nyuv2.py", "content": "import os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.pipelines import AnnotationMask\nfrom unidepth.datasets.utils import DatasetFromList\nfrom unidepth.utils import identity\n\n\nclass NYUv2Depth(ImageDataset):\n CAM_INTRINSIC = {\n \"ALL\": torch.tensor(\n [\n [5.1885790117450188e02, 0, 3.2558244941119034e02],\n [0, 5.1946961112127485e02, 2.5373616633400465e02],\n [0, 0, 1],\n ]\n )\n }\n min_depth = 0.005\n max_depth = 10.0\n depth_scale = 1000.0\n log_mean = 0.9140\n log_std = 0.4825\n test_split = \"nyu_test.txt\"\n train_split = \"nyu_train.txt\"\n hdf5_paths = [\"nyuv2.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.masker = AnnotationMask(\n min_value=0.0,\n max_value=self.max_depth if test_mode else None,\n custom_fn=self.eval_mask if test_mode else lambda x, *args, **kwargs: x,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename, _ = line.strip().split(\" \")\n sample = [\n image_filename,\n depth_filename,\n ]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n return results\n\n def get_intrinsics(self, idx, image_name):\n return self.CAM_INTRINSIC[\"ALL\"].clone()\n\n def eval_mask(self, valid_mask, info={}):\n border_mask = torch.zeros_like(valid_mask)\n border_mask[..., 45:-9, 41:-39] = 1\n return torch.logical_and(valid_mask, border_mask)\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"quality\"] = [2] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/oasis.py", "content": "import os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass OASISv2(ImageDataset):\n min_depth = 0.01\n max_depth = 400.0\n depth_scale = 1000.0\n test_split = \"val.txt\"\n train_split = \"train.txt\"\n hdf5_paths = [\"Oasis2.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n\n dataset = []\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.split_file), \"w\") as f:\n # f.write(txt_string)\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n sample = [image_filename, depth_filename]\n dataset.append(sample)\n h5file.close()\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"ssi\"] = [True]\n results[\"valid_camera\"] = [False]\n return results\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n }\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/pipelines/__init__.py", "content": "from .formating import AnnotationMask, Collect\nfrom .transforms import (Compose, ContextCrop, Crop, GaussianBlur, KittiCrop,\n PanoCrop, PanoRoll, RandomAutoContrast,\n RandomBrightness, RandomColor, RandomColorJitter,\n RandomContrast, RandomEqualize, RandomFiller,\n RandomFlip, RandomGamma, RandomGrayscale,\n RandomInvert, RandomMasking, RandomPosterize,\n RandomSaturation, RandomSharpness, RandomShear,\n RandomSolarize, RandomTranslate, Rotate)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/pipelines/formating.py", "content": "from collections.abc import Sequence\n\nimport numpy as np\nimport torch\n\n\nclass Collect(object):\n def __init__(\n self,\n keys,\n meta_keys=(\n \"filename\",\n \"keyframe_idx\",\n \"sequence_name\",\n \"image_filename\",\n \"depth_filename\",\n \"image_ori_shape\",\n \"camera\",\n \"original_camera\",\n \"sfm\",\n \"image_shape\",\n \"resized_shape\",\n \"scale_factor\",\n \"rotation\",\n \"resize_factor\",\n \"flip\",\n \"flip_direction\",\n \"dataset_name\",\n \"paddings\",\n \"max_value\",\n \"log_mean\",\n \"log_std\",\n \"image_rescale\",\n \"focal_rescale\",\n \"depth_rescale\",\n ),\n ):\n self.keys = keys\n self.meta_keys = meta_keys\n\n def __call__(self, results):\n data_keys = [key for field in self.keys for key in results.get(field, [])]\n data = {\n key: {\n sequence_key: results[key][sequence_key]\n for sequence_key in results[\"sequence_fields\"]\n }\n for key in data_keys\n }\n data[\"img_metas\"] = {\n key: value for key, value in results.items() if key not in data_keys\n }\n return data\n\n def __repr__(self):\n return (\n self.__class__.__name__ + f\"(keys={self.keys}, meta_keys={self.meta_keys})\"\n )\n\n\nclass AnnotationMask(object):\n def __init__(self, min_value, max_value, custom_fn=lambda x: x):\n self.min_value = min_value\n self.max_value = max_value\n self.custom_fn = custom_fn\n\n def __call__(self, results):\n for key in results.get(\"gt_fields\", []):\n if key + \"_mask\" in results[\"mask_fields\"]:\n if \"flow\" in key:\n for sequence_idx in results.get(\"sequence_fields\", []):\n boundaries = (results[key][sequence_idx] >= -1) & (\n results[key][sequence_idx] <= 1\n )\n boundaries = boundaries[:, :1] & boundaries[:, 1:]\n results[key + \"_mask\"][sequence_idx] = (\n results[key + \"_mask\"][sequence_idx] & boundaries\n )\n continue\n for sequence_idx in results.get(\"sequence_fields\", []):\n mask = results[key][sequence_idx] > self.min_value\n if self.max_value is not None:\n mask = mask & (results[key][sequence_idx] < self.max_value)\n mask = self.custom_fn(mask, info=results)\n if key + \"_mask\" not in results:\n results[key + \"_mask\"] = {}\n results[key + \"_mask\"][sequence_idx] = mask.bool()\n results[\"mask_fields\"].add(key + \"_mask\")\n return results\n\n def __repr__(self):\n return (\n self.__class__.__name__\n + f\"(min_value={self.min_value}, max_value={ self.max_value})\"\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/pipelines/transforms.py", "content": "import os\nimport random\nfrom copy import deepcopy\nfrom math import ceil, exp, log, log2, log10, tanh\nfrom typing import Dict, List, Tuple\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nimport torchvision.transforms.v2.functional as TF\n\nfrom unidepth.utils.geometric import downsample\n\n\nclass PanoCrop:\n def __init__(self, crop_v=0.1):\n self.crop_v = crop_v\n\n def _crop_data(self, results, crop_size):\n \"\"\"Function to randomly crop images, bounding boxes, masks, semantic\n segmentation maps.\n\n Args:\n results (dict): Result dict from loading pipeline.\n crop_size (tuple): Expected absolute size after cropping, (h, w).\n allow_negative_crop (bool): Whether to allow a crop that does not\n contain any bbox area. Default to False.\n\n Returns:\n dict: Randomly cropped results, 'image_shape' key in result dict is\n updated according to crop size.\n \"\"\"\n offset_w, offset_h = crop_size\n left, top, right, bottom = offset_w[0], offset_h[0], offset_w[1], offset_h[1]\n H, W = results[\"image\"].shape[-2:]\n for key in results.get(\"image_fields\", [\"image\"]):\n img = results[key][..., top : H - bottom, left : W - right]\n results[key] = img\n results[\"image_shape\"] = tuple(img.shape)\n\n for key in results.get(\"gt_fields\", []):\n results[key] = results[key][..., top : H - bottom, left : W - right]\n\n for key in results.get(\"mask_fields\", []):\n results[key] = results[key][..., top : H - bottom, left : W - right]\n\n results[\"camera\"].crop(left, top, right, bottom)\n return results\n\n def __call__(self, results):\n H, W = results[\"image\"].shape[-2:]\n crop_w = (0, 0)\n crop_h = (int(H * self.crop_v), int(H * self.crop_v))\n results = self._crop_data(results, (crop_w, crop_h))\n return results\n\n\nclass PanoRoll:\n def __init__(self, roll=[-0.5, 0.5]):\n self.roll = roll\n\n def __call__(self, results):\n W = results[\"image\"].shape[-1]\n roll = random.randint(int(W * self.roll[0]), int(W * self.roll[1]))\n for key in results.get(\"image_fields\", [\"image\"]):\n img = results[key]\n img = torch.roll(img, roll, dims=-1)\n results[key] = img\n for key in results.get(\"gt_fields\", []):\n results[key] = torch.roll(results[key], roll, dims=-1)\n for key in results.get(\"mask_fields\", []):\n results[key] = torch.roll(results[key], roll, dims=-1)\n return results\n\n\nclass RandomFlip:\n \"\"\"Flip the points & bbox.\n\n If the input dict contains the key \"flip\", then the flag will be used,\n otherwise it will be randomly decided by a ratio specified in the init\n method.\n\n Args:\n flip_ratio_bev_horizontal (float, optional): The flipping probability\n in horizontal direction. Defaults to 0.0.\n flip_ratio_bev_vertical (float, optional): The flipping probability\n in vertical direction. Defaults to 0.0.\n \"\"\"\n\n def __init__(self, direction=\"horizontal\", prob=0.5, **kwargs):\n self.flip_ratio = prob\n valid_directions = [\"horizontal\", \"vertical\", \"diagonal\"]\n if isinstance(direction, str):\n assert direction in valid_directions\n elif isinstance(direction, list):\n assert set(direction).issubset(set(valid_directions))\n else:\n raise ValueError(\"direction must be either str or list of str\")\n self.direction = direction\n\n def __call__(self, results):\n \"\"\"Call function to flip points, values in the ``bbox3d_fields`` and\n also flip 2D image and its annotations.\n\n Args:\n results (dict): Result dict from loading pipeline.\n\n Returns:\n dict: Flipped results, 'flip', 'flip_direction',\n \"\"\"\n if \"flip\" not in results:\n if isinstance(self.direction, list):\n # None means non-flip\n direction_list = self.direction + [None]\n else:\n # None means non-flip\n direction_list = [self.direction, None]\n\n if isinstance(self.flip_ratio, list):\n non_flip_ratio = 1 - sum(self.flip_ratio)\n flip_ratio_list = self.flip_ratio + [non_flip_ratio]\n else:\n non_flip_ratio = 1 - self.flip_ratio\n # exclude non-flip\n single_ratio = self.flip_ratio / (len(direction_list) - 1)\n flip_ratio_list = [single_ratio] * (len(direction_list) - 1) + [\n non_flip_ratio\n ]\n\n cur_dir = np.random.choice(direction_list, p=flip_ratio_list)\n\n results[\"flip\"] = cur_dir is not None\n\n if \"flip_direction\" not in results:\n results[\"flip_direction\"] = cur_dir\n\n if results[\"flip\"]:\n # flip image\n if results[\"flip_direction\"] != \"vertical\":\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.hflip(results[key])\n for key in results.get(\"mask_fields\", []):\n results[key] = TF.hflip(results[key])\n for key in results.get(\"gt_fields\", []):\n results[key] = TF.hflip(results[key])\n if \"flow\" in key: # flip u direction\n results[key][:, 0] = -results[key][:, 0]\n\n H, W = results[\"image\"].shape[-2:]\n results[\"camera\"] = results[\"camera\"].flip(\n H=H, W=W, direction=\"horizontal\"\n )\n # results[\"K\"][..., 0, 2] = results[\"image\"].shape[-1] - results[\"K\"][..., 0, 2]\n # flip: - t_x rotate around y by: pi - angle_y * 2\n flip_transform = torch.tensor(\n [[-1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],\n dtype=torch.float32,\n ).unsqueeze(0)\n repeats = (results[\"cam2w\"].shape[0],) + (1,) * (\n results[\"cam2w\"].ndim - 1\n )\n results[\"cam2w\"] = flip_transform.repeat(*repeats) @ results[\"cam2w\"]\n\n if results[\"flip_direction\"] != \"horizontal\":\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.vflip(results[key])\n for key in results.get(\"mask_fields\", []):\n results[key] = TF.vflip(results[key])\n for key in results.get(\"gt_fields\", []):\n results[key] = TF.vflip(results[key])\n results[\"K\"][..., 1, 2] = (\n results[\"image\"].shape[-2] - results[\"K\"][..., 1, 2]\n )\n results[\"flip\"] = [results[\"flip\"]] * len(results[\"image\"])\n return results\n\n def __repr__(self):\n \"\"\"str: Return a string that describes the module.\"\"\"\n repr_str = self.__class__.__name__\n repr_str += f\" flip_ratio={self.flip_ratio})\"\n return repr_str\n\n\nclass Crop:\n def __init__(\n self,\n crop_size,\n crop_type=\"absolute\",\n crop_offset=(0, 0),\n ):\n if crop_type not in [\n \"relative_range\",\n \"relative\",\n \"absolute\",\n \"absolute_range\",\n ]:\n raise ValueError(f\"Invalid crop_type {crop_type}.\")\n if crop_type in [\"absolute\", \"absolute_range\"]:\n assert crop_size[0] > 0 and crop_size[1] > 0\n assert isinstance(crop_size[0], int) and isinstance(crop_size[1], int)\n else:\n assert 0 < crop_size[0] <= 1 and 0 < crop_size[1] <= 1\n self.crop_size = crop_size\n self.crop_type = crop_type\n self.offset_h, self.offset_w = (\n crop_offset[: len(crop_offset) // 2],\n crop_offset[len(crop_offset) // 2 :],\n )\n\n def _get_crop_size(self, image_shape):\n h, w = image_shape\n if self.crop_type == \"absolute\":\n return (min(self.crop_size[0], h), min(self.crop_size[1], w))\n elif self.crop_type == \"absolute_range\":\n assert self.crop_size[0] <= self.crop_size[1]\n crop_h = np.random.randint(\n min(h, self.crop_size[0]), min(h, self.crop_size[1]) + 1\n )\n crop_w = np.random.randint(\n min(w, self.crop_size[0]), min(w, self.crop_size[1]) + 1\n )\n return crop_h, crop_w\n elif self.crop_type == \"relative\":\n crop_h, crop_w = self.crop_size\n return int(h * crop_h + 0.5), int(w * crop_w + 0.5)\n elif self.crop_type == \"relative_range\":\n crop_size = np.asarray(self.crop_size, dtype=np.float32)\n crop_h, crop_w = crop_size + np.random.rand(2) * (1 - crop_size)\n return int(h * crop_h + 0.5), int(w * crop_w + 0.5)\n\n def _crop_data(self, results, crop_size):\n assert crop_size[0] > 0 and crop_size[1] > 0\n for key in results.get(\"image_fields\", [\"image\"]):\n img = results[key]\n img = TF.crop(\n img, self.offset_h[0], self.offset_w[0], crop_size[0], crop_size[1]\n )\n results[key] = img\n results[\"image_shape\"] = tuple(img.shape)\n\n for key in results.get(\"gt_fields\", []):\n gt = results[key]\n results[key] = TF.crop(\n gt, self.offset_h[0], self.offset_w[0], crop_size[0], crop_size[1]\n )\n\n # crop semantic seg\n for key in results.get(\"mask_fields\", []):\n mask = results[key]\n results[key] = TF.crop(\n mask, self.offset_h[0], self.offset_w[0], crop_size[0], crop_size[1]\n )\n\n results[\"K\"][..., 0, 2] = results[\"K\"][..., 0, 2] - self.offset_w[0]\n results[\"K\"][..., 1, 2] = results[\"K\"][..., 1, 2] - self.offset_h[0]\n return results\n\n def __call__(self, results):\n image_shape = results[\"image\"].shape[-2:]\n crop_size = self._get_crop_size(image_shape)\n results = self._crop_data(results, crop_size)\n\n return results\n\n def __repr__(self):\n repr_str = self.__class__.__name__\n repr_str += f\"(crop_size={self.crop_size}, \"\n repr_str += f\"crop_type={self.crop_type}, \"\n return repr_str\n\n\nclass KittiCrop:\n def __init__(self, crop_size):\n self.crop_size = crop_size\n\n def _crop_data(self, results, crop_size):\n \"\"\"Function to randomly crop images, bounding boxes, masks, semantic\n segmentation maps.\n\n Args:\n results (dict): Result dict from loading pipeline.\n crop_size (tuple): Expected absolute size after cropping, (h, w).\n allow_negative_crop (bool): Whether to allow a crop that does not\n contain any bbox area. Default to False.\n\n Returns:\n dict: Randomly cropped results, 'image_shape' key in result dict is\n updated according to crop size.\n \"\"\"\n assert crop_size[0] > 0 and crop_size[1] > 0\n for key in results.get(\"image_fields\", [\"image\"]):\n img = results[key]\n h, w = img.shape[-2:]\n offset_h, offset_w = int(h - self.crop_size[0]), int(\n (w - self.crop_size[1]) / 2\n )\n\n # crop the image\n img = TF.crop(img, offset_h, offset_w, crop_size[0], crop_size[1])\n results[key] = img\n results[\"image_shape\"] = tuple(img.shape)\n\n for key in results.get(\"gt_fields\", []):\n gt = results[key]\n results[key] = TF.crop(gt, offset_h, offset_w, crop_size[0], crop_size[1])\n\n # crop semantic seg\n for key in results.get(\"mask_fields\", []):\n mask = results[key]\n results[key] = TF.crop(mask, offset_h, offset_w, crop_size[0], crop_size[1])\n\n results[\"camera\"].crop(offset_w, offset_h)\n return results\n\n def __call__(self, results):\n \"\"\"Call function to randomly crop images, bounding boxes, masks,\n semantic segmentation maps.\n\n Args:\n results (dict): Result dict from loading pipeline.\n\n Returns:\n dict: Randomly cropped results, 'image_shape' key in result dict is\n updated according to crop size.\n \"\"\"\n results = self._crop_data(results, self.crop_size)\n\n return results\n\n def __repr__(self):\n repr_str = self.__class__.__name__\n repr_str += f\"(crop_size={self.crop_size}, \"\n return repr_str\n\n\nclass RandomMasking:\n def __init__(\n self,\n mask_ratio,\n mask_patch=16,\n prob=0.5,\n warmup_steps=50000,\n sampling=\"random\",\n curriculum=False,\n ):\n self.mask_patch = mask_patch\n self.prob = prob\n self.mask_ratio = mask_ratio\n self.warmup_steps = max(1, warmup_steps)\n self.hard_bound = 1\n self.idx = 0\n self.curriculum = curriculum\n self.sampling = sampling\n self.low_bound = 0.0\n self.up_bound = 0.0\n\n def __call__(self, results):\n B, _, H, W = results[\"image\"].shape\n device = results[\"image\"].device\n down_size = H // self.mask_patch, W // self.mask_patch\n if np.random.random() > self.prob: # fill with dummy\n return self._nop(results, down_size, device)\n\n validity_mask = results[\"validity_mask\"].float().reshape(B, -1, H, W)\n validity_mask = F.interpolate(validity_mask, size=down_size).bool()\n validity_mask = validity_mask.reshape(B, 1, *down_size)\n is_random = self.is_warmup or results.get(\"guidance\") is None\n\n if not is_random:\n guidance = F.interpolate(results[\"guidance\"], size=(H, W), mode=\"bilinear\")\n results[\"guidance\"] = -F.max_pool2d(\n -guidance, kernel_size=self.mask_patch, stride=self.mask_patch\n )\n\n if is_random and self.sampling == \"inverse\":\n sampling = self.inverse_sampling\n elif is_random and self.sampling == \"random\":\n sampling = self.random_sampling\n else:\n sampling = self.guided_sampling\n mask_ratio = np.random.uniform(self.low_bound, self.up_bound)\n for key in results.get(\"image_fields\", [\"image\"]):\n mask = sampling(results, mask_ratio, down_size, validity_mask, device)\n results[key + \"_mask\"] = mask\n return results\n\n def _nop(self, results, down_size, device):\n B = results[\"image\"].shape[0]\n for key in results.get(\"image_fields\", [\"image\"]):\n mask_blocks = torch.zeros(size=(B, 1, *down_size), device=device)\n results[key + \"_mask\"] = mask_blocks\n return results\n\n def random_sampling(self, results, mask_ratio, down_size, validity_mask, device):\n B = results[\"image\"].shape[0]\n prob_blocks = torch.rand(size=(B, 1, *down_size), device=device)\n mask_blocks = torch.logical_and(prob_blocks < mask_ratio, validity_mask)\n return mask_blocks\n\n def inverse_sampling(self, results, mask_ratio, down_size, validity_mask, device):\n # from PIL import Image\n # from unidepth.utils import colorize\n def area_sample(depth, fx, fy):\n dtype = depth.dtype\n B = depth.shape[0]\n H, W = down_size\n depth = downsample(depth, depth.shape[-2] // H)\n depth[depth > 200] = 50 # set sky as if depth 50 meters\n pixel_area3d = depth / torch.sqrt(fx * fy)\n\n # Set invalid as -1 (no div problem) -> then clip to 0.0\n pixel_area3d[depth == 0.0] = -1\n prob_density = (1 / pixel_area3d).clamp(min=0.0).square()\n prob_density = prob_density / prob_density.sum(\n dim=(-1, -2), keepdim=True\n ).clamp(min=1e-5)\n # Image.fromarray((prob_density[0] * 255 * 100).clamp(min=0.0, max=255.0).squeeze().cpu().byte().numpy()).save(\"prob_density.png\")\n\n # Sample locations based on prob_density\n prob_density_flat = prob_density.view(B, -1)\n\n # Get the avgerage valid locations, of those we mask self.mask_ratio\n valid_locations = (prob_density_flat > 0).to(dtype).sum(dim=1)\n\n masks = []\n for i in range(B):\n num_samples = int(valid_locations[i] * mask_ratio)\n mask = torch.zeros_like(prob_density_flat[i])\n # Sample indices\n if num_samples > 0:\n sampled_indices_flat = torch.multinomial(\n prob_density_flat[i], num_samples, replacement=False\n )\n mask.scatter_(0, sampled_indices_flat, 1)\n masks.append(mask)\n return torch.stack(masks).bool().view(B, 1, H, W)\n\n def random_sample(validity_mask):\n prob_blocks = torch.rand(\n size=(validity_mask.shape[0], 1, *down_size), device=device\n )\n mask = torch.logical_and(prob_blocks < mask_ratio, validity_mask)\n return mask\n\n fx = results[\"K\"][..., 0, 0].view(-1, 1, 1, 1) / self.mask_patch\n fy = results[\"K\"][..., 1, 1].view(-1, 1, 1, 1) / self.mask_patch\n\n valid = ~results[\"ssi\"] & ~results[\"si\"] & results[\"valid_camera\"]\n mask_blocks = torch.zeros_like(validity_mask)\n if valid.any():\n out = area_sample(results[\"depth\"][valid], fx[valid], fy[valid])\n mask_blocks[valid] = out\n if (~valid).any():\n mask_blocks[~valid] = random_sample(validity_mask[~valid])\n\n # mask_blocks_ = (mask_blocks.float() * 255).squeeze(1).byte().cpu().numpy()\n # Image.fromarray(mask_blocks_[0]).save(\"mask1.png\")\n # Image.fromarray(mask_blocks_[-1]).save(\"mask2.png\")\n # dd = results[\"depth\"]\n # Image.fromarray(colorize(dd[0].squeeze().cpu().numpy())).save(\"depth1_p.png\")\n # Image.fromarray(colorize(dd[-1].squeeze().cpu().numpy())).save(\"depth2_p.png\")\n # dd = downsample(dd, dd.shape[-2] // down_size[0])\n # Image.fromarray(colorize(dd[0].squeeze().cpu().numpy())).save(\"depth1.png\")\n # Image.fromarray(colorize(dd[-1].squeeze().cpu().numpy())).save(\"depth2.png\")\n # raise ValueError\n\n return mask_blocks\n\n def guided_sampling(self, results, mask_ratio, down_size, validity_mask, device):\n # get the lowest (based on guidance) \"mask_ratio\" quantile of the patches that are in validity mask\n B = results[\"image\"].shape[0]\n guidance = results[\"guidance\"]\n mask_blocks = torch.zeros(size=(B, 1, *down_size), device=device)\n for b in range(B):\n low_bound = torch.quantile(\n guidance[b][validity_mask[b]], max(0.0, self.hard_bound - mask_ratio)\n )\n up_bound = torch.quantile(\n guidance[b][validity_mask[b]], min(1.0, self.hard_bound)\n )\n mask_blocks[b] = torch.logical_and(\n guidance[b] < up_bound, guidance[b] > low_bound\n )\n mask_blocks = torch.logical_and(mask_blocks, validity_mask)\n return mask_blocks\n\n def step(self):\n self.idx += 1\n # schedule hard from 1.0 to self.mask_ratio\n if self.curriculum:\n step = max(0, self.idx / self.warmup_steps / 2 - 0.5)\n self.hard_bound = 1 - (1 - self.mask_ratio) * tanh(step)\n self.up_bound = self.mask_ratio * tanh(step)\n self.low_bound = 0.2 * tanh(step)\n\n @property\n def is_warmup(self):\n return self.idx < self.warmup_steps\n\n\nclass Rotate:\n def __init__(\n self, angle, center=None, img_fill_val=(123.68, 116.28, 103.53), prob=0.5\n ):\n if isinstance(img_fill_val, (float, int)):\n img_fill_val = tuple([float(img_fill_val)] * 3)\n elif isinstance(img_fill_val, tuple):\n assert len(img_fill_val) == 3, (\n \"image_fill_val as tuple must \"\n f\"have 3 elements. got {len(img_fill_val)}.\"\n )\n img_fill_val = tuple([float(val) for val in img_fill_val])\n else:\n raise ValueError(\"image_fill_val must be float or tuple with 3 elements.\")\n assert np.all(\n [0 <= val <= 255 for val in img_fill_val]\n ), f\"all elements of img_fill_val should between range [0,255] got {img_fill_val}.\"\n assert 0 <= prob <= 1.0, f\"The probability should be in range [0,1]bgot {prob}.\"\n self.center = center\n self.img_fill_val = img_fill_val\n self.prob = prob\n self.random = not isinstance(angle, (float, int))\n self.angle = angle\n\n def _rotate(self, results, angle, center=None, fill_val=0.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n img = results[key]\n img_rotated = TF.rotate(\n img,\n angle,\n center=center,\n interpolation=TF.InterpolationMode.NEAREST_EXACT,\n fill=self.img_fill_val,\n )\n results[key] = img_rotated.to(img.dtype)\n results[\"image_shape\"] = results[key].shape\n\n for key in results.get(\"mask_fields\", []):\n results[key] = TF.rotate(\n results[key],\n angle,\n center=center,\n interpolation=TF.InterpolationMode.NEAREST_EXACT,\n fill=fill_val,\n )\n\n for key in results.get(\"gt_fields\", []):\n results[key] = TF.rotate(\n results[key],\n angle,\n center=center,\n interpolation=TF.InterpolationMode.NEAREST_EXACT,\n fill=fill_val,\n )\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n\n angle = (\n (self.angle[1] - self.angle[0]) * np.random.rand() + self.angle[0]\n if self.random\n else np.random.choice([-1, 1], size=1) * self.angle\n )\n self._rotate(results, angle, None, fill_val=0.0)\n results[\"rotation\"] = angle\n return results\n\n\nclass RandomColor:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def _adjust_color_img(self, results, factor=1.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.adjust_hue(results[key], factor) # .to(img.dtype)\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n factor = (\n ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else self.level\n )\n self._adjust_color_img(results, factor)\n return results\n\n\nclass RandomSaturation:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def _adjust_saturation_img(self, results, factor=1.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n # NOTE defaultly the image should be BGR format\n results[key] = TF.adjust_saturation(results[key], factor) # .to(img.dtype)\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n factor = (\n 2 ** ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else 2**self.level\n )\n self._adjust_saturation_img(results, factor)\n return results\n\n\nclass RandomSharpness:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def _adjust_sharpeness_img(self, results, factor=1.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n # NOTE defaultly the image should be BGR format\n results[key] = TF.adjust_sharpness(results[key], factor) # .to(img.dtype)\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n factor = (\n 2 ** ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else 2**self.level\n )\n self._adjust_sharpeness_img(results, factor)\n return results\n\n def __repr__(self):\n repr_str = self.__class__.__name__\n repr_str += f\"(level={self.level}, \"\n repr_str += f\"prob={self.prob})\"\n return repr_str\n\n\nclass RandomSolarize:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def _adjust_solarize_img(self, results, factor=255.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.solarize(results[key], factor) # .to(img.dtype)\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n factor = (\n ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else self.level\n )\n self._adjust_solarize_img(results, factor)\n return results\n\n\nclass RandomPosterize:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def _posterize_img(self, results, factor=1.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.posterize(results[key], int(factor)) # .to(img.dtype)\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n factor = (\n ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else self.level\n )\n self._posterize_img(results, factor)\n return results\n\n\nclass RandomEqualize:\n def __init__(self, prob=0.5):\n assert 0 <= prob <= 1.0, \"The probability should be in range [0,1].\"\n self.prob = prob\n\n def _imequalize(self, results):\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.equalize(results[key]) # .to(img.dtype)\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n self._imequalize(results)\n return results\n\n\nclass RandomBrightness:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def _adjust_brightness_img(self, results, factor=1.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.adjust_brightness(results[key], factor) # .to(img.dtype)\n\n def __call__(self, results, level=None):\n if np.random.random() > self.prob:\n return results\n factor = (\n 2 ** ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else 2**self.level\n )\n self._adjust_brightness_img(results, factor)\n return results\n\n\nclass RandomContrast:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def _adjust_contrast_img(self, results, factor=1.0):\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.adjust_contrast(results[key], factor) # .to(img.dtype)\n\n def __call__(self, results, level=None):\n if np.random.random() > self.prob:\n return results\n factor = (\n 2 ** ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else 2**self.level\n )\n self._adjust_contrast_img(results, factor)\n return results\n\n\nclass RandomGamma:\n def __init__(self, level, prob=0.5):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n\n def __call__(self, results, level=None):\n if np.random.random() > self.prob:\n return results\n factor = (self.level[1] - self.level[0]) * np.random.rand() + self.level[0]\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n results[key] = TF.adjust_gamma(results[key], 1 + factor)\n return results\n\n\nclass RandomInvert:\n def __init__(self, prob=0.5):\n self.prob = prob\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n results[key] = TF.invert(results[key]) # .to(img.dtype)\n return results\n\n\nclass RandomAutoContrast:\n def __init__(self, prob=0.5):\n self.prob = prob\n\n def _autocontrast_img(self, results):\n for key in results.get(\"image_fields\", [\"image\"]):\n img = results[key]\n results[key] = TF.autocontrast(img) # .to(img.dtype)\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n self._autocontrast_img(results)\n return results\n\n\nclass RandomShear(object):\n def __init__(\n self,\n level,\n prob=0.5,\n direction=\"horizontal\",\n ):\n self.random = not isinstance(level, (float, int))\n self.level = level\n self.prob = prob\n self.direction = direction\n\n def _shear_img(self, results, magnitude):\n for key in results.get(\"image_fields\", [\"image\"]):\n img_sheared = TF.affine(\n results[key],\n angle=0.0,\n translate=[0.0, 0.0],\n scale=1.0,\n shear=magnitude,\n interpolation=TF.InterpolationMode.BILINEAR,\n fill=0.0,\n )\n results[key] = img_sheared\n\n def _shear_masks(self, results, magnitude):\n for key in results.get(\"mask_fields\", []):\n mask_sheared = TF.affine(\n results[key],\n angle=0.0,\n translate=[0.0, 0.0],\n scale=1.0,\n shear=magnitude,\n interpolation=TF.InterpolationMode.NEAREST_EXACT,\n fill=0.0,\n )\n results[key] = mask_sheared\n\n def _shear_gt(\n self,\n results,\n magnitude,\n ):\n for key in results.get(\"gt_fields\", []):\n mask_sheared = TF.affine(\n results[key],\n angle=0.0,\n translate=[0.0, 0.0],\n scale=1.0,\n shear=magnitude,\n interpolation=TF.InterpolationMode.NEAREST_EXACT,\n fill=0.0,\n )\n results[key] = mask_sheared\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n magnitude = (\n ((self.level[1] - self.level[0]) * np.random.rand() + self.level[0])\n if self.random\n else np.random.choice([-1, 1], size=1) * self.level\n )\n if self.direction == \"horizontal\":\n magnitude = [magnitude, 0.0]\n else:\n magnitude = [0.0, magnitude]\n self._shear_img(results, magnitude)\n self._shear_masks(results, magnitude)\n self._shear_gt(results, magnitude)\n return results\n\n\nclass RandomTranslate(object):\n def __init__(\n self,\n range,\n prob=0.5,\n direction=\"horizontal\",\n ):\n self.range = range\n self.prob = prob\n self.direction = direction\n\n def _translate_img(self, results, magnitude):\n \"\"\"Shear the image.\n\n Args:\n results (dict): Result dict from loading pipeline.\n magnitude (int | float): The magnitude used for shear.\n direction (str): The direction for shear, either \"horizontal\"\n or \"vertical\".\n interpolation (str): Same as in :func:`mmcv.imshear`.\n \"\"\"\n for key in results.get(\"image_fields\", [\"image\"]):\n img_sheared = TF.affine(\n results[key],\n angle=0.0,\n translate=magnitude,\n scale=1.0,\n shear=[0.0, 0.0],\n interpolation=TF.InterpolationMode.BILINEAR,\n fill=(123.68, 116.28, 103.53),\n )\n results[key] = img_sheared\n\n def _translate_mask(self, results, magnitude):\n \"\"\"Shear the masks.\"\"\"\n for key in results.get(\"mask_fields\", []):\n mask_sheared = TF.affine(\n results[key],\n angle=0.0,\n translate=magnitude,\n scale=1.0,\n shear=[0.0, 0.0],\n interpolation=TF.InterpolationMode.NEAREST_EXACT,\n fill=0.0,\n )\n results[key] = mask_sheared\n\n def _translate_gt(\n self,\n results,\n magnitude,\n ):\n \"\"\"Shear the segmentation maps.\"\"\"\n for key in results.get(\"gt_fields\", []):\n mask_sheared = TF.affine(\n results[key],\n angle=0.0,\n translate=magnitude,\n scale=1.0,\n shear=[0.0, 0.0],\n interpolation=TF.InterpolationMode.NEAREST_EXACT,\n fill=0.0,\n )\n results[key] = mask_sheared\n\n def __call__(self, results):\n \"\"\"Call function to shear images, bounding boxes, masks and semantic\n segmentation maps.\n\n Args:\n results (dict): Result dict from loading pipeline.\n\n Returns:\n dict: Sheared results.\n \"\"\"\n if np.random.random() > self.prob:\n return results\n magnitude = (self.range[1] - self.range[0]) * np.random.rand() + self.range[0]\n if self.direction == \"horizontal\":\n magnitude = [magnitude * results[\"image\"].shape[1], 0]\n else:\n magnitude = [0, magnitude * results[\"image\"].shape[0]]\n self._translate_img(results, magnitude)\n self._translate_mask(results, magnitude)\n self._translate_gt(results, magnitude)\n results[\"K\"][..., 0, 2] = results[\"K\"][..., 0, 2] + magnitude[0]\n results[\"K\"][..., 1, 2] = results[\"K\"][..., 1, 2] + magnitude[1]\n return results\n\n def __repr__(self):\n repr_str = self.__class__.__name__\n repr_str += f\"(range={self.range}, \"\n repr_str += f\"prob={self.prob}, \"\n repr_str += f\"direction={self.direction}, \"\n return repr_str\n\n\nclass RandomColorJitter:\n def __init__(self, level, prob=0.9):\n self.level = level\n self.prob = prob\n self.list_transform = [\n self._adjust_brightness_img,\n # self._adjust_sharpness_img,\n self._adjust_contrast_img,\n self._adjust_saturation_img,\n self._adjust_color_img,\n ]\n\n def _adjust_contrast_img(self, results, factor=1.0):\n \"\"\"Adjust the image contrast.\"\"\"\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n img = results[key]\n results[key] = TF.adjust_contrast(img, factor)\n\n def _adjust_sharpness_img(self, results, factor=1.0):\n \"\"\"Adjust the image contrast.\"\"\"\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n img = results[key]\n results[key] = TF.adjust_sharpness(img, factor)\n\n def _adjust_brightness_img(self, results, factor=1.0):\n \"\"\"Adjust the brightness of image.\"\"\"\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n img = results[key]\n results[key] = TF.adjust_brightness(img, factor)\n\n def _adjust_saturation_img(self, results, factor=1.0):\n \"\"\"Apply Color transformation to image.\"\"\"\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n img = results[key]\n results[key] = TF.adjust_saturation(img, factor / 2.0)\n\n def _adjust_color_img(self, results, factor=1.0):\n \"\"\"Apply Color transformation to image.\"\"\"\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n img = results[key]\n results[key] = TF.adjust_hue(img, (factor - 1.0) / 4.0)\n\n def __call__(self, results):\n \"\"\"Call function for color transformation.\n Args:\n results (dict): Results dict from loading pipeline.\n\n Returns:\n dict: Results after the transformation.\n \"\"\"\n random.shuffle(self.list_transform)\n for op in self.list_transform:\n if np.random.random() < self.prob:\n factor = 1.0 + (\n (self.level[1] - self.level[0]) * np.random.random() + self.level[0]\n )\n op(results, factor)\n return results\n\n\nclass RandomGrayscale:\n def __init__(self, prob=0.1, num_output_channels=3):\n super().__init__()\n self.prob = prob\n self.num_output_channels = num_output_channels\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n results[key] = TF.rgb_to_grayscale(\n results[key], num_output_channels=self.num_output_channels\n )\n return results\n\n\ndef masked_nearest_interpolation(input, mask, target_size):\n \"\"\"\n Resize the depth map using bilinear interpolation, considering only valid pixels within NxN neighbors.\n\n Args:\n depth (torch.Tensor): The depth map tensor of shape (H, W).\n mask (torch.Tensor): The mask tensor of shape (H, W) where 1 indicates valid depth and 0 indicates missing depth.\n target_size (tuple): The desired output size (target_H, target_W).\n\n Returns:\n torch.Tensor: The resized depth map.\n \"\"\"\n B, C, H, W = input.shape\n target_H, target_W = target_size\n mask = mask.float()\n\n # Generate a grid of coordinates in the target space\n grid_y, grid_x = torch.meshgrid(\n torch.linspace(0, H - 1, target_H),\n torch.linspace(0, W - 1, target_W),\n indexing=\"ij\",\n )\n grid_y = grid_y.to(input.device)\n grid_x = grid_x.to(input.device)\n\n # Calculate the floor and ceil of the grid coordinates to get the bounding box\n x0 = torch.floor(grid_x).long().clamp(0, W - 1)\n x1 = (x0 + 1).clamp(0, W - 1)\n y0 = torch.floor(grid_y).long().clamp(0, H - 1)\n y1 = (y0 + 1).clamp(0, H - 1)\n\n # Gather depth values at the four corners\n Ia = input[..., y0, x0]\n Ib = input[..., y1, x0]\n Ic = input[..., y0, x1]\n Id = input[..., y1, x1]\n\n # Gather corresponding mask values\n ma = mask[..., y0, x0]\n mb = mask[..., y1, x0]\n mc = mask[..., y0, x1]\n md = mask[..., y1, x1]\n\n # Calculate distances to each neighbor\n # The distances are calculated from the center (grid_x, grid_y) to each corner\n dist_a = (grid_x - x0.float()) ** 2 + (grid_y - y0.float()) ** 2 # Top-left\n dist_b = (grid_x - x0.float()) ** 2 + (grid_y - y1.float()) ** 2 # Bottom-left\n dist_c = (grid_x - x1.float()) ** 2 + (grid_y - y0.float()) ** 2 # Top-right\n dist_d = (grid_x - x1.float()) ** 2 + (grid_y - y1.float()) ** 2 # Bottom-right\n\n # Stack the neighbors, their masks, and distances\n stacked_values = torch.stack(\n [Ia, Ib, Ic, Id], dim=-1\n ) # Shape: (B, C, target_H, target_W, 4)\n stacked_masks = torch.stack(\n [ma, mb, mc, md], dim=-1\n ) # Shape: (B, 1, target_H, target_W, 4)\n stacked_distances = torch.stack(\n [dist_a, dist_b, dist_c, dist_d], dim=-1\n ) # Shape: (target_H, target_W, 4)\n stacked_distances = (\n stacked_distances.unsqueeze(0).unsqueeze(1).repeat(B, 1, 1, 1, 1)\n ) # Shape: (B, 1, target_H, target_W, 4)\n\n # Set distances to infinity for invalid neighbors (so that invalid neighbors are never chosen)\n stacked_distances[stacked_masks == 0] = float(\"inf\")\n\n # Find the index of the nearest valid neighbor (the one with the smallest distance)\n nearest_indices = stacked_distances.argmin(dim=-1, keepdim=True)[\n ..., :1\n ] # Shape: (B, 1, target_H, target_W, 1)\n\n # Select the corresponding depth value using the nearest valid neighbor index\n interpolated_depth = torch.gather(\n stacked_values, dim=-1, index=nearest_indices.repeat(1, C, 1, 1, 1)\n ).squeeze(-1)\n\n # Set depth to zero where no valid neighbors were found\n interpolated_depth = interpolated_depth * stacked_masks.sum(dim=-1).clip(\n min=0.0, max=1.0\n )\n\n return interpolated_depth\n\n\nclass ContextCrop:\n def __init__(\n self,\n image_shape,\n keep_original=False,\n test_min_ctx=1.0,\n train_ctx_range=[0.5, 1.5],\n shape_constraints={},\n ):\n self.image_shape = image_shape\n self.keep_original = keep_original\n self.test_min_ctx = test_min_ctx\n self.train_ctx_range = train_ctx_range\n self.shape_mult = shape_constraints[\"shape_mult\"]\n self.sample = shape_constraints[\"sample\"]\n self.ratio_bounds = shape_constraints[\"ratio_bounds\"]\n pixels_min = shape_constraints[\"pixels_min\"] / (\n self.shape_mult * self.shape_mult\n )\n pixels_max = shape_constraints[\"pixels_max\"] / (\n self.shape_mult * self.shape_mult\n )\n self.pixels_bounds = (pixels_min, pixels_max)\n self.ctx = None\n\n def _transform_img(self, results, shapes):\n for key in results.get(\"image_fields\", [\"image\"]):\n img = self.crop(results[key], **shapes)\n img = TF.resize(\n img,\n results[\"resized_shape\"],\n interpolation=TF.InterpolationMode.BICUBIC,\n antialias=True,\n )\n results[key] = img\n\n def _transform_masks(self, results, shapes):\n for key in results.get(\"mask_fields\", []):\n mask = self.crop(results[key].float(), **shapes).byte()\n mask = masked_nearest_interpolation(\n mask, mask > 0, results[\"resized_shape\"]\n )\n results[key] = mask\n\n def _transform_gt(self, results, shapes):\n for key in results.get(\"gt_fields\", []):\n gt = self.crop(results[key], **shapes)\n gt = masked_nearest_interpolation(gt, gt > 0, results[\"resized_shape\"])\n results[key] = gt\n\n @staticmethod\n def crop(img, height, width, top, left) -> torch.Tensor:\n h, w = img.shape[-2:]\n right = left + width\n bottom = top + height\n padding_ltrb = [\n max(-left + min(0, right), 0),\n max(-top + min(0, bottom), 0),\n max(right - max(w, left), 0),\n max(bottom - max(h, top), 0),\n ]\n image_cropped = img[..., max(top, 0) : bottom, max(left, 0) : right]\n return TF.pad(image_cropped, padding_ltrb)\n\n def test_closest_shape(self, image_shape):\n h, w = image_shape\n input_ratio = w / h\n if self.sample:\n input_pixels = int(ceil(h / self.shape_mult * w / self.shape_mult))\n pixels = max(\n min(input_pixels, self.pixels_bounds[1]), self.pixels_bounds[0]\n )\n ratio = min(max(input_ratio, self.ratio_bounds[0]), self.ratio_bounds[1])\n h = round((pixels / ratio) ** 0.5)\n w = h * ratio\n self.image_shape[0] = int(h) * self.shape_mult\n self.image_shape[1] = int(w) * self.shape_mult\n\n def _get_crop_shapes(self, image_shape, ctx=None):\n h, w = image_shape\n input_ratio = w / h\n if self.keep_original:\n self.test_closest_shape(image_shape)\n ctx = 1.0\n elif ctx is None:\n ctx = float(\n torch.empty(1)\n .uniform_(self.train_ctx_range[0], self.train_ctx_range[1])\n .item()\n )\n output_ratio = self.image_shape[1] / self.image_shape[0]\n\n if output_ratio <= input_ratio: # out like 4:3 in like kitti\n if (\n ctx >= 1\n ): # fully in -> use just max_length with sqrt(ctx), here max is width\n new_w = w * ctx**0.5\n # sporge un po in una sola dim\n # we know that in_width will stick out before in_height, partial overshoot (sporge)\n # new_h > old_h via area -> new_h ** 2 * ratio_new = old_h ** 2 * ratio_old * ctx\n elif output_ratio / input_ratio * ctx > 1:\n new_w = w * ctx\n else: # fully contained -> use area\n new_w = w * (ctx * output_ratio / input_ratio) ** 0.5\n new_h = new_w / output_ratio\n else:\n if ctx >= 1:\n new_h = h * ctx**0.5\n elif input_ratio / output_ratio * ctx > 1:\n new_h = h * ctx\n else:\n new_h = h * (ctx * input_ratio / output_ratio) ** 0.5\n new_w = new_h * output_ratio\n return (int(ceil(new_h - 0.5)), int(ceil(new_w - 0.5))), ctx\n\n def __call__(self, results):\n h, w = results[\"image\"].shape[-2:]\n results[\"image_ori_shape\"] = (h, w)\n\n results.get(\"mask_fields\", set()).add(\"validity_mask\")\n if \"validity_mask\" not in results:\n results[\"validity_mask\"] = torch.ones(\n (results[\"image\"].shape[0], 1, h, w),\n dtype=torch.uint8,\n device=results[\"image\"].device,\n )\n\n n_iter = 1 if self.keep_original or not self.sample else 100\n\n min_valid_area = 0.5\n results[\"camera_fields\"].add(\"camera_original\")\n results[\"camera_original\"] = results[\"camera\"].clone()\n max_hfov, max_vfov = results[\"camera\"].max_fov[0] # it is a 1-dim list\n ctx = None\n for ii in range(n_iter):\n\n (height, width), ctx = self._get_crop_shapes((h, w), ctx=self.ctx or ctx)\n margin_h = h - height\n margin_w = w - width\n\n # keep it centered in y direction\n top = margin_h // 2\n left = margin_w // 2\n if not self.keep_original:\n left = left + np.random.randint(\n -self.shape_mult // 2, self.shape_mult // 2 + 1\n )\n top = top + np.random.randint(\n -self.shape_mult // 2, self.shape_mult // 2 + 1\n )\n\n right = left + width\n bottom = top + height\n x_zoom = self.image_shape[0] / height\n paddings = [\n max(-left + min(0, right), 0),\n max(bottom - max(h, top), 0),\n max(right - max(w, left), 0),\n max(-top + min(0, bottom), 0),\n ]\n\n valid_area = (\n h\n * w\n / (h + paddings[1] + paddings[3])\n / (w + paddings[0] + paddings[2])\n )\n\n new_hfov, new_vfov = results[\"camera_original\"].get_new_fov(\n new_shape=(height, width), original_shape=(h, w)\n )[0]\n\n if (\n valid_area >= min_valid_area\n and new_hfov < max_hfov\n and new_vfov < max_vfov\n ):\n results[\"camera\"] = results[\"camera\"].crop(\n left, top, right=w - right, bottom=h - bottom\n )\n results[\"camera\"] = results[\"camera\"].resize(x_zoom)\n break\n ctx = (\n ctx * 0.96\n ) # if not enough valid area, try again with less ctx (more zoom)\n\n # save ctx for next iteration of sequences?\n self.ctx = ctx\n\n results[\"resized_shape\"] = self.image_shape\n results[\"paddings\"] = paddings # left ,top ,right, bottom\n results[\"image_rescale\"] = x_zoom\n results[\"scale_factor\"] = results.get(\"scale_factor\", 1.0) * x_zoom\n\n shapes = dict(height=height, width=width, top=top, left=left)\n self._transform_img(results, shapes)\n if not self.keep_original:\n self._transform_gt(results, shapes)\n self._transform_masks(results, shapes)\n else:\n # only validity_mask (rgb's masks follows rgb transform) #FIXME\n mask = results[\"validity_mask\"].float()\n mask = self.crop(mask, **shapes).byte()\n mask = TF.resize(\n mask,\n results[\"resized_shape\"],\n interpolation=TF.InterpolationMode.NEAREST,\n )\n results[\"validity_mask\"] = mask\n\n # keep original images before photo-augment\n results[\"image_original\"] = results[\"image\"].clone()\n results[\"image_fields\"].add(\n *[\n field.replace(\"image\", \"image_original\")\n for field in results[\"image_fields\"]\n ]\n )\n\n # repeat for batch resized shape and paddings\n results[\"paddings\"] = [results[\"paddings\"]] * results[\"image\"].shape[0]\n results[\"resized_shape\"] = [results[\"resized_shape\"]] * results[\"image\"].shape[\n 0\n ]\n return results\n\n\nclass RandomFiller:\n def __init__(self, *args, **kwargs):\n super().__init__()\n\n def _transform(self, results):\n def fill_noise(size, device):\n return torch.normal(0, 2.0, size=size, device=device)\n\n def fill_black(size, device):\n return -4 * torch.ones(size, device=device, dtype=torch.float32)\n\n def fill_white(size, device):\n return 4 * torch.ones(size, device=device, dtype=torch.float32)\n\n def fill_zero(size, device):\n return torch.zeros(size, device=device, dtype=torch.float32)\n\n B, C = results[\"image\"].shape[:2]\n mismatch = B // results[\"validity_mask\"].shape[0]\n if mismatch:\n results[\"validity_mask\"] = results[\"validity_mask\"].repeat(\n mismatch, 1, 1, 1\n )\n validity_mask = results[\"validity_mask\"].repeat(1, C, 1, 1).bool()\n filler_fn = np.random.choice([fill_noise, fill_black, fill_white, fill_zero])\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key][~validity_mask] = filler_fn(\n size=results[key][~validity_mask].shape, device=results[key].device\n )\n\n def __call__(self, results):\n # generate mask for filler\n if \"validity_mask\" not in results:\n paddings = results.get(\"padding_size\", [0] * 4)\n height, width = results[\"image\"].shape[-2:]\n results.get(\"mask_fields\", []).add(\"validity_mask\")\n results[\"validity_mask\"] = torch.zeros_like(results[\"image\"][:, :1])\n results[\"validity_mask\"][\n ...,\n paddings[1] : height - paddings[3],\n paddings[0] : width - paddings[2],\n ] = 1.0\n self._transform(results)\n return results\n\n\nclass GaussianBlur:\n def __init__(self, kernel_size, sigma=(0.1, 2.0), prob=0.9):\n super().__init__()\n self.kernel_size = kernel_size\n self.sigma = sigma\n self.prob = prob\n self.padding = kernel_size // 2\n\n def apply(self, x, kernel):\n # Pad the input tensor\n x = F.pad(\n x, (self.padding, self.padding, self.padding, self.padding), mode=\"reflect\"\n )\n # Apply the convolution with the Gaussian kernel\n return F.conv2d(x, kernel, stride=1, padding=0, groups=x.size(1))\n\n def _create_kernel(self, sigma):\n # Create a 1D Gaussian kernel\n kernel_1d = torch.exp(\n -torch.arange(-self.padding, self.padding + 1) ** 2 / (2 * sigma**2)\n )\n kernel_1d = kernel_1d / kernel_1d.sum()\n\n # Expand the kernel to 2D and match size of the input\n kernel_2d = kernel_1d.unsqueeze(0) * kernel_1d.unsqueeze(1)\n kernel_2d = kernel_2d.view(1, 1, self.kernel_size, self.kernel_size).expand(\n 3, 1, -1, -1\n )\n return kernel_2d\n\n def __call__(self, results):\n if np.random.random() > self.prob:\n return results\n sigma = (self.sigma[1] - self.sigma[0]) * np.random.rand() + self.sigma[0]\n kernel = self._create_kernel(sigma)\n for key in results.get(\"image_fields\", [\"image\"]):\n if \"original\" not in key:\n results[key] = self.apply(results[key], kernel)\n return results\n\n\nclass Compose:\n def __init__(self, transforms):\n self.transforms = deepcopy(transforms)\n\n def __call__(self, results):\n for t in self.transforms:\n results = t(results)\n return results\n\n def __setattr__(self, name: str, value) -> None:\n super().__setattr__(name, value)\n for t in self.transforms:\n setattr(t, name, value)\n\n def __repr__(self):\n format_string = self.__class__.__name__ + \"(\"\n for t in self.transforms:\n format_string += f\"\\n {t}\"\n format_string += \"\\n)\"\n return format_string\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/point_odyssey.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass PointOdyssey(SequenceDataset):\n min_depth = 0.01\n max_depth = 250.0\n depth_scale = 1000.0\n test_split = \"test.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences_clean.json\"\n hdf5_paths = [f\"PointOdyssey.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/proteus.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Proteus(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n default_fps = 5\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"Proteus.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/samplers copy.py", "content": "import itertools\nimport warnings\nfrom operator import itemgetter\nfrom typing import Any, Optional\n\nimport numpy as np\nimport torch\nfrom torch.utils.data import Sampler\n\nfrom unidepth.utils import get_dist_info\n\n\ndef _get_numpy_dtype(size: int) -> Any:\n return np.int32 if size <= 2**31 else np.int64\n\n\ndef _get_torch_dtype(size: int) -> Any:\n return torch.int32 if size <= 2**31 else torch.int64\n\n\ndef _generate_randperm_indices(*, size: int, generator: torch.Generator):\n \"\"\"Generate the indices of a random permutation.\"\"\"\n dtype = _get_torch_dtype(size)\n # This is actually matching PyTorch's CPU implementation, see: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TensorFactories.cpp#L900-L921\n perm = torch.arange(size, dtype=dtype)\n for i in range(size):\n j = torch.randint(i, size, size=(1,), generator=generator).item()\n\n # Always swap even if no-op\n value = perm[j].item()\n perm[j] = perm[i].item()\n perm[i] = value\n yield value\n\n\n# The following function is somewhat equivalent to _new_shuffle_tensor_slice below,\n# but avoids a full in-place random permutation generation.\ndef _shuffle_tensor_slice(\n *, tensor: torch.Tensor, start: int = 0, step: int = 1, generator: torch.Generator\n) -> np.ndarray:\n stop = len(tensor)\n count = stop // step\n drop_count = stop - step * count\n if drop_count:\n warnings.warn(f\"# of dropped samples: {drop_count}\")\n\n dtype = _get_numpy_dtype(stop)\n result = np.empty(count, dtype=dtype)\n\n for i in range(count):\n j = (\n torch.randint(0, i + 1, size=(1,), generator=generator).item()\n if i > 0\n else 0\n )\n\n result[i] = result[j]\n result[j] = tensor[start + i * step].item()\n\n return result\n\n\ndef _new_shuffle_tensor_slice(\n *, tensor: torch.Tensor, start: int = 0, step: int = 1, generator: torch.Generator\n) -> np.ndarray:\n stop = len(tensor)\n count = stop // step\n dtype = torch.int64 # Needed for using randperm result as indices\n count = stop // step\n drop_count = stop - step * count\n if drop_count:\n warnings.warn(f\"# of dropped samples: {drop_count}\")\n indices = torch.randperm(count, dtype=dtype, generator=generator)\n return tensor[start::step][indices].numpy()\n\n\ndef _make_seed(seed: int, start: int, iter_count: int) -> int:\n # NOTE: Tried a few variants (including iter_count << 32), this one worked best.\n return seed + start + (iter_count << 24)\n\n\nclass ShardedInfiniteSampler(Sampler):\n def __init__(\n self,\n *,\n sample_count: int,\n shuffle: bool = False,\n seed: int = 0,\n start: Optional[int] = None,\n step: Optional[int] = None,\n advance: int = 0,\n use_new_shuffle_tensor_slice: bool = False,\n ):\n self._sample_count = sample_count\n self._seed = seed\n self._shuffle = shuffle\n rank, world_size = get_dist_info()\n self._start = rank if start is None else start\n self._step = world_size if step is None else step\n self._advance = advance\n self._iter_count = 0\n self._shuffle_tensor_slice_fn = (\n _new_shuffle_tensor_slice\n if use_new_shuffle_tensor_slice\n else _shuffle_tensor_slice\n )\n\n def __iter__(self):\n iter_count = self._advance // self._sample_count\n if iter_count > 0:\n self._advance -= iter_count * self._sample_count\n self._iter_count += iter_count\n\n if self._shuffle:\n iterator = self._shuffled_iterator()\n else:\n iterator = self._iterator()\n\n yield from itertools.islice(iterator, self._advance, None)\n\n def _iterator(self):\n assert not self._shuffle\n\n while True:\n iterable = range(self._sample_count)\n yield from itertools.islice(iterable, self._start, None, self._step)\n\n def _shuffled_iterator(self):\n assert self._shuffle\n\n # Instantiate a generator here (rather than in the ctor) to be keep the class\n # picklable (requirement of mp.spawn)\n generator = torch.Generator()\n\n # Always shuffle everything first\n generator.manual_seed(self._seed)\n dtype = _get_torch_dtype(self._sample_count)\n perm = torch.randperm(self._sample_count, dtype=dtype, generator=generator)\n\n while True:\n # Re-seed on each iteration to allow skipping whole permutations\n seed = _make_seed(self._seed, self._start, self._iter_count)\n generator.manual_seed(seed)\n\n iterable = self._shuffle_tensor_slice_fn(\n tensor=perm, start=self._start, step=self._step, generator=generator\n )\n yield from iterable\n self._iter_count += 1\n\n\nclass DistributedSamplerNoDuplicate(torch.utils.data.DistributedSampler):\n \"\"\"A distributed sampler that doesn't add duplicates. Arguments are the same as DistributedSampler\"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n if not self.drop_last and len(self.dataset) % self.num_replicas != 0:\n # some ranks may have less samples, that's fine\n if self.rank >= len(self.dataset) % self.num_replicas:\n self.num_samples -= 1\n self.total_size = len(self.dataset)\n\n\nclass DatasetFromSampler(torch.utils.data.Dataset):\n \"\"\"Dataset to create indexes from `Sampler`.\n\n Args:\n sampler: PyTorch sampler\n \"\"\"\n\n def __init__(self, sampler: Sampler):\n \"\"\"Initialisation for DatasetFromSampler.\"\"\"\n self.sampler = sampler\n self.sampler_list = None\n\n def __getitem__(self, index: int):\n \"\"\"Gets element of the dataset.\n\n Args:\n index: index of the element in the dataset\n\n Returns:\n Single element by index\n \"\"\"\n if self.sampler_list is None:\n self.sampler_list = list(self.sampler)\n return self.sampler_list[index]\n\n def __len__(self) -> int:\n \"\"\"\n Returns:\n int: length of the dataset\n \"\"\"\n return len(self.sampler)\n\n\nclass DistributedSamplerWrapper(torch.utils.data.DistributedSampler):\n \"\"\"\n Wrapper over `Sampler` for distributed training\n Allows you to use any sampler in distributed mode.\n\n It is especially useful in conjunction with\n `torch.nn.parallel.DistributedDataParallel`. In such case, each\n process can pass a DistributedSamplerWrapper instance as a DataLoader\n sampler, and load a subset of subsampled data of the original dataset\n that is exclusive to it.\n\n .. note::\n Sampler is assumed to be of constant size.\n \"\"\"\n\n def __init__(\n self,\n sampler,\n num_replicas: Optional[int] = None,\n rank: Optional[int] = None,\n shuffle: bool = True,\n ):\n \"\"\"\n\n Args:\n sampler: Sampler used for subsampling\n num_replicas (int, optional): Number of processes participating in\n distributed training\n rank (int, optional): Rank of the current process\n within ``num_replicas``\n shuffle (bool, optional): If true (default),\n sampler will shuffle the indices\n \"\"\"\n super(DistributedSamplerWrapper, self).__init__(\n DatasetFromSampler(sampler),\n num_replicas=num_replicas,\n rank=rank,\n shuffle=shuffle,\n )\n self.sampler = sampler\n\n def __iter__(self):\n self.dataset = DatasetFromSampler(self.sampler)\n indexes_of_indexes = super().__iter__()\n subsampler_indexes = self.dataset\n return iter(itemgetter(*indexes_of_indexes)(subsampler_indexes))\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/samplers.py", "content": "import torch\n\n\nclass DistributedSamplerNoDuplicate(torch.utils.data.DistributedSampler):\n \"\"\"A distributed sampler that doesn't add duplicates. Arguments are the same as DistributedSampler\"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n if not self.drop_last and len(self.dataset) % self.num_replicas != 0:\n # some ranks may have less samples, that's fine\n if self.rank >= len(self.dataset) % self.num_replicas:\n self.num_samples -= 1\n self.total_size = len(self.dataset)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/scannet.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass ScanNet(SequenceDataset):\n min_depth = 0.005\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"test.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"ScanNetS.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/scannetpp.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass ScanNetpp(SequenceDataset):\n min_depth = 0.001\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"val_iphone.txt\"\n train_split = \"train_iphone.txt\"\n sequences_file = \"sequences_iphone_clean.json\"\n hdf5_paths = [f\"ScanNetpp_viz.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n\n\nclass ScanNetpp_F(SequenceDataset):\n min_depth = 0.001\n max_depth = 10.0\n depth_scale = 1000.0\n train_split = \"train.txt\"\n test_split = \"val_split.txt\"\n\n sequences_file = \"sequences_split.json\"\n hdf5_paths = [f\"ScanNetpp_F.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=(\n decode_fields if not test_mode else [*decode_fields, \"points\"]\n ),\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/sequence_dataset.py", "content": "import json\nimport os\nfrom functools import partial\nfrom typing import Any, Dict, Tuple\n\nimport h5py\nimport numpy as np\nimport tables\nimport torch\nimport torchvision.transforms.v2.functional as TF\n\nfrom unidepth.datasets.base_dataset import BaseDataset\nfrom unidepth.datasets.utils import DatasetFromList\nfrom unidepth.datasets.utils_decode import (decode_camera, decode_depth,\n decode_flow, decode_K, decode_mask,\n decode_numpy, decode_rgb,\n decode_tensor)\nfrom unidepth.utils.distributed import is_main_process\n\n\nclass SequenceDataset(BaseDataset):\n DECODE_FNS = {\n \"image\": partial(decode_rgb, name=\"image\"),\n \"points\": partial(decode_numpy, name=\"points\"),\n \"K\": partial(decode_K, name=\"camera\"),\n \"camera_params\": partial(decode_camera, name=\"camera\"),\n \"cam2w\": partial(decode_tensor, name=\"cam2w\"),\n \"depth\": partial(decode_depth, name=\"depth\"),\n \"flow_fwd\": partial(decode_flow, name=\"flow_fwd\"),\n \"flow_bwd\": partial(decode_flow, name=\"flow_bwd\"),\n \"flow_fwd_mask\": partial(decode_mask, name=\"flow_fwd_mask\"),\n \"flow_bwd_mask\": partial(decode_mask, name=\"flow_bwd_mask\"),\n }\n default_fps = 5\n\n def __init__(\n self,\n image_shape: Tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: Dict[str, Any],\n resize_method: str,\n mini: float,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\", \"flow_fwd\", \"flow_fwd_mask\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.num_frames = num_frames\n self.original_num_frames = num_frames\n self.decode_fields = decode_fields\n self.inplace_fields = inplace_fields\n self.fps = self.default_fps\n self.fps_range = kwargs.get(\"fps_range\", None)\n if self.fps_range is not None:\n self.fps_range[1] = min(self.default_fps, self.fps_range[1])\n\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\").strip()\n sequences = np.array(h5file[self.sequences_file]).tostring().decode(\"ascii\")\n sequences = json.loads(sequences)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n if len(line.strip().split(\" \")) == 1:\n print(line)\n continue\n sequence_name, num_samples = line.strip().split(\" \")\n dataset.append(\n {\n \"sequence_name\": sequence_name,\n \"num_samples\": int(num_samples),\n \"chunk_idx\": 0,\n }\n )\n\n # filter dataset based on attr \"invalid_sequences\"\n invalid_sequences = getattr(self, \"invalid_sequences\", [])\n dataset = [\n sample\n for sample in dataset\n if sample[\"sequence_name\"] not in invalid_sequences\n ]\n\n self.dataset = DatasetFromList(dataset)\n self.sequences = DatasetFromList(\n [sequences[sample[\"sequence_name\"]] for sample in dataset]\n )\n self.log_load_dataset()\n\n def get_random_idxs(self, num_samples_sequence):\n if self.num_frames == 1:\n return [np.random.randint(0, num_samples_sequence)], 0\n\n # Check if we can satisfy the required number of frames\n if self.num_frames > num_samples_sequence:\n raise ValueError(\n \"Cannot sample more frames than available in the sequence.\"\n )\n\n # Restrict FPS range to be within default FPS\n min_fps, max_fps = self.fps_range\n max_fps = min(max_fps, self.default_fps)\n if min_fps > self.default_fps:\n sampled_fps = self.default_fps\n else:\n # Compute minimal viable FPS\n min_required_fps = (\n self.num_frames / num_samples_sequence\n ) * self.default_fps\n min_fps = max(min_fps, min_required_fps)\n\n # Sample an FPS from the viable range\n sampled_fps = np.random.uniform(min_fps, max_fps)\n\n # Compute the stride based on the sampled FPS\n stride = self.default_fps / sampled_fps\n max_start_index = num_samples_sequence - int(stride * (self.num_frames - 1))\n\n # Ensure a valid starting position\n if max_start_index <= 0:\n raise ValueError(\n \"No valid start position allows sampling num_frames with the chosen FPS.\"\n )\n\n start_index = np.random.randint(0, max_start_index + 1)\n\n # Compute indices based on the sampled FPS\n indices = [int(start_index + i * stride) for i in range(self.num_frames)]\n\n return indices, np.random.randint(0, len(indices))\n\n def get_test_idxs(self, num_samples_sequence, keyframe_idx):\n if self.num_frames == 1:\n return [\n keyframe_idx if keyframe_idx is not None else num_samples_sequence // 2\n ], 0\n\n if self.num_frames == -1:\n cap_idxs = min(32, num_samples_sequence) # CAP 32 images\n idxs = list(\n range(max(0, num_samples_sequence - cap_idxs), num_samples_sequence, 1)\n )\n return idxs, keyframe_idx\n\n # pick closest keyframe_idx st they are around it or capped by the 0 and max num_samples_sequence\n keyframe_idx = (\n keyframe_idx if keyframe_idx is not None else num_samples_sequence - 1\n )\n excess_tail = 0 - min(0, keyframe_idx - self.num_frames // 2)\n excess_head = (\n max(num_samples_sequence, keyframe_idx + (self.num_frames - 1) // 2)\n - num_samples_sequence\n )\n start = keyframe_idx - self.num_frames // 2 + excess_tail - excess_head\n end = keyframe_idx + (self.num_frames - 1) // 2 + excess_head - excess_tail\n idxs = list(range(start, 1 + end))\n\n return idxs, idxs.index(keyframe_idx)\n\n def get_single_sequence(self, idx):\n self.num_frames = self.original_num_frames\n # sequence_name = self.dataset[idx][\"sequence_name\"]\n sample = self.sequences[idx]\n chunk_idx = int(sample.get(\"chunk_idx\", 0))\n h5_path = os.path.join(self.data_root, self.hdf5_paths[chunk_idx])\n\n num_samples_sequence = len(sample[\"image\"])\n if self.num_frames > 0 and num_samples_sequence < self.num_frames:\n raise IndexError(f\"Sequence {idx} has less than {self.num_frames} frames\")\n keyframe_idx = None\n\n if not self.test_mode:\n idxs, keyframe_idx = self.get_random_idxs(num_samples_sequence)\n else:\n idxs, keyframe_idx = self.get_test_idxs(\n num_samples_sequence, sample.get(\"keyframe_idx\", None)\n )\n\n self.num_frames = len(idxs)\n results = {}\n results = self.pre_pipeline(results)\n results[\"sequence_fields\"] = [(i, 0) for i in range(self.num_frames)]\n results[\"keyframe_idx\"] = keyframe_idx\n with tables.File(\n h5_path,\n mode=\"r\",\n libver=\"latest\",\n swmr=True,\n ) as h5file_chunk:\n\n for i, j in enumerate(idxs):\n results[(i, 0)] = {\n k: v.copy() for k, v in results.items() if \"fields\" in k\n }\n for inplace_field in self.inplace_fields:\n inplace_field_ = inplace_field.replace(\"intrinsics\", \"K\").replace(\n \"extrinsics\", \"cam2w\"\n )\n results = self.DECODE_FNS[inplace_field_](\n results, sample[inplace_field][j], idx=i, sample=sample, j=j\n )\n\n for i, j in enumerate(idxs):\n for decode_field in self.decode_fields:\n results = self.DECODE_FNS[decode_field](\n results,\n h5file_chunk,\n sample[decode_field][j],\n idx=i,\n depth_scale=self.depth_scale,\n )\n\n results[\"filename\"] = sample[\"image\"][j]\n\n results = self.preprocess(results)\n if not self.test_mode:\n results = self.augment(results)\n results = self.postprocess(results)\n return results\n\n def preprocess(self, results):\n results = self.replicate(results)\n for i, seq in enumerate(results[\"sequence_fields\"]):\n results[seq] = self.resizer(results[seq])\n self.resizer.ctx = None if self.num_copies > 1 else self.resizer.ctx\n num_pts = torch.count_nonzero(results[seq][\"depth\"] > 0)\n if num_pts < 50:\n raise IndexError(f\"Too few points in depth map ({num_pts})\")\n\n for key in results[seq].get(\"image_fields\", [\"image\"]):\n results[seq][key] = results[seq][key].to(torch.float32) / 255\n\n # update fields common in sequence\n for key in [\n \"image_fields\",\n \"gt_fields\",\n \"mask_fields\",\n \"camera_fields\",\n ]:\n if key in results[(0, 0)]:\n results[key] = results[(0, 0)][key]\n\n results = self.pack_batch(results)\n return results\n\n def postprocess(self, results):\n # # normalize after because color aug requires [0,255]?\n for key in results.get(\"image_fields\", [\"image\"]):\n results[key] = TF.normalize(results[key], **self.normalization_stats)\n results = self.filler(results)\n results = self.unpack_batch(results)\n results = self.masker(results)\n results = self.collecter(results)\n return results\n\n def __getitem__(self, idx):\n try:\n if isinstance(idx, (list, tuple)):\n results = [self.get_single_sequence(i) for i in idx]\n else:\n results = self.get_single_sequence(idx)\n except Exception as e:\n print(f\"Error loading sequence {idx} for {self.__class__.__name__}: {e}\")\n idx = np.random.randint(0, len(self.dataset))\n results = self[idx]\n return results\n\n def log_load_dataset(self):\n if is_main_process():\n info = f\"Loaded {self.__class__.__name__} with {sum([len(x['image']) for x in self.sequences])} images in {len(self)} sequences.\"\n print(info)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/sintel copy.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Sintel(SequenceDataset):\n min_depth = 0.001\n max_depth = 1000.0\n depth_scale = 1000.0\n test_split = \"training.txt\"\n train_split = \"training.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"Sintel.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\", \"flow_fwd\", \"flow_fwd_mask\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/sintel.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Sintel(SequenceDataset):\n min_depth = 0.001\n max_depth = 1000.0\n depth_scale = 1000.0\n test_split = \"training.txt\"\n train_split = \"training.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"Sintel.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames\n results[\"synthetic\"] = [True] * self.num_frames\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/sunrgbd.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass SUNRGBD(ImageDataset):\n min_depth = 0.005\n max_depth = 8.0\n depth_scale = 1000.0\n test_split = \"alltest.txt\"\n train_split = \"alltrain.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"SUNRGB.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/synscapes.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Synscapes(SequenceDataset):\n min_depth = 0.1\n max_depth = 1000.0\n depth_scale = 256.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"Synscapes.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/tartanair.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass TartanAir(SequenceDataset):\n min_depth = 0.01\n max_depth = 512.0\n depth_scale = 1000.0\n default_fps = 15\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"TartanAir.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/taskonomy.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass Taskonomy(ImageDataset):\n min_depth = 0.005\n max_depth = 15.0\n depth_scale = 512.0\n test_split = \"val.txt\"\n train_split = \"train_clean.txt\"\n intrisics_file = \"intrinsics.json\"\n hdf5_paths = [\"Taskonomy.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\") # [:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.split_file), \"w\") as f:\n # f.write(txt_string)\n # with open(os.path.join(os.environ[\"TMPDIR\"], self.intrisics_file), \"w\") as f:\n # json.dump(intrinsics, f)\n\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename, chunk_idx = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val, chunk_idx]\n dataset.append(sample)\n h5file.close()\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n if self.test_mode and not self.benchmark:\n dataset = self.chunk(dataset, chunk_dim=1, pct=0.01)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def get_mapper(self):\n return {\n \"image_filename\": 0,\n \"depth_filename\": 1,\n \"K\": 2,\n }\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/tat_rmvd.py", "content": "import json\nimport os\nfrom copy import deepcopy\nfrom typing import Any\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.pipelines import AnnotationMask, KittiCrop\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\nfrom unidepth.datasets.utils import DatasetFromList\nfrom unidepth.utils import identity\n\n\nclass TATRMVD(SequenceDataset):\n min_depth = 0.001\n max_depth = 50.0\n depth_scale = 1000.0\n default_fps = 6\n test_split = \"test.txt\"\n train_split = \"test.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"tanks_and_temples_rmvd.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"si\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [2] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/theo.py", "content": "from typing import Any\n\nimport torch\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Theo(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n default_fps = 5\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"THEO.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"camera_params\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def preprocess(self, results):\n self.resizer.ctx = None\n for i, seq in enumerate(results[\"sequence_fields\"]):\n # Create a mask where the distance from the center is less than H/2\n H, W = results[seq][\"image\"].shape[-2:]\n x = torch.linspace(-(W - 1) / 2, (W - 1) / 2, W)\n y = torch.linspace(-(H - 1) / 2, (H - 1) / 2, H)\n xv, yv = torch.meshgrid(x, y, indexing=\"xy\")\n distance_from_center = torch.sqrt(xv**2 + yv**2).reshape(1, 1, H, W)\n results[seq][\"validity_mask\"] = distance_from_center < (H - 1) / 2\n\n return super().preprocess(results)\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/unrealstereo4k.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass UnrealStereo4K(SequenceDataset):\n min_depth = 0.01\n max_depth = 200.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"UnrealStereo4K.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/urbansyn.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass UrbanSyn(SequenceDataset):\n min_depth = 0.1\n max_depth = 1000.0\n depth_scale = 256.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"UrbanSyn.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/utils.py", "content": "import copy\nimport multiprocessing as mp\nimport pickle\nfrom collections import defaultdict\nfrom typing import Any, Dict, List\n\nimport numpy as np\nimport torch\nimport torch.utils.data\n\nfrom unidepth.utils.distributed import (all_gather, get_local_rank,\n get_local_size, get_rank,\n get_world_size)\n\n\nclass ConcatDataset(torch.utils.data.ConcatDataset):\n def __init__(self, datasets, shape_constraints: dict[str, list[int]] = {}):\n super().__init__(datasets)\n\n self.sample = shape_constraints[\"sample\"]\n self.shape_mult = shape_constraints[\"shape_mult\"]\n self.ratio_bounds = shape_constraints[\"ratio_bounds\"]\n self.pixels_max = shape_constraints[\"pixels_max\"]\n self.pixels_min = shape_constraints[\"pixels_min\"]\n\n self.height_min = shape_constraints[\"height_min\"]\n self.width_min = shape_constraints[\"width_min\"]\n\n def sample_shape(self):\n if not self.sample:\n return\n # 1: sample image ratio\n ratio = np.random.uniform(*self.ratio_bounds)\n pixels_min = self.pixels_min // (self.shape_mult * self.shape_mult)\n pixels_max = self.pixels_max // (self.shape_mult * self.shape_mult)\n # 2: sample image height or width, if ratio > 1 or < 1\n if ratio > 1:\n height_min = max(self.height_min, np.sqrt(pixels_min / ratio))\n height = np.random.uniform(height_min, np.sqrt(pixels_max / ratio))\n width = height * ratio\n else:\n width_min = max(self.width_min, np.sqrt(pixels_min * ratio))\n width = np.random.uniform(width_min, np.sqrt(pixels_max * ratio))\n height = width / ratio\n # 3: get final shape based on the shape_mult\n shape = [int(height) * self.shape_mult, int(width) * self.shape_mult]\n for dataset in self.datasets:\n setattr(dataset, \"image_shape\", shape)\n setattr(dataset.resizer, \"image_shape\", shape)\n\n def __getitem__(self, idxs):\n self.sample_shape()\n return [super(ConcatDataset, self).__getitem__(idx) for idx in idxs]\n\n\ndef _paddings(image_shape, network_shape):\n cur_h, cur_w = image_shape\n h, w = network_shape\n pad_top, pad_bottom = (h - cur_h) // 2, h - cur_h - (h - cur_h) // 2\n pad_left, pad_right = (w - cur_w) // 2, w - cur_w - (w - cur_w) // 2\n return pad_left, pad_right, pad_top, pad_bottom\n\n\ndef collate_fn(in_data: List[List[Dict[str, Any]]], is_batched: bool = True):\n out_data = defaultdict(list)\n img_metas = []\n in_data = in_data[0] if is_batched else in_data\n\n # get max_shape and paddings\n shapes = [tensor.shape[-2:] for x in in_data for tensor in x[\"depth\"].values()]\n max_shape_tuple = tuple(max(elements) for elements in zip(*shapes))\n paddings = [\n [\n _paddings(tensor.shape[-2:], max_shape_tuple)\n for tensor in x[\"depth\"].values()\n ]\n for x in in_data\n ]\n\n for x in in_data: # here iter over batches\n padding = paddings.pop(0)\n for k, v in x.items():\n if \"img_metas\" not in k:\n values = list(v.values())\n v = torch.cat(values)\n out_data[k].append(v)\n else:\n v[\"depth_paddings\"] = padding\n img_metas.append(v)\n\n output_dict = {\n \"data\": {k: torch.stack(v, dim=0) for k, v in out_data.items()},\n \"img_metas\": img_metas,\n }\n # camera are always flattened and the stack/cat so if list of B times (T, 3, 3) cameras\n # it goes to (B * T, 3, 3), to be consistent with the image shape -> reshape\n if \"camera\" in output_dict[\"data\"]:\n output_dict[\"data\"][\"camera\"] = output_dict[\"data\"][\"camera\"].reshape(\n *output_dict[\"data\"][\"image\"].shape[:2]\n )\n return output_dict\n\n\ndef local_scatter(array: list[Any]):\n \"\"\"\n Scatter an array from local leader to all local workers.\n The i-th local worker gets array[i].\n\n Args:\n array: Array with same size of #local workers.\n \"\"\"\n if get_world_size() == 1:\n return array[0]\n if get_local_rank() == 0:\n assert len(array) == get_local_size()\n all_gather(array)\n else:\n all_data = all_gather(None)\n array = all_data[get_rank() - get_local_rank()]\n return array[get_local_rank()]\n\n\nclass DatasetFromList(torch.utils.data.Dataset): # type: ignore\n \"\"\"Wrap a list to a torch Dataset.\n\n We serialize and wrap big python objects in a torch.Dataset due to a\n memory leak when dealing with large python objects using multiple workers.\n See: https://github.com/pytorch/pytorch/issues/13246\n \"\"\"\n\n def __init__(self, lst: List[Any], deepcopy: bool = False, serialize: bool = True):\n \"\"\"Creates an instance of the class.\n\n Args:\n lst: a list which contains elements to produce.\n deepcopy: whether to deepcopy the element when producing it, s.t.\n the result can be modified in place without affecting the source\n in the list.\n serialize: whether to hold memory using serialized objects. When\n enabled, data loader workers can use shared RAM from master\n process instead of making a copy.\n \"\"\"\n self._copy = deepcopy\n self._serialize = serialize\n\n def _serialize(data: Any):\n buffer = pickle.dumps(data, protocol=pickle.HIGHEST_PROTOCOL)\n return torch.frombuffer(buffer, dtype=torch.uint8)\n\n if self._serialize:\n # load only on 0th rank\n if get_local_rank() == 0:\n _lst = [_serialize(x) for x in lst]\n self._addr = torch.cumsum(\n torch.tensor([len(x) for x in _lst], dtype=torch.int64), dim=0\n )\n self._lst = torch.concatenate(_lst)\n # Move data to shared memory, obtain a handle to send to each local worker.\n handles = [None] + [\n bytes(mp.reduction.ForkingPickler.dumps((self._addr, self._lst)))\n for _ in range(get_local_size() - 1)\n ]\n else:\n handles = None\n\n # Each worker receives the handle from local leader (rank 0)\n # then materialize the tensor from shared memory\n handle = local_scatter(handles)\n if get_local_rank() > 0:\n self._addr, self._lst = mp.reduction.ForkingPickler.loads(handle)\n\n else:\n self._lst = lst\n\n def __len__(self) -> int:\n \"\"\"Return len of list.\"\"\"\n if self._serialize:\n return len(self._addr)\n return len(self._lst)\n\n def __getitem__(self, idx: int) -> Any:\n \"\"\"Return item of list at idx.\"\"\"\n if self._serialize:\n start_addr = 0 if idx == 0 else self._addr[idx - 1]\n end_addr = self._addr[idx]\n bytes_ = memoryview(self._lst[start_addr:end_addr].numpy())\n return pickle.loads(bytes_)\n if self._copy:\n return copy.deepcopy(self._lst[idx])\n\n return self._lst[idx]\n\n\ndef get_weights(\n train_datasets: dict[str, torch.utils.data.Dataset], sampling: dict[str, float]\n) -> torch.Tensor:\n from .image_dataset import ImageDataset\n from .sequence_dataset import SequenceDataset\n\n weights = []\n num_samples = 0\n info_weights = {}\n for dataset_name, dataset in train_datasets.items():\n assert (\n dataset_name in sampling\n ), f\"Dataset {dataset_name} not found in {sampling.keys()}\"\n\n if isinstance(dataset, ImageDataset):\n # sum of all samples has weight as in sampling s.t. sampling dataset in general is as in sampling\n # inside is uniform\n weight = sampling[dataset_name] / len(dataset)\n weights.append(torch.full((len(dataset),), weight).double())\n num_samples += len(dataset)\n\n elif isinstance(dataset, SequenceDataset):\n # local weight is num_samples, but global must be as in sampling\n # hence is num_samples / (sum num_samples / sampling[dataset_name])\n # s.t. sampling anything from the dataset is\n # sum(num_samples / (sum num_samples / sampling[dataset_name]))\n # -> sampling[dataset_name]\n numerator = [int(data[\"num_samples\"]) for data in dataset.dataset]\n weights.append(\n sampling[dataset_name]\n * torch.tensor(numerator).double()\n / sum(numerator)\n )\n num_samples += sum(numerator)\n\n else:\n weight = sampling[dataset_name] / len(dataset)\n weights.append(torch.full((len(dataset),), weight).double())\n\n info_weights[dataset_name] = weights[-1][-1]\n\n return torch.cat(weights), num_samples\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/utils_decode.py", "content": "import io\n\nimport cv2\nimport numpy as np\nimport torch\nimport torchvision\nimport torchvision.transforms.v2.functional as TF\nfrom PIL import Image\n\nfrom unidepth.utils.camera import (EUCM, MEI, BatchCamera, Fisheye624, Pinhole,\n Spherical)\n\n\ndef decode_depth(results, h5file, value, idx, depth_scale, name=\"depth\", **kwargs):\n file = h5file.get_node(\"/\" + value).read()\n decoded_data = Image.open(io.BytesIO(file))\n decoded_data = TF.pil_to_tensor(decoded_data).squeeze()\n\n if decoded_data.ndim == 3: # 24 channel loading\n decoded_channels = [\n (decoded_data[0] & 0xFF).to(torch.int32),\n (decoded_data[1] & 0xFF).to(torch.int32),\n (decoded_data[2] & 0xFF).to(torch.int32),\n ]\n # Reshape and extract the original depth map\n decoded_data = (\n decoded_channels[0]\n | (decoded_channels[1] << 8)\n | (decoded_channels[2] << 16)\n )\n\n decoded_data = decoded_data.to(torch.float32)\n results.get(\"gt_fields\", set()).add(name)\n results[(idx, 0)].get(\"gt_fields\", set()).add(name)\n results[f\"{name}_ori_shape\"] = decoded_data.shape\n results[(idx, 0)][name] = (\n decoded_data.view(1, 1, *decoded_data.shape).contiguous() / depth_scale\n )\n return results\n\n\ndef decode_numpy(results, h5file, value, idx, name=\"points\", **kwargs):\n file = h5file.get_node(\"/\" + value).read()\n decoded_data = np.load(io.BytesIO(file), allow_pickle=False)\n decoded_data = torch.from_numpy(decoded_data).to(torch.float32)\n if decoded_data.ndim > 2:\n decoded_data = decoded_data.permute(2, 0, 1)\n results.get(\"gt_fields\", set()).add(name)\n results[(idx, 0)].get(\"gt_fields\", set()).add(name)\n results[(idx, 0)][name] = decoded_data.unsqueeze(0)\n return results\n\n\ndef decode_tensor(results, value, idx, name, **kwargs):\n results.get(\"camera_fields\", set()).add(name)\n results[(idx, 0)].get(\"camera_fields\", set()).add(name)\n results[(idx, 0)][name] = torch.tensor(value).unsqueeze(0)\n return results\n\n\ndef decode_camera(results, value, idx, name, sample, j, **kwargs):\n results.get(\"camera_fields\", set()).add(name)\n results[(idx, 0)].get(\"camera_fields\", set()).add(name)\n camera = eval(sample[\"camera_model\"][j])(params=torch.tensor(value).unsqueeze(0))\n results[(idx, 0)][name] = BatchCamera.from_camera(camera)\n return results\n\n\ndef decode_K(results, value, idx, name, **kwargs):\n results.get(\"camera_fields\", set()).add(name)\n results[(idx, 0)].get(\"camera_fields\", set()).add(name)\n camera = Pinhole(K=torch.tensor(value).unsqueeze(0))\n results[(idx, 0)][name] = BatchCamera.from_camera(camera)\n return results\n\n\ndef decode_mask(results, h5file, value, idx, name, **kwargs):\n file = h5file.get_node(\"/\" + value).read()\n mask = torchvision.io.decode_image(torch.from_numpy(file)).bool().squeeze()\n results.get(\"mask_fields\", set()).add(name)\n results[(idx, 0)].get(\"mask_fields\", set()).add(name)\n results[f\"{name}_ori_shape\"] = mask.shape[-2:]\n results[(idx, 0)][name] = mask.view(1, 1, *mask.shape).contiguous()\n return results\n\n\ndef decode_rgb(results, h5file, value, idx, name=\"image\", **kwargs):\n file = h5file.get_node(\"/\" + value).read()\n image = (\n torchvision.io.decode_image(torch.from_numpy(file)).to(torch.uint8).squeeze()\n )\n results.get(\"image_fields\", set()).add(name)\n results[(idx, 0)].get(\"image_fields\", set()).add(name)\n results[f\"{name}_ori_shape\"] = image.shape[-2:]\n if image.ndim == 2:\n image = image.unsqueeze(0).repeat(3, 1, 1)\n results[(idx, 0)][name] = image.unsqueeze(0)\n return results\n\n\ndef decode_flow(results, h5file, value, idx, name, **kwargs):\n file = h5file.get_node(\"/\" + value).read()\n image = (\n torchvision.io.decode_image(torch.from_numpy(file)).to(torch.uint8).squeeze()\n )\n decoded_channels = [\n (image[0] & 0xFF).to(torch.int16),\n (image[1] & 0xFF).to(torch.int16),\n (image[2] & 0xFF).to(torch.int16),\n ]\n\n # Reshape and extract the original 2-channel flow map\n flow = torch.zeros((2, image.shape[1], image.shape[2]), dtype=torch.int16)\n flow[0] = (decoded_channels[0] | decoded_channels[1] << 8) & 0xFFF\n flow[1] = (decoded_channels[1] >> 4 | decoded_channels[2] << 4) & 0xFFF\n\n results.get(\"gt_fields\", set()).add(name)\n results[(idx, 0)].get(\"gt_fields\", set()).add(name)\n results[f\"{name}_ori_shape\"] = flow.shape[-2:]\n flow = flow.unsqueeze(0).contiguous().float()\n results[(idx, 0)][name] = (0.5 + flow) / 4095.0 * 2 - 1\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/vkitti.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass VKITTI(SequenceDataset):\n min_depth = 0.01\n max_depth = 255.0\n depth_scale = 256.0\n test_split = \"training.txt\"\n train_split = \"training.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"VKITTI2.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\", \"flow_fwd\", \"flow_fwd_mask\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [0] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/void.py", "content": "import json\nimport os\n\nimport h5py\nimport numpy as np\nimport torch\n\nfrom unidepth.datasets.image_dataset import ImageDataset\nfrom unidepth.datasets.utils import DatasetFromList\n\n\nclass VOID(ImageDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 256.0\n test_split = \"void_val.txt\"\n train_split = \"void_train.txt\"\n intrisics_file = \"void_intrinsics.json\"\n hdf5_paths = [\"void.hdf5\"]\n\n def __init__(\n self,\n image_shape,\n split_file,\n test_mode,\n crop=None,\n benchmark=False,\n augmentations_db={},\n normalize=True,\n resize_method=\"hard\",\n mini=1.0,\n **kwargs,\n ):\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n **kwargs,\n )\n self.test_mode = test_mode\n\n self.crop = crop\n self.load_dataset()\n\n def load_dataset(self):\n h5file = h5py.File(\n os.path.join(self.data_root, self.hdf5_paths[0]),\n \"r\",\n libver=\"latest\",\n swmr=True,\n )\n txt_file = np.array(h5file[self.split_file])\n txt_string = txt_file.tostring().decode(\"ascii\")[:-1] # correct the -1\n intrinsics = np.array(h5file[self.intrisics_file]).tostring().decode(\"ascii\")\n intrinsics = json.loads(intrinsics)\n h5file.close()\n dataset = []\n for line in txt_string.split(\"\\n\"):\n image_filename, depth_filename = line.strip().split(\" \")\n intrinsics_val = torch.tensor(intrinsics[image_filename]).squeeze()[:, :3]\n sample = [image_filename, depth_filename, intrinsics_val]\n dataset.append(sample)\n\n if not self.test_mode:\n dataset = self.chunk(dataset, chunk_dim=1, pct=self.mini)\n\n self.dataset = DatasetFromList(dataset)\n self.log_load_dataset()\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_copies\n results[\"quality\"] = [2] * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/waymo.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass Waymo(SequenceDataset):\n min_depth = 0.05\n max_depth = 70.0\n depth_scale = 256.0\n test_split = \"validation.txt\"\n train_split = \"training.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [f\"Waymo_viz.hdf5\"]\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [False] * self.num_frames * self.num_copies\n results[\"synthetic\"] = [False] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/datasets/wildrgbd.py", "content": "from typing import Any\n\nfrom unidepth.datasets.sequence_dataset import SequenceDataset\n\n\nclass WildRGBD(SequenceDataset):\n min_depth = 0.01\n max_depth = 10.0\n depth_scale = 1000.0\n test_split = \"train.txt\"\n train_split = \"train.txt\"\n sequences_file = \"sequences.json\"\n hdf5_paths = [\"WildRGBD.hdf5\"]\n default_fps = 30\n\n def __init__(\n self,\n image_shape: tuple[int, int],\n split_file: str,\n test_mode: bool,\n normalize: bool,\n augmentations_db: dict[str, Any],\n resize_method: str,\n mini: float = 1.0,\n num_frames: int = 1,\n benchmark: bool = False,\n decode_fields: list[str] = [\"image\", \"depth\"],\n inplace_fields: list[str] = [\"K\", \"cam2w\"],\n **kwargs,\n ) -> None:\n super().__init__(\n image_shape=image_shape,\n split_file=split_file,\n test_mode=test_mode,\n benchmark=benchmark,\n normalize=normalize,\n augmentations_db=augmentations_db,\n resize_method=resize_method,\n mini=mini,\n num_frames=num_frames,\n decode_fields=decode_fields,\n inplace_fields=inplace_fields,\n **kwargs,\n )\n\n def pre_pipeline(self, results):\n results = super().pre_pipeline(results)\n results[\"dense\"] = [True] * self.num_frames * self.num_copies\n results[\"quality\"] = [1] * self.num_frames * self.num_copies\n return results\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/__init__.py", "content": "from .activation import GEGLU, SwiGLU\nfrom .attention import AttentionBlock, AttentionDecoderBlock, AttentionLayer\nfrom .convnext import CvnxtBlock\nfrom .mlp import MLP\nfrom .nystrom_attention import NystromBlock\nfrom .positional_encoding import PositionEmbeddingSine\nfrom .upsample import (ConvUpsample, ConvUpsampleShuffle,\n ConvUpsampleShuffleResidual, ResUpsampleBil)\n\n__all__ = [\n \"SwiGLU\",\n \"GEGLU\",\n \"CvnxtBlock\",\n \"AttentionBlock\",\n \"NystromBlock\",\n \"PositionEmbeddingSine\",\n \"ConvUpsample\",\n \"MLP\",\n \"ConvUpsampleShuffle\",\n \"AttentionDecoderBlock\",\n \"ConvUpsampleShuffleResidual\",\n]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/activation.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass SwiGLU(nn.Module):\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x, gates = x.chunk(2, dim=-1)\n return x * F.silu(gates)\n\n\nclass GEGLU(nn.Module):\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x, gates = x.chunk(2, dim=-1)\n return x * F.gelu(gates)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/attention.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom functools import partial\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\n\nfrom .layer_scale import LayerScale\nfrom .mlp import MLP\n\n\nclass SimpleAttention(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int = 4,\n dropout: float = 0.0,\n cosine: bool = False,\n context_dim: int | None = None,\n ):\n super().__init__()\n self.dropout = dropout\n self.num_heads = num_heads\n self.hidden_dim = dim\n context_dim = context_dim or dim\n\n self.kv = nn.Linear(context_dim, dim * 2, bias=False)\n self.q = nn.Linear(dim, dim, bias=False)\n self.norm_attnx = nn.LayerNorm(dim)\n self.norm_attnctx = nn.LayerNorm(context_dim)\n self.cosine = cosine\n self.out = nn.Linear(dim, dim)\n\n def forward(\n self,\n x: torch.Tensor,\n attn_bias: torch.Tensor | None = None,\n context: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n pos_embed_context: torch.Tensor | None = None,\n rope: nn.Module | None = None,\n ) -> torch.Tensor:\n context = x if context is None else context\n x = self.norm_attnx(x)\n context = self.norm_attnctx(context)\n k, v = rearrange(\n self.kv(context), \"b n (kv h d) -> b h n d kv\", h=self.num_heads, kv=2\n ).unbind(dim=-1)\n q = rearrange(self.q(x), \"b n (h d) -> b h n d\", h=self.num_heads)\n\n if rope is not None:\n q = rope(q)\n k = rope(k)\n else:\n if pos_embed is not None:\n pos_embed = rearrange(\n pos_embed, \"b n (h d) -> b h n d\", h=self.num_heads\n )\n q = q + pos_embed\n if pos_embed_context is not None:\n pos_embed_context = rearrange(\n pos_embed_context, \"b n (h d) -> b h n d\", h=self.num_heads\n )\n k = k + pos_embed_context\n\n if self.cosine:\n q, k = map(partial(F.normalize, p=2, dim=-1), (q, k)) # cosine sim\n x = F.scaled_dot_product_attention(\n q, k, v, dropout_p=self.dropout, attn_mask=attn_bias\n )\n x = rearrange(x, \"b h n d -> b n (h d)\")\n x = self.out(x)\n return x\n\n\nclass AttentionBlock(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int = 4,\n expansion: int = 4,\n dropout: float = 0.0,\n cosine: bool = False,\n gated: bool = False,\n layer_scale: float = 1.0,\n context_dim: int | None = None,\n use_bias: bool = True,\n ):\n super().__init__()\n self.dropout = dropout\n self.num_heads = num_heads\n self.hidden_dim = dim\n context_dim = context_dim or dim\n self.mlp = MLP(dim, expansion=expansion, dropout=dropout, gated=gated)\n self.kv = nn.Linear(context_dim, dim * 2, bias=use_bias)\n self.q = nn.Linear(dim, dim, bias=use_bias)\n self.norm_attnx = nn.LayerNorm(dim)\n self.norm_attnctx = nn.LayerNorm(context_dim)\n self.cosine = cosine\n self.out = nn.Linear(dim, dim, bias=use_bias)\n self.ls1 = LayerScale(dim, layer_scale) if layer_scale > 0.0 else nn.Identity()\n self.ls2 = LayerScale(dim, layer_scale) if layer_scale > 0.0 else nn.Identity()\n\n def attn(\n self,\n x: torch.Tensor,\n attn_bias: torch.Tensor | None = None,\n context: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n pos_embed_context: torch.Tensor | None = None,\n ) -> torch.Tensor:\n x = self.norm_attnx(x)\n context = self.norm_attnctx(context)\n k, v = rearrange(\n self.kv(context), \"b n (kv h d) -> b h n d kv\", h=self.num_heads, kv=2\n ).unbind(dim=-1)\n q = rearrange(self.q(x), \"b n (h d) -> b h n d\", h=self.num_heads)\n\n if pos_embed is not None:\n pos_embed = rearrange(pos_embed, \"b n (h d) -> b h n d\", h=self.num_heads)\n q = q + pos_embed\n if pos_embed_context is not None:\n pos_embed_context = rearrange(\n pos_embed_context, \"b n (h d) -> b h n d\", h=self.num_heads\n )\n k = k + pos_embed_context\n\n if self.cosine:\n q, k = map(partial(F.normalize, p=2, dim=-1), (q, k)) # cosine sim\n\n x = F.scaled_dot_product_attention(\n q, k, v, dropout_p=self.dropout, attn_mask=attn_bias\n )\n x = rearrange(x, \"b h n d -> b n (h d)\")\n x = self.out(x)\n return x\n\n def forward(\n self,\n x: torch.Tensor,\n attn_bias: torch.Tensor | None = None,\n context: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n pos_embed_context: torch.Tensor | None = None,\n ) -> torch.Tensor:\n context = x if context is None else context\n x = (\n self.ls1(\n self.attn(\n x,\n attn_bias=attn_bias,\n context=context,\n pos_embed=pos_embed,\n pos_embed_context=pos_embed_context,\n )\n )\n + x\n )\n x = self.ls2(self.mlp(x)) + x\n return x\n\n\nclass AttentionLayer(nn.Module):\n def __init__(\n self,\n num_blocks: int,\n dim: int,\n num_heads: int = 4,\n expansion: int = 4,\n dropout: float = 0.0,\n cosine: bool = False,\n gated: bool = False,\n layer_scale: float = 1.0,\n context_dim: int | None = None,\n use_bias: bool = True,\n ):\n super().__init__()\n self.layers = nn.ModuleList(\n [\n AttentionBlock(\n dim=dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n cosine=cosine,\n gated=gated,\n layer_scale=layer_scale,\n context_dim=context_dim,\n use_bias=use_bias,\n )\n for _ in range(num_blocks)\n ]\n )\n\n def forward(\n self,\n x: torch.Tensor,\n context: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n pos_embed_context: torch.Tensor | None = None,\n attn_bias: torch.Tensor | None = None,\n ) -> torch.Tensor:\n for layer in self.layers:\n x = layer(\n x,\n context=context,\n pos_embed=pos_embed,\n pos_embed_context=pos_embed_context,\n attn_bias=attn_bias,\n )\n return x\n\n\nclass AttentionDecoderBlock(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int = 4,\n expansion: int = 4,\n dropout: float = 0.0,\n cosine: bool = False,\n gated: bool = False,\n layer_scale: float = 1.0,\n context_dim: int | None = None,\n single_head_ca: bool = True,\n ):\n super().__init__()\n self.dropout = dropout\n self.num_heads = num_heads\n self.hidden_dim = dim\n self.single_head_ca = single_head_ca\n context_dim = context_dim or dim\n self.mlp = MLP(dim, expansion=expansion, dropout=dropout, gated=gated)\n self.kv_ca = nn.Linear(context_dim, dim * 2)\n self.q_ca = nn.Linear(dim, dim)\n self.kv_sa = nn.Linear(dim, dim * 2)\n self.q_sa = nn.Linear(dim, dim)\n self.norm_x_sa = nn.LayerNorm(dim)\n self.norm_x_ca = nn.LayerNorm(dim)\n self.norm_ctx_ca = nn.LayerNorm(context_dim)\n self.cosine = cosine\n self.out_ca = nn.Linear(dim, dim)\n self.out_sa = nn.Linear(dim, dim)\n self.ls1 = LayerScale(dim, layer_scale) if layer_scale > 0.0 else nn.Identity()\n self.ls2 = LayerScale(dim, layer_scale) if layer_scale > 0.0 else nn.Identity()\n self.ls3 = LayerScale(dim, layer_scale) if layer_scale > 0.0 else nn.Identity()\n\n def cross_attn(\n self,\n x: torch.Tensor,\n attn_bias: torch.Tensor | None = None,\n context: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n pos_embed_context: torch.Tensor | None = None,\n rope: nn.Module | None = None,\n ) -> torch.Tensor:\n num_heads = 1 if self.single_head_ca else self.num_heads\n x = self.norm_x_ca(x)\n context = self.norm_ctx_ca(context)\n k, v = rearrange(\n self.kv_ca(context), \"b n (kv h d) -> b h n d kv\", h=num_heads, kv=2\n ).unbind(dim=-1)\n q = rearrange(self.q_ca(x), \"b n (h d) -> b h n d\", h=num_heads)\n\n if rope is not None:\n q = rope(q)\n k = rope(k)\n else:\n if pos_embed is not None:\n pos_embed = rearrange(pos_embed, \"b n (h d) -> b h n d\", h=num_heads)\n q = q + pos_embed\n if pos_embed_context is not None:\n pos_embed_context = rearrange(\n pos_embed_context, \"b n (h d) -> b h n d\", h=num_heads\n )\n k = k + pos_embed_context\n\n if self.cosine:\n q, k = map(partial(F.normalize, p=2, dim=-1), (q, k)) # cosine sim\n x = F.scaled_dot_product_attention(\n q, k, v, dropout_p=self.dropout, attn_mask=attn_bias\n )\n x = rearrange(x, \"b h n d -> b n (h d)\")\n x = self.out_ca(x)\n return x\n\n def self_attn(\n self,\n x: torch.Tensor,\n attn_bias: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n rope: nn.Module | None = None,\n ) -> torch.Tensor:\n x = self.norm_x_sa(x)\n k, v = rearrange(\n self.kv_sa(x), \"b n (kv h d) -> b h n d kv\", h=self.num_heads, kv=2\n ).unbind(dim=-1)\n q = rearrange(self.q_sa(x), \"b n (h d) -> b h n d\", h=self.num_heads)\n\n if rope is not None:\n q = rope(q)\n k = rope(k)\n elif pos_embed is not None:\n pos_embed = rearrange(pos_embed, \"b n (h d) -> b h n d\", h=self.num_heads)\n q = q + pos_embed\n\n if self.cosine:\n q, k = map(partial(F.normalize, p=2, dim=-1), (q, k)) # cosine sim\n x = F.scaled_dot_product_attention(\n q, k, v, dropout_p=self.dropout, attn_mask=attn_bias\n )\n x = rearrange(x, \"b h n d -> b n (h d)\")\n x = self.out_sa(x)\n return x\n\n def forward(\n self,\n x: torch.Tensor,\n attn_bias: torch.Tensor | None = None,\n context: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n pos_embed_context: torch.Tensor | None = None,\n rope: nn.Module | None = None,\n ) -> torch.Tensor:\n context = x if context is None else context\n x = (\n self.ls1(\n self.cross_attn(\n x,\n rope=rope,\n attn_bias=attn_bias,\n context=context,\n pos_embed=pos_embed,\n pos_embed_context=pos_embed_context,\n )\n )\n + x\n )\n x = (\n self.ls2(\n self.self_attn(x, rope=rope, attn_bias=attn_bias, pos_embed=pos_embed)\n )\n + x\n )\n x = self.ls3(self.mlp(x)) + x\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/convnext.py", "content": "import torch\nimport torch.nn as nn\n\n\nclass CvnxtBlock(nn.Module):\n def __init__(\n self,\n dim,\n kernel_size=7,\n layer_scale=1.0,\n expansion=4,\n dilation=1,\n padding_mode: str = \"zeros\",\n ):\n super().__init__()\n self.dwconv = nn.Conv2d(\n dim,\n dim,\n kernel_size=kernel_size,\n padding=dilation * (kernel_size - 1) // 2,\n groups=dim,\n dilation=dilation,\n padding_mode=padding_mode,\n ) # depthwise conv\n self.norm = nn.LayerNorm(dim)\n self.pwconv1 = nn.Linear(dim, expansion * dim)\n self.act = nn.GELU()\n self.pwconv2 = nn.Linear(expansion * dim, dim)\n self.gamma = (\n nn.Parameter(layer_scale * torch.ones((dim))) if layer_scale > 0.0 else 1.0\n )\n\n def forward(self, x):\n input = x\n x = self.dwconv(x)\n x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)\n x = self.norm(x)\n x = self.pwconv1(x)\n x = self.act(x)\n x = self.pwconv2(x)\n\n x = self.gamma * x\n x = input + x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/drop_path.py", "content": "import torch\nimport torch.nn as nn\n\n\ndef drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False):\n if drop_prob == 0.0 or not training:\n return x\n keep_prob = 1 - drop_prob\n shape = (x.shape[0],) + (1,) * (\n x.ndim - 1\n ) # work with diff dim tensors, not just 2D ConvNets\n random_tensor = x.new_empty(shape).bernoulli_(keep_prob)\n if keep_prob > 0.0:\n random_tensor.div_(keep_prob)\n output = x * random_tensor\n return output\n\n\nclass DropPath(nn.Module):\n def __init__(self, drop_prob=None):\n super(DropPath, self).__init__()\n self.drop_prob = drop_prob\n\n def forward(self, x):\n return drop_path(x, self.drop_prob, self.training)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/layer_scale.py", "content": "import torch\nimport torch.nn as nn\n\n\nclass LayerScale(nn.Module):\n def __init__(\n self,\n dim: int,\n init_values: float | torch.Tensor = 1e-5,\n inplace: bool = False,\n ) -> None:\n super().__init__()\n self.inplace = inplace\n self.gamma = nn.Parameter(init_values * torch.ones(dim))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return x.mul_(self.gamma) if self.inplace else x * self.gamma\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/mlp.py", "content": "import torch\nimport torch.nn as nn\n\nfrom unidepth.utils.misc import default\n\nfrom .activation import SwiGLU\n\n\nclass MLP(nn.Module):\n def __init__(\n self,\n input_dim: int,\n expansion: int = 4,\n dropout: float = 0.0,\n gated: bool = False,\n output_dim: int | None = None,\n ):\n super().__init__()\n if gated:\n expansion = int(expansion * 2 / 3)\n hidden_dim = int(input_dim * expansion)\n output_dim = default(output_dim, input_dim)\n self.norm = nn.LayerNorm(input_dim)\n self.proj1 = nn.Linear(input_dim, hidden_dim)\n self.proj2 = nn.Linear(hidden_dim, output_dim)\n self.act = nn.GELU() if not gated else SwiGLU()\n self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x = self.norm(x)\n x = self.proj1(x)\n x = self.act(x)\n x = self.proj2(x)\n x = self.dropout(x)\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/nystrom.py", "content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport logging\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\n\nfrom xformers.components.attention import Attention, AttentionConfig, register_attention\nfrom xformers.components.attention.core import (\n scaled_dot_product_attention,\n scaled_query_key_softmax,\n)\nfrom xformers.components.attention.utils import (\n bool_mask_to_additive,\n iterative_pinv,\n reshape_key_padding_mask,\n)\n\nlogger = logging.getLogger(\"xformers\")\n\n\n@dataclass\nclass NystromSelfAttentionConfig(AttentionConfig):\n \"\"\"\n num_heads Number of heads.\n num_landmarks Number of landmarks to use for softmax approximation. 64 often sufficient for a good\n approximation according to https://arxiv.org/pdf/2102.03902.pdf.\n causal Apply a causal mask, in that the attention cannot be applied to the future.\n use_razavi_pinverse If true, use iterative method from (Razavi et al. 2014) to approximate the Moore-Penrose\n inverse, otherwise use standard torch inverse.\n pinverse_original_init True if using original initialization when calculating Moore-Penrose pseudo inverse using\n method from (Razavi et al. 2014).\n False if using exact coefficient computation (leads to faster convergence).\n inv_iterations Number of iterations for calculating the Moore-Penrose pseudo inverse.\n v_skip_connection A module that will take V as input and will be added as a skip connection to the\n softmax approximation. A skip connection is added in the paper to help with training.\n conv_kernel_size Kernel size for convolution optionally added to help in training.\n If v_skip_connection is not specified, this will be used to define the default\n depth wise convolution used as a skip connection.\n If both conv_kernel_size and v_skip_connection are None, no skip connection will\n be added.\n landmark_pooling Which module to use when computing landmarks. Default is AdaptiveAvgPool2d.\n \"\"\"\n\n num_heads: int\n num_landmarks: Optional[int]\n landmark_pooling: Optional[nn.Module]\n causal: Optional[bool]\n pinverse_original_init: Optional[bool]\n inv_iterations: Optional[int]\n v_skip_connection: Optional[nn.Module]\n conv_kernel_size: Optional[int]\n use_razavi_pinverse: Optional[bool]\n\n\nclass AvgPool(nn.Module):\n def __init__(self, n: int):\n super().__init__()\n self.n = n\n\n def forward(self, x: torch.Tensor):\n # Average independently for every segment in the sequence dimension\n seq_len = x.shape[1]\n head_dim = x.shape[2]\n segments = seq_len // self.n\n assert segments > 0, \"num_landmarks should be smaller than the sequence length\"\n\n # Dimensions are a match\n if seq_len % self.n == 0:\n return x.reshape(\n -1,\n self.n,\n segments,\n head_dim,\n ).mean(dim=-2)\n\n # Handle the last segment boundary being off\n n_round = self.n - seq_len % self.n\n\n x_avg_round = (\n x[:, : n_round * segments, :]\n .reshape(-1, n_round, segments, head_dim)\n .mean(dim=-2)\n )\n x_avg_off = (\n x[:, n_round * segments :, :]\n .reshape(-1, self.n - n_round, segments + 1, head_dim)\n .mean(dim=-2)\n )\n return torch.cat((x_avg_round, x_avg_off), dim=-2)\n\n\n@register_attention(\"nystrom\", NystromSelfAttentionConfig)\nclass NystromAttention(Attention):\n # TODO: update defaults for use_razavi_pinverse and inv_iterations\n def __init__(\n self,\n dropout: float,\n num_heads: int,\n num_landmarks: int = 64,\n landmark_pooling: Optional[nn.Module] = None,\n causal: bool = False,\n use_razavi_pinverse: bool = True,\n pinverse_original_init: bool = False,\n inv_iterations: int = 6, # recommended default in paper was 6.\n v_skip_connection: Optional[nn.Module] = None,\n conv_kernel_size: Optional[int] = None,\n *args,\n **kwargs,\n ):\n \"\"\"\n Nystrom attention mechanism, from Nystromformer_.\n ::\n\n \"A Nystrom-based Algorithm for Approximating Self-Attention.\"\n Xiong, Y., Zeng, Z., Chakraborty, R., Tan, M., Fung, G., Li, Y., Singh, V. (2021)\n\n Reference codebase: https://github.com/mlpen/Nystromformer\n\n .. _Nystromformer: https://arxiv.org/pdf/2102.03902.pdf\n\n \"\"\"\n super().__init__()\n # merged key padding mask and attention mask is not accepted\n self.requires_separate_masks = True\n self.num_landmarks = num_landmarks\n # TODO: should be able to not have to pass in num_heads\n self.num_heads = num_heads\n self.use_razavi_pinverse = use_razavi_pinverse\n self.pinverse_original_init = pinverse_original_init\n self.inv_iterations = inv_iterations\n self.attn_drop = nn.Dropout(dropout)\n self.skip_connection = v_skip_connection\n self.causal = causal\n\n if self.skip_connection is None and conv_kernel_size is not None:\n self.skip_connection = nn.Conv2d(\n in_channels=self.num_heads,\n out_channels=self.num_heads,\n kernel_size=(conv_kernel_size, 1),\n padding=(conv_kernel_size // 2, 0),\n bias=False,\n groups=self.num_heads,\n )\n\n if landmark_pooling is not None:\n self.landmark_pooling = landmark_pooling\n else:\n self.landmark_pooling = AvgPool(n=self.num_landmarks)\n\n # Optional lower triangular masks for causal attention\n self.causal_mask_1: Optional[torch.Tensor] = None\n self.causal_mask_2: Optional[torch.Tensor] = None\n self.causal_mask_3: Optional[torch.Tensor] = None\n\n # This attention does not support attention masks\n self.supports_attention_mask = False\n self.supports_key_padding_mask = True\n\n def forward(\n self,\n q: torch.Tensor,\n k: torch.Tensor,\n v: torch.Tensor,\n key_padding_mask: Optional[torch.Tensor] = None,\n *args,\n **kwargs,\n ):\n r\"\"\"\n key_padding_mask Only a key padding mask is accepted here. The size must be (batch size, sequence length) or\n (batch size * num_heads, 1, sequence length). If dimensions are not correct, the mask will\n be ignored. An additive mask is expected, meaning float values using \"-inf\" to mask values\n \"\"\"\n\n batched_dim = k.size(0)\n seq_len = k.size(-2)\n tt = {\"dtype\": q.dtype, \"device\": q.device}\n\n if key_padding_mask is not None:\n if key_padding_mask.dtype == torch.bool:\n logger.warning(\n \"Bool mask found, but an additive mask is expected. Converting but this is slow\"\n )\n\n key_padding_mask = bool_mask_to_additive(key_padding_mask)\n\n if key_padding_mask.ndim == 2:\n key_padding_mask = reshape_key_padding_mask(\n key_padding_mask, batched_dim\n )\n\n zeros = torch.zeros_like(key_padding_mask)\n ones = torch.ones_like(key_padding_mask)\n is_masked = torch.isinf(-key_padding_mask)\n\n # _mask takes 1 if the token is not padded, otherwise 0.\n _mask = torch.where(is_masked, zeros, ones)\n _mask = _mask.transpose(2, 1)\n assert _mask.shape == (batched_dim, q.shape[1], 1)\n\n # Mask q and k before pooling\n # https://github.com/mlpen/Nystromformer/blob/main/code/attention_nystrom.py#L31\n q = q * _mask\n k = k * _mask\n\n assert key_padding_mask.size() == (batched_dim, 1, seq_len), (\n f\"key_padding_mask has invalid dimensions {key_padding_mask.size()}.\"\n f\" Must have dimensions {batched_dim, 1, seq_len} or (batch_size, {seq_len}).\"\n )\n\n if self.num_landmarks >= seq_len:\n mask: Optional[torch.Tensor] = None\n\n if self.causal:\n mask = self._triu_mask(batched_dim, seq_len, seq_len, **tt)\n\n if key_padding_mask is not None:\n mask = key_padding_mask if mask is None else mask + key_padding_mask\n\n x = scaled_dot_product_attention(q=q, k=k, v=v, att_mask=mask)\n\n else:\n q_landmarks = self.landmark_pooling(q)\n k_landmarks = self.landmark_pooling(k)\n\n if self.causal and (\n self.causal_mask_1 is None\n or (batched_dim, seq_len, self.num_landmarks)\n != self.causal_mask_1.size()\n ):\n self.causal_mask_1 = self._triu_mask(\n batched_dim, seq_len, self.num_landmarks, **tt\n )\n self.causal_mask_2 = self._triu_mask(\n batched_dim, self.num_landmarks, self.num_landmarks, **tt\n )\n self.causal_mask_3 = self._triu_mask(\n batched_dim, self.num_landmarks, seq_len, **tt\n )\n\n mask_3: Optional[torch.Tensor] = self.causal_mask_3\n if key_padding_mask is not None:\n mask_3 = (\n key_padding_mask if mask_3 is None else mask_3 + key_padding_mask\n )\n\n kernel_1 = scaled_query_key_softmax(q=q, k=k_landmarks, att_mask=None)\n kernel_2 = scaled_query_key_softmax(\n q=q_landmarks, k=k_landmarks, att_mask=None\n )\n kernel_3 = scaled_dot_product_attention(\n q=q_landmarks, k=k, v=v, att_mask=mask_3\n )\n\n kernel_2_inv = (\n iterative_pinv(\n kernel_2, self.inv_iterations, self.pinverse_original_init\n )\n if self.use_razavi_pinverse\n else torch.linalg.pinv(kernel_2)\n )\n\n x = torch.matmul(\n torch.matmul(\n kernel_1,\n kernel_2_inv,\n ),\n kernel_3,\n )\n\n if self.skip_connection:\n # Assumption here is that v is 3D.\n v_conv = self.skip_connection(\n v.reshape(-1, self.num_heads, v.size(-2), v.size(-1))\n )\n x += v_conv.reshape(-1, v_conv.size(-2), v_conv.size(-1))\n x = self.attn_drop(x)\n return x\n\n def _triu_mask(self, dim_1: int, dim_2: int, dim_3: int, **kwargs) -> torch.Tensor:\n device = kwargs[\"device\"]\n dtype = kwargs[\"dtype\"]\n\n return torch.triu(\n torch.ones(dim_2, dim_3, dtype=dtype, device=device) * float(\"-inf\"),\n diagonal=1,\n ).expand(\n dim_1, -1, -1\n ) # micro optim, save memory on the batch dimension\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/nystrom_attention.py", "content": "from functools import partial\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\n\nfrom .nystrom import NystromAttention\nfrom .attention import AttentionBlock\n\n\nclass NystromBlock(AttentionBlock):\n def __init__(\n self,\n dim: int,\n num_heads: int = 4,\n expansion: int = 4,\n dropout: float = 0.0,\n cosine: bool = False,\n gated: bool = False,\n layer_scale: float = 1.0,\n context_dim: int | None = None,\n ):\n super().__init__(\n dim=dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n cosine=cosine,\n gated=gated,\n layer_scale=layer_scale,\n context_dim=context_dim,\n )\n self.attention_fn = NystromAttention(\n num_landmarks=128, num_heads=num_heads, dropout=dropout\n )\n\n def attn(\n self,\n x: torch.Tensor,\n attn_bias: torch.Tensor | None = None,\n context: torch.Tensor | None = None,\n pos_embed: torch.Tensor | None = None,\n pos_embed_context: torch.Tensor | None = None,\n rope: nn.Module | None = None,\n ) -> torch.Tensor:\n x = self.norm_attnx(x)\n context = self.norm_attnctx(context)\n k, v = rearrange(\n self.kv(context), \"b n (kv h d) -> b n h d kv\", h=self.num_heads, kv=2\n ).unbind(dim=-1)\n q = rearrange(self.q(x), \"b n (h d) -> b n h d\", h=self.num_heads)\n\n if rope is not None:\n q = rope(q)\n k = rope(k)\n else:\n if pos_embed is not None:\n pos_embed = rearrange(\n pos_embed, \"b n (h d) -> b n h d\", h=self.num_heads\n )\n q = q + pos_embed\n if pos_embed_context is not None:\n pos_embed_context = rearrange(\n pos_embed_context, \"b n (h d) -> b n h d\", h=self.num_heads\n )\n k = k + pos_embed_context\n\n if self.cosine:\n q, k = map(partial(F.normalize, p=2, dim=-1), (q, k)) # cosine sim\n x = self.attention_fn(q, k, v, key_padding_mask=attn_bias)\n x = rearrange(x, \"b n h d -> b n (h d)\")\n x = self.out(x)\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/positional_encoding.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom math import pi\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nfrom einops import rearrange, repeat\n\n\nclass PositionEmbeddingSine(nn.Module):\n def __init__(\n self, num_pos_feats=64, temperature=10000, normalize=False, scale=None\n ):\n super().__init__()\n self.num_pos_feats = num_pos_feats\n self.temperature = temperature\n self.normalize = normalize\n if scale is not None and normalize is False:\n raise ValueError(\"normalize should be True if scale is passed\")\n if scale is None:\n scale = 2 * pi\n self.scale = scale\n\n def forward(\n self, x: torch.Tensor, mask: Optional[torch.Tensor] = None\n ) -> torch.Tensor:\n if mask is None:\n mask = torch.zeros(\n (x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool\n )\n not_mask = ~mask\n y_embed = not_mask.cumsum(1, dtype=torch.float32)\n x_embed = not_mask.cumsum(2, dtype=torch.float32)\n if self.normalize:\n eps = 1e-6\n y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale\n x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale\n\n dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)\n dim_t = self.temperature ** (\n 2 * torch.div(dim_t, 2, rounding_mode=\"floor\") / self.num_pos_feats\n )\n\n pos_x = x_embed[:, :, :, None] / dim_t\n pos_y = y_embed[:, :, :, None] / dim_t\n pos_x = torch.stack(\n (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4\n ).flatten(3)\n pos_y = torch.stack(\n (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4\n ).flatten(3)\n pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)\n return pos\n\n def __repr__(self, _repr_indent=4):\n head = \"Positional encoding \" + self.__class__.__name__\n body = [\n \"num_pos_feats: {}\".format(self.num_pos_feats),\n \"temperature: {}\".format(self.temperature),\n \"normalize: {}\".format(self.normalize),\n \"scale: {}\".format(self.scale),\n ]\n # _repr_indent = 4\n lines = [head] + [\" \" * _repr_indent + line for line in body]\n return \"\\n\".join(lines)\n\n\nclass LearnedSinusoidalPosEmb(nn.Module):\n def __init__(self, dim):\n super().__init__()\n assert (dim % 2) == 0\n half_dim = dim // 2\n self.weights = nn.Parameter(torch.randn(half_dim))\n\n def forward(self, x):\n x = rearrange(x, \"b -> b 1\")\n freqs = x * rearrange(self.weights, \"d -> 1 d\") * 2 * pi\n fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)\n fouriered = torch.cat((x, fouriered), dim=-1)\n return fouriered\n\n\ndef generate_fourier_features(x, max_freq=64, num_bands=16):\n x = x.unsqueeze(-1)\n device, dtype, orig_x = x.device, x.dtype, x\n\n scales = torch.linspace(\n -max_freq / 2, max_freq / 2, num_bands, device=device, dtype=dtype\n )\n scales = scales[(*((None,) * (len(x.shape) - 1)), Ellipsis)]\n\n x = x * scales * pi\n x = torch.cat([x.sin(), x.cos()], dim=-1)\n x = torch.cat((x, orig_x), dim=-1)\n return x.flatten(-2)\n\n\ndef broadcat(tensors, dim=-1):\n num_tensors = len(tensors)\n shape_lens = set(list(map(lambda t: len(t.shape), tensors)))\n assert len(shape_lens) == 1, \"tensors must all have the same number of dimensions\"\n shape_len = list(shape_lens)[0]\n dim = (dim + shape_len) if dim < 0 else dim\n dims = list(zip(*map(lambda t: list(t.shape), tensors)))\n expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]\n assert all(\n [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]\n ), \"invalid dimensions for broadcastable concatentation\"\n max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))\n expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))\n expanded_dims.insert(dim, (dim, dims[dim]))\n expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))\n tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))\n return torch.cat(tensors, dim=dim)\n\n\ndef rotate_half(x):\n x = rearrange(x, \"... (d r) -> ... d r\", r=2)\n x1, x2 = x.unbind(dim=-1)\n x = torch.stack((-x2, x1), dim=-1)\n return rearrange(x, \"... d r -> ... (d r)\")\n\n\nclass VisionRotaryEmbedding(nn.Module):\n def __init__(\n self,\n dim,\n pt_seq_len,\n ft_seq_len=None,\n custom_freqs=None,\n freqs_for=\"lang\",\n theta=10000,\n max_freq=10,\n num_freqs=1,\n ):\n super().__init__()\n if custom_freqs:\n freqs = custom_freqs\n elif freqs_for == \"lang\":\n freqs = 1.0 / (\n theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)\n )\n elif freqs_for == \"pixel\":\n freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi\n elif freqs_for == \"constant\":\n freqs = torch.ones(num_freqs).float()\n else:\n raise ValueError(f\"unknown modality {freqs_for}\")\n\n if ft_seq_len is None:\n ft_seq_len = pt_seq_len\n t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len\n\n freqs_h = torch.einsum(\"..., f -> ... f\", t, freqs)\n freqs_h = repeat(freqs_h, \"... n -> ... (n r)\", r=2)\n\n freqs_w = torch.einsum(\"..., f -> ... f\", t, freqs)\n freqs_w = repeat(freqs_w, \"... n -> ... (n r)\", r=2)\n\n freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)\n\n self.register_buffer(\"freqs_cos\", freqs.cos())\n self.register_buffer(\"freqs_sin\", freqs.sin())\n\n print(\"======== shape of rope freq\", self.freqs_cos.shape, \"========\")\n\n def forward(self, t, start_index=0):\n rot_dim = self.freqs_cos.shape[-1]\n end_index = start_index + rot_dim\n assert (\n rot_dim <= t.shape[-1]\n ), f\"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}\"\n t_left, t, t_right = (\n t[..., :start_index],\n t[..., start_index:end_index],\n t[..., end_index:],\n )\n t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)\n return torch.cat((t_left, t, t_right), dim=-1)\n\n\nclass VisionRotaryEmbeddingFast(nn.Module):\n def __init__(\n self,\n dim,\n pt_seq_len,\n ft_seq_len=None,\n custom_freqs=None,\n freqs_for=\"lang\",\n theta=10000,\n max_freq=10,\n num_freqs=1,\n ):\n super().__init__()\n if custom_freqs:\n freqs = custom_freqs\n elif freqs_for == \"lang\":\n freqs = 1.0 / (\n theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)\n )\n elif freqs_for == \"pixel\":\n freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi\n elif freqs_for == \"constant\":\n freqs = torch.ones(num_freqs).float()\n else:\n raise ValueError(f\"unknown modality {freqs_for}\")\n\n if ft_seq_len is None:\n ft_seq_len = pt_seq_len\n t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len\n\n freqs = torch.einsum(\"..., f -> ... f\", t, freqs)\n freqs = repeat(freqs, \"... n -> ... (n r)\", r=2)\n freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)\n\n freqs_cos = freqs.cos().view(-1, freqs.shape[-1])\n freqs_sin = freqs.sin().view(-1, freqs.shape[-1])\n\n self.register_buffer(\"freqs_cos\", freqs_cos)\n self.register_buffer(\"freqs_sin\", freqs_sin)\n\n def forward(self, t):\n return t * self.freqs_cos + rotate_half(t) * self.freqs_sin\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/layers/upsample.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nfrom einops import rearrange\n\nfrom .convnext import CvnxtBlock\n\n\nclass ConvUpsample(nn.Module):\n def __init__(\n self,\n hidden_dim,\n num_layers: int = 2,\n expansion: int = 4,\n layer_scale: float = 1.0,\n kernel_size: int = 7,\n **kwargs,\n ):\n super().__init__()\n self.convs = nn.ModuleList([])\n for _ in range(num_layers):\n self.convs.append(\n CvnxtBlock(\n hidden_dim,\n kernel_size=kernel_size,\n expansion=expansion,\n layer_scale=layer_scale,\n )\n )\n self.up = nn.Sequential(\n nn.Conv2d(hidden_dim, hidden_dim // 2, kernel_size=1, padding=0),\n nn.UpsamplingBilinear2d(scale_factor=2),\n nn.Conv2d(hidden_dim // 2, hidden_dim // 2, kernel_size=3, padding=1),\n )\n\n def forward(self, x: torch.Tensor):\n for conv in self.convs:\n x = conv(x)\n x = self.up(x)\n x = rearrange(x, \"b c h w -> b (h w) c\")\n return x\n\n\nclass ConvUpsampleShuffle(nn.Module):\n def __init__(\n self,\n hidden_dim,\n num_layers: int = 2,\n expansion: int = 4,\n layer_scale: float = 1.0,\n kernel_size: int = 7,\n **kwargs,\n ):\n super().__init__()\n self.convs = nn.ModuleList([])\n for _ in range(num_layers):\n self.convs.append(\n CvnxtBlock(\n hidden_dim,\n kernel_size=kernel_size,\n expansion=expansion,\n layer_scale=layer_scale,\n )\n )\n self.up = nn.Sequential(\n nn.PixelShuffle(2),\n nn.Conv2d(hidden_dim // 4, hidden_dim // 2, kernel_size=3, padding=1),\n )\n\n def forward(self, x: torch.Tensor):\n for conv in self.convs:\n x = conv(x)\n x = self.up(x)\n x = rearrange(x, \"b c h w -> b (h w) c\")\n return x\n\n\nclass ConvUpsampleShuffleResidual(nn.Module):\n def __init__(\n self,\n hidden_dim,\n num_layers: int = 2,\n expansion: int = 4,\n layer_scale: float = 1.0,\n kernel_size: int = 7,\n padding_mode: str = \"zeros\",\n **kwargs,\n ):\n super().__init__()\n self.convs = nn.ModuleList([])\n for _ in range(num_layers):\n self.convs.append(\n CvnxtBlock(\n hidden_dim,\n kernel_size=kernel_size,\n expansion=expansion,\n layer_scale=layer_scale,\n padding_mode=padding_mode,\n )\n )\n self.up = nn.Sequential(\n nn.PixelShuffle(2),\n nn.Conv2d(\n hidden_dim // 4,\n hidden_dim // 4,\n kernel_size=7,\n padding=3,\n padding_mode=padding_mode,\n groups=hidden_dim // 4,\n ),\n nn.ReLU(),\n nn.Conv2d(\n hidden_dim // 4,\n hidden_dim // 2,\n kernel_size=3,\n padding=1,\n padding_mode=padding_mode,\n ),\n )\n self.residual = nn.Sequential(\n nn.Conv2d(hidden_dim, hidden_dim // 2, kernel_size=1, padding=0),\n nn.UpsamplingBilinear2d(scale_factor=2),\n )\n\n def forward(self, x: torch.Tensor):\n for conv in self.convs:\n x = conv(x)\n x = self.up(x) + self.residual(x)\n x = rearrange(x, \"b c h w -> b (h w) c\")\n return x\n\n\nclass ResidualConvUnit(nn.Module):\n def __init__(\n self,\n dim,\n kernel_size: int = 3,\n padding_mode: str = \"zeros\",\n dilation: int = 1,\n layer_scale: float = 1.0,\n use_norm: bool = False,\n ):\n super().__init__()\n self.conv1 = nn.Conv2d(\n dim,\n dim,\n kernel_size=kernel_size,\n padding=dilation * (kernel_size - 1) // 2,\n dilation=dilation,\n padding_mode=padding_mode,\n )\n self.conv2 = nn.Conv2d(\n dim,\n dim,\n kernel_size=kernel_size,\n padding=dilation * (kernel_size - 1) // 2,\n dilation=dilation,\n padding_mode=padding_mode,\n )\n self.activation = nn.LeakyReLU()\n self.gamma = (\n nn.Parameter(layer_scale * torch.ones(1, dim, 1, 1))\n if layer_scale > 0.0\n else 1.0\n )\n self.norm1 = nn.GroupNorm(dim // 16, dim) if use_norm else nn.Identity()\n self.norm2 = nn.GroupNorm(dim // 16, dim) if use_norm else nn.Identity()\n\n def forward(self, x):\n out = self.activation(x)\n out = self.conv1(out)\n out = self.norm1(out)\n out = self.activation(out)\n out = self.conv2(out)\n out = self.norm2(out)\n return self.gamma * out + x\n\n\nclass ResUpsampleBil(nn.Module):\n def __init__(\n self,\n hidden_dim,\n output_dim: int = None,\n num_layers: int = 2,\n kernel_size: int = 3,\n layer_scale: float = 1.0,\n padding_mode: str = \"zeros\",\n use_norm: bool = False,\n **kwargs,\n ):\n super().__init__()\n output_dim = output_dim if output_dim is not None else hidden_dim // 2\n self.convs = nn.ModuleList([])\n for _ in range(num_layers):\n self.convs.append(\n ResidualConvUnit(\n hidden_dim,\n kernel_size=kernel_size,\n layer_scale=layer_scale,\n padding_mode=padding_mode,\n use_norm=use_norm,\n )\n )\n self.up = nn.Sequential(\n nn.Conv2d(\n hidden_dim,\n output_dim,\n kernel_size=1,\n padding=0,\n padding_mode=padding_mode,\n ),\n nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n )\n\n def forward(self, x: torch.Tensor):\n for conv in self.convs:\n x = conv(x)\n x = self.up(x)\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/__init__.py", "content": "from .unidepthv1 import UniDepthV1\nfrom .unidepthv2 import UniDepthV2, UniDepthV2old\n\n__all__ = [\n \"UniDepthV1\",\n \"UniDepthV2old\",\n \"UniDepthV2\",\n]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/__init__.py", "content": "from .convnext import ConvNeXt\nfrom .convnext2 import ConvNeXtV2\nfrom .dinov2 import _make_dinov2_model\n\n__all__ = [\n \"ConvNeXt\",\n \"ConvNeXtV2\",\n \"_make_dinov2_model\",\n]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/convnext.py", "content": "from collections import OrderedDict\nfrom functools import partial\nfrom typing import Callable, Optional, Sequence, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nfrom timm.layers import (AvgPool2dSame, DropPath, GlobalResponseNormMlp,\n LayerNorm, LayerNorm2d, Mlp, create_conv2d,\n get_act_layer, make_divisible, to_ntuple,\n trunc_normal_)\nfrom torch.utils.checkpoint import checkpoint\n\n\ndef get_num_layer_for_convnext(var_name):\n \"\"\"\n Divide [3, 3, 27, 3] layers into 12 groups; each group is three\n consecutive blocks, including possible neighboring downsample layers;\n adapted from https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py\n \"\"\"\n if var_name.startswith(\"downsample_layers\"):\n stage_id = int(var_name.split(\".\")[1])\n if stage_id == 0:\n layer_id = 0\n elif stage_id == 1 or stage_id == 2:\n layer_id = stage_id + 1\n elif stage_id == 3:\n layer_id = 12\n\n elif var_name.startswith(\"stages\"):\n stage_id = int(var_name.split(\".\")[1])\n block_id = int(var_name.split(\".\")[3])\n if stage_id == 0 or stage_id == 1:\n layer_id = stage_id + 1\n elif stage_id == 2:\n layer_id = 3 + block_id // 3\n elif stage_id == 3:\n layer_id = 12\n\n elif var_name.startswith(\"stem\"):\n return 0\n else:\n layer_id = 12\n return layer_id + 1\n\n\ndef get_parameter_groups(model, lr, wd=1e-5, ld=0.9, skip_list=None):\n parameter_group_names = {}\n parameter_group_vars = {}\n skip = set()\n if skip_list is not None:\n skip = skip_list\n if hasattr(model, \"no_weight_decay\"):\n skip.update(model.no_weight_decay())\n num_layers = 12\n layer_scale = list(ld ** (num_layers + 1 - i) for i in range(num_layers + 2))\n for name, param in model.named_parameters():\n if not param.requires_grad:\n continue # frozen weights\n if len(param.shape) == 1 or name.endswith(\".bias\") or name in skip:\n group_name = \"no_decay\"\n this_wd = 0.0\n else:\n group_name = \"decay\"\n this_wd = wd\n\n layer_id = get_num_layer_for_convnext(name)\n group_name = \"layer_%d_%s\" % (layer_id, group_name)\n\n if group_name not in parameter_group_names:\n scale = layer_scale[layer_id]\n cur_lr = lr * scale\n\n parameter_group_names[group_name] = {\n \"weight_decay\": this_wd,\n \"weight_decay_init\": this_wd,\n \"weight_decay_base\": this_wd,\n \"params\": [],\n \"lr_init\": cur_lr,\n \"lr_base\": lr,\n \"lr\": cur_lr,\n }\n parameter_group_vars[group_name] = {\n \"weight_decay\": this_wd,\n \"weight_decay_init\": this_wd,\n \"weight_decay_base\": this_wd,\n \"params\": [],\n \"lr_init\": cur_lr,\n \"lr_base\": lr,\n \"lr\": cur_lr,\n }\n if this_wd == 0.0:\n parameter_group_names[group_name][\"weight_decay_final\"] = 0.0\n parameter_group_vars[group_name][\"weight_decay_final\"] = 0.0\n parameter_group_vars[group_name][\"params\"].append(param)\n parameter_group_names[group_name][\"params\"].append(name)\n # from unidepth.utils import is_main_process\n # import json\n # if is_main_process():\n # print(\"Param groups = %s\" % json.dumps(parameter_group_names, indent=2))\n return list(parameter_group_vars.values()), [\n v[\"lr\"] for k, v in parameter_group_vars.items()\n ]\n\n\nclass Downsample(nn.Module):\n def __init__(self, in_chs, out_chs, stride=1, dilation=1):\n super().__init__()\n avg_stride = stride if dilation == 1 else 1\n if stride > 1 or dilation > 1:\n avg_pool_fn = (\n AvgPool2dSame if avg_stride == 1 and dilation > 1 else nn.AvgPool2d\n )\n self.pool = avg_pool_fn(\n 2, avg_stride, ceil_mode=True, count_include_pad=False\n )\n else:\n self.pool = nn.Identity()\n\n if in_chs != out_chs:\n self.conv = create_conv2d(in_chs, out_chs, 1, stride=1)\n else:\n self.conv = nn.Identity()\n\n def forward(self, x):\n x = self.pool(x)\n x = self.conv(x)\n return x\n\n\nclass ConvNeXtBlock(nn.Module):\n \"\"\"ConvNeXt Block\n There are two equivalent implementations:\n (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)\n (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back\n\n Unlike the official impl, this one allows choice of 1 or 2, 1x1 conv can be faster with appropriate\n choice of LayerNorm impl, however as model size increases the tradeoffs appear to change and nn.Linear\n is a better choice. This was observed with PyTorch 1.10 on 3090 GPU, it could change over time & w/ different HW.\n \"\"\"\n\n def __init__(\n self,\n in_chs: int,\n out_chs: Optional[int] = None,\n kernel_size: int = 7,\n stride: int = 1,\n dilation: Union[int, Tuple[int, int]] = (1, 1),\n mlp_ratio: float = 4,\n conv_mlp: bool = False,\n conv_bias: bool = True,\n use_grn: bool = False,\n ls_init_value: Optional[float] = 1e-6,\n act_layer: Union[str, Callable] = \"gelu\",\n norm_layer: Optional[Callable] = None,\n drop_path: float = 0.0,\n ):\n \"\"\"\n\n Args:\n in_chs: Block input channels.\n out_chs: Block output channels (same as in_chs if None).\n kernel_size: Depthwise convolution kernel size.\n stride: Stride of depthwise convolution.\n dilation: Tuple specifying input and output dilation of block.\n mlp_ratio: MLP expansion ratio.\n conv_mlp: Use 1x1 convolutions for MLP and a NCHW compatible norm layer if True.\n conv_bias: Apply bias for all convolution (linear) layers.\n use_grn: Use GlobalResponseNorm in MLP (from ConvNeXt-V2)\n ls_init_value: Layer-scale init values, layer-scale applied if not None.\n act_layer: Activation layer.\n norm_layer: Normalization layer (defaults to LN if not specified).\n drop_path: Stochastic depth probability.\n \"\"\"\n super().__init__()\n out_chs = out_chs or in_chs\n dilation = to_ntuple(2)(dilation)\n act_layer = get_act_layer(act_layer)\n if not norm_layer:\n norm_layer = LayerNorm2d if conv_mlp else LayerNorm\n mlp_layer = partial(\n GlobalResponseNormMlp if use_grn else Mlp, use_conv=conv_mlp\n )\n self.use_conv_mlp = conv_mlp\n self.conv_dw = create_conv2d(\n in_chs,\n out_chs,\n kernel_size=kernel_size,\n stride=stride,\n dilation=dilation[0],\n depthwise=True,\n bias=conv_bias,\n )\n self.norm = norm_layer(out_chs)\n self.mlp = mlp_layer(out_chs, int(mlp_ratio * out_chs), act_layer=act_layer)\n self.gamma = (\n nn.Parameter(ls_init_value * torch.ones(out_chs))\n if ls_init_value is not None\n else None\n )\n if in_chs != out_chs or stride != 1 or dilation[0] != dilation[1]:\n self.shortcut = Downsample(\n in_chs, out_chs, stride=stride, dilation=dilation[0]\n )\n else:\n self.shortcut = nn.Identity()\n self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n def forward(self, x):\n shortcut = x\n x = self.conv_dw(x.contiguous())\n if self.use_conv_mlp:\n x = self.norm(x)\n x = self.mlp(x)\n else:\n x = x.permute(0, 2, 3, 1).contiguous()\n x = self.norm(x)\n x = self.mlp(x)\n x = x.permute(0, 3, 1, 2).contiguous()\n if self.gamma is not None:\n x = x.mul(self.gamma.reshape(1, -1, 1, 1))\n\n x = self.drop_path(x) + self.shortcut(shortcut)\n return x.contiguous()\n\n\nclass ConvNeXtStage(nn.Module):\n def __init__(\n self,\n in_chs,\n out_chs,\n kernel_size=7,\n stride=2,\n depth=2,\n dilation=(1, 1),\n drop_path_rates=None,\n ls_init_value=1.0,\n conv_mlp=False,\n conv_bias=True,\n use_grn=False,\n act_layer=\"gelu\",\n norm_layer=None,\n norm_layer_cl=None,\n ):\n super().__init__()\n self.grad_checkpointing = False\n\n if in_chs != out_chs or stride > 1 or dilation[0] != dilation[1]:\n ds_ks = 2 if stride > 1 or dilation[0] != dilation[1] else 1\n pad = (\n \"same\" if dilation[1] > 1 else 0\n ) # same padding needed if dilation used\n self.downsample = nn.Sequential(\n norm_layer(in_chs),\n create_conv2d(\n in_chs,\n out_chs,\n kernel_size=ds_ks,\n stride=stride,\n dilation=dilation[0],\n padding=pad,\n bias=conv_bias,\n ),\n )\n in_chs = out_chs\n else:\n self.downsample = nn.Identity()\n\n drop_path_rates = drop_path_rates or [0.0] * depth\n stage_blocks = []\n for i in range(depth):\n stage_blocks.append(\n ConvNeXtBlock(\n in_chs=in_chs,\n out_chs=out_chs,\n kernel_size=kernel_size,\n dilation=dilation[1],\n drop_path=drop_path_rates[i],\n ls_init_value=ls_init_value,\n conv_mlp=conv_mlp,\n conv_bias=conv_bias,\n use_grn=use_grn,\n act_layer=act_layer,\n norm_layer=norm_layer if conv_mlp else norm_layer_cl,\n )\n )\n in_chs = out_chs\n self.blocks = nn.ModuleList(stage_blocks)\n\n def forward(self, x):\n xs = []\n x = self.downsample(x)\n for block in self.blocks:\n if self.grad_checkpointing:\n x = checkpoint(block, x)\n else:\n x = block(x)\n xs.append(x)\n return xs\n\n\nclass ConvNeXt(nn.Module):\n def __init__(\n self,\n in_chans: int = 3,\n output_stride: int = 32,\n depths: Tuple[int, ...] = (3, 3, 9, 3),\n dims: Tuple[int, ...] = (96, 192, 384, 768),\n kernel_sizes: Union[int, Tuple[int, ...]] = 7,\n ls_init_value: Optional[float] = 1e-6,\n stem_type: str = \"patch\",\n patch_size: int = 4,\n conv_mlp: bool = False,\n conv_bias: bool = True,\n use_grn: bool = False,\n act_layer: Union[str, Callable] = \"gelu\",\n norm_layer: Optional[Union[str, Callable]] = None,\n norm_eps: Optional[float] = None,\n drop_path_rate: float = 0.0,\n output_idx=[],\n use_checkpoint=False,\n ):\n \"\"\"\n Args:\n in_chans: Number of input image channels.\n num_classes: Number of classes for classification head.\n global_pool: Global pooling type.\n output_stride: Output stride of network, one of (8, 16, 32).\n depths: Number of blocks at each stage.\n dims: Feature dimension at each stage.\n kernel_sizes: Depthwise convolution kernel-sizes for each stage.\n ls_init_value: Init value for Layer Scale, disabled if None.\n stem_type: Type of stem.\n patch_size: Stem patch size for patch stem.\n head_init_scale: Init scaling value for classifier weights and biases.\n head_norm_first: Apply normalization before global pool + head.\n head_hidden_size: Size of MLP hidden layer in head if not None and head_norm_first == False.\n conv_mlp: Use 1x1 conv in MLP, improves speed for small networks w/ chan last.\n conv_bias: Use bias layers w/ all convolutions.\n use_grn: Use Global Response Norm (ConvNeXt-V2) in MLP.\n act_layer: Activation layer type.\n norm_layer: Normalization layer type.\n drop_rate: Head pre-classifier dropout rate.\n drop_path_rate: Stochastic depth drop rate.\n \"\"\"\n super().__init__()\n self.num_layers = len(depths)\n self.depths = output_idx\n self.embed_dims = [\n int(dim) for i, dim in enumerate(dims) for _ in range(depths[i])\n ]\n self.embed_dim = dims[0]\n\n assert output_stride in (8, 16, 32)\n kernel_sizes = to_ntuple(4)(kernel_sizes)\n if norm_layer is None:\n norm_layer = LayerNorm2d\n norm_layer_cl = norm_layer if conv_mlp else LayerNorm\n if norm_eps is not None:\n norm_layer = partial(norm_layer, eps=norm_eps)\n norm_layer_cl = partial(norm_layer_cl, eps=norm_eps)\n else:\n assert (\n conv_mlp\n ), \"If a norm_layer is specified, conv MLP must be used so all norm expect rank-4, channels-first input\"\n norm_layer_cl = norm_layer\n if norm_eps is not None:\n norm_layer_cl = partial(norm_layer_cl, eps=norm_eps)\n\n self.feature_info = []\n\n assert stem_type in (\"patch\", \"overlap\", \"overlap_tiered\")\n if stem_type == \"patch\":\n # NOTE: this stem is a minimal form of ViT PatchEmbed, as used in SwinTransformer w/ patch_size = 4\n self.stem = nn.Sequential(\n nn.Conv2d(\n in_chans,\n dims[0],\n kernel_size=patch_size,\n stride=patch_size,\n bias=conv_bias,\n ),\n norm_layer(dims[0]),\n )\n stem_stride = patch_size\n else:\n mid_chs = make_divisible(dims[0] // 2) if \"tiered\" in stem_type else dims[0]\n self.stem = nn.Sequential(\n nn.Conv2d(\n in_chans,\n mid_chs,\n kernel_size=3,\n stride=2,\n padding=1,\n bias=conv_bias,\n ),\n nn.Conv2d(\n mid_chs, dims[0], kernel_size=3, stride=2, padding=1, bias=conv_bias\n ),\n norm_layer(dims[0]),\n )\n stem_stride = 4\n\n self.stages = nn.Sequential()\n dp_rates = [\n x.tolist()\n for x in torch.linspace(0, drop_path_rate, sum(depths)).split(depths)\n ]\n stages = []\n prev_chs = dims[0]\n curr_stride = stem_stride\n dilation = 1\n # 4 feature resolution stages, each consisting of multiple residual blocks\n for i in range(4):\n stride = 2 if curr_stride == 2 or i > 0 else 1\n if curr_stride >= output_stride and stride > 1:\n dilation *= stride\n stride = 1\n curr_stride *= stride\n first_dilation = 1 if dilation in (1, 2) else 2\n out_chs = dims[i]\n stages.append(\n ConvNeXtStage(\n prev_chs,\n out_chs,\n kernel_size=kernel_sizes[i],\n stride=stride,\n dilation=(first_dilation, dilation),\n depth=depths[i],\n drop_path_rates=dp_rates[i],\n ls_init_value=ls_init_value,\n conv_mlp=conv_mlp,\n conv_bias=conv_bias,\n use_grn=use_grn,\n act_layer=act_layer,\n norm_layer=norm_layer,\n norm_layer_cl=norm_layer_cl,\n )\n )\n prev_chs = out_chs\n # NOTE feature_info use currently assumes stage 0 == stride 1, rest are stride 2\n self.feature_info += [\n dict(num_chs=prev_chs, reduction=curr_stride, module=f\"stages.{i}\")\n ]\n self.stages = nn.ModuleList(stages)\n self.mask_token = nn.Parameter(torch.zeros(1, self.embed_dim, 1, 1))\n self.num_features = prev_chs\n self.apply(self._init_weights)\n self.set_grad_checkpointing(use_checkpoint)\n\n def _init_weights(self, module):\n if isinstance(module, nn.Conv2d):\n trunc_normal_(module.weight, std=0.02)\n if module.bias is not None:\n nn.init.zeros_(module.bias)\n elif isinstance(module, nn.Linear):\n trunc_normal_(module.weight, std=0.02)\n nn.init.zeros_(module.bias)\n\n def forward(self, x, masks=None):\n outs = []\n x = self.stem(x)\n if masks is not None:\n masks = torch.nn.functional.interpolate(\n masks.float(), size=x.shape[-2:], mode=\"nearest\"\n )\n x = torch.where(masks.bool(), self.mask_token.to(x.dtype), x).contiguous()\n for stage in self.stages:\n xs = stage(x)\n outs.extend([x.permute(0, 2, 3, 1).contiguous() for x in xs])\n x = xs[-1]\n return outs, [x.mean(dim=(1, 2)).unsqueeze(1).contiguous() for x in outs]\n\n @torch.jit.ignore\n def group_matcher(self, coarse=False):\n return dict(\n stem=r\"^stem\",\n blocks=(\n r\"^stages\\.(\\d+)\"\n if coarse\n else [\n (r\"^stages\\.(\\d+)\\.downsample\", (0,)), # blocks\n (r\"^stages\\.(\\d+)\\.blocks\\.(\\d+)\", None),\n (r\"^norm_pre\", (99999,)),\n ]\n ),\n )\n\n @torch.jit.ignore\n def set_grad_checkpointing(self, enable=True):\n for s in self.stages:\n s.grad_checkpointing = enable\n\n def freeze(self) -> None:\n for module in self.modules():\n module.eval()\n for parameters in self.parameters():\n parameters.requires_grad = False\n\n def get_params(self, lr, wd, ld, *args, **kwargs):\n encoder_p, encoder_lr = get_parameter_groups(self, lr, wd, ld)\n return encoder_p, encoder_lr\n\n def no_weight_decay(self):\n return {\"mask_token\"}\n\n @classmethod\n def build(cls, config):\n obj = globals()[config[\"model\"][\"encoder\"][\"name\"]](config)\n return obj\n\n\ndef checkpoint_filter_fn(state_dict, model):\n \"\"\"Remap FB checkpoints -> timm\"\"\"\n if \"head.norm.weight\" in state_dict or \"norm_pre.weight\" in state_dict:\n return state_dict # non-FB checkpoint\n if \"model\" in state_dict:\n state_dict = state_dict[\"model\"]\n\n out_dict = {}\n if \"visual.trunk.stem.0.weight\" in state_dict:\n out_dict = {\n k.replace(\"visual.trunk.\", \"\"): v\n for k, v in state_dict.items()\n if k.startswith(\"visual.trunk.\")\n }\n if \"visual.head.proj.weight\" in state_dict:\n out_dict[\"head.fc.weight\"] = state_dict[\"visual.head.proj.weight\"]\n out_dict[\"head.fc.bias\"] = torch.zeros(\n state_dict[\"visual.head.proj.weight\"].shape[0]\n )\n elif \"visual.head.mlp.fc1.weight\" in state_dict:\n out_dict[\"head.pre_logits.fc.weight\"] = state_dict[\n \"visual.head.mlp.fc1.weight\"\n ]\n out_dict[\"head.pre_logits.fc.bias\"] = state_dict[\"visual.head.mlp.fc1.bias\"]\n out_dict[\"head.fc.weight\"] = state_dict[\"visual.head.mlp.fc2.weight\"]\n out_dict[\"head.fc.bias\"] = torch.zeros(\n state_dict[\"visual.head.mlp.fc2.weight\"].shape[0]\n )\n return out_dict\n\n import re\n\n for k, v in state_dict.items():\n k = k.replace(\"downsample_layers.0.\", \"stem.\")\n k = re.sub(r\"stages.([0-9]+).([0-9]+)\", r\"stages.\\1.blocks.\\2\", k)\n k = re.sub(\n r\"downsample_layers.([0-9]+).([0-9]+)\", r\"stages.\\1.downsample.\\2\", k\n )\n k = k.replace(\"dwconv\", \"conv_dw\")\n k = k.replace(\"pwconv\", \"mlp.fc\")\n if \"grn\" in k:\n k = k.replace(\"grn.beta\", \"mlp.grn.bias\")\n k = k.replace(\"grn.gamma\", \"mlp.grn.weight\")\n v = v.reshape(v.shape[-1])\n k = k.replace(\"head.\", \"head.fc.\")\n if k.startswith(\"norm.\"):\n k = k.replace(\"norm\", \"head.norm\")\n if v.ndim == 2 and \"head\" not in k:\n model_shape = model.state_dict()[k].shape\n v = v.reshape(model_shape)\n out_dict[k] = v\n\n return out_dict\n\n\nHF_URL = {\n \"convnext_xxlarge_pt\": (\n \"laion/CLIP-convnext_xxlarge-laion2B-s34B-b82K-augreg-soup\",\n \"open_clip_pytorch_model.bin\",\n ),\n \"convnext_large_pt\": (\n \"laion/CLIP-convnext_large_d_320.laion2B-s29B-b131K-ft-soup\",\n \"open_clip_pytorch_model.bin\",\n ),\n \"convnext_large\": (\n \"timm/convnext_large_mlp.clip_laion2b_soup_ft_in12k_in1k_384\",\n \"pytorch_model.bin\",\n ),\n}\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/convnext2.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom timm.models.layers import DropPath, trunc_normal_\n\n\ndef get_num_layer_for_convnext_single(var_name, depths):\n \"\"\"\n Each layer is assigned distinctive layer ids\n \"\"\"\n if var_name.startswith(\"downsample_layers\"):\n stage_id = int(var_name.split(\".\")[1])\n layer_id = sum(depths[:stage_id]) + 1\n return layer_id\n\n elif var_name.startswith(\"stages\"):\n stage_id = int(var_name.split(\".\")[1])\n block_id = int(var_name.split(\".\")[2])\n layer_id = sum(depths[:stage_id]) + block_id + 1\n return layer_id\n\n else:\n return sum(depths) + 1\n\n\ndef get_num_layer_for_convnext(var_name):\n \"\"\"\n Divide [3, 3, 27, 3] layers into 12 groups; each group is three\n consecutive blocks, including possible neighboring downsample layers;\n adapted from https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py\n \"\"\"\n num_max_layer = 12\n if var_name.startswith(\"downsample_layers\"):\n stage_id = int(var_name.split(\".\")[1])\n if stage_id == 0:\n layer_id = 0\n elif stage_id == 1 or stage_id == 2:\n layer_id = stage_id + 1\n elif stage_id == 3:\n layer_id = 12\n return layer_id\n\n elif var_name.startswith(\"stages\"):\n stage_id = int(var_name.split(\".\")[1])\n block_id = int(var_name.split(\".\")[2])\n if stage_id == 0 or stage_id == 1:\n layer_id = stage_id + 1\n elif stage_id == 2:\n layer_id = 3 + block_id // 3\n elif stage_id == 3:\n layer_id = 12\n return layer_id\n else:\n return num_max_layer + 1\n\n\ndef get_parameter_groups(model, lr, wd=1e-5, ld=0.9, skip_list=()):\n parameter_group_names = {}\n parameter_group_vars = {}\n skip = {}\n if skip_list is not None:\n skip = skip_list\n elif hasattr(model, \"no_weight_decay\"):\n skip = model.no_weight_decay()\n num_layers = 12 # sum(model.depths)\n layer_scale = list(ld ** (num_layers + 1 - i) for i in range(num_layers + 2))\n for name, param in model.named_parameters():\n if not param.requires_grad:\n continue # frozen weights\n if (\n len(param.shape) == 1\n or name.endswith(\".bias\")\n or name in skip\n or name.endswith(\".gamma\")\n or name.endswith(\".beta\")\n ):\n group_name = \"no_decay\"\n this_weight_decay = 0.0\n else:\n group_name = \"decay\"\n this_weight_decay = wd\n\n # layer_id = get_num_layer_for_convnext_single(name, model.depths)\n layer_id = get_num_layer_for_convnext(name)\n group_name = \"layer_%d_%s\" % (layer_id, group_name)\n\n if group_name not in parameter_group_names:\n scale = layer_scale[layer_id]\n cur_lr = lr * scale\n\n parameter_group_names[group_name] = {\n \"weight_decay\": this_weight_decay,\n \"params\": [],\n \"lr_scale\": scale,\n \"lr\": cur_lr,\n }\n parameter_group_vars[group_name] = {\n \"weight_decay\": this_weight_decay,\n \"params\": [],\n \"lr_scale\": scale,\n \"lr\": cur_lr,\n }\n parameter_group_vars[group_name][\"params\"].append(param)\n parameter_group_names[group_name][\"params\"].append(name)\n # if is_main_process():\n # print(\"Param groups = %s\" % json.dumps(parameter_group_names, indent=2))\n return list(parameter_group_vars.values()), [\n v[\"lr\"] for k, v in parameter_group_vars.items()\n ]\n\n\nclass LayerNorm(nn.Module):\n \"\"\"LayerNorm that supports two data formats: channels_last (default) or channels_first.\n The ordering of the dimensions in the inputs. channels_last corresponds to inputs with\n shape (batch_size, height, width, channels) while channels_first corresponds to inputs\n with shape (batch_size, channels, height, width).\n \"\"\"\n\n def __init__(self, normalized_shape, eps=1e-6, data_format=\"channels_last\"):\n super().__init__()\n self.weight = nn.Parameter(torch.ones(normalized_shape))\n self.bias = nn.Parameter(torch.zeros(normalized_shape))\n self.eps = eps\n self.data_format = data_format\n if self.data_format not in [\"channels_last\", \"channels_first\"]:\n raise NotImplementedError\n self.normalized_shape = (normalized_shape,)\n\n def forward(self, x):\n if self.data_format == \"channels_last\":\n return F.layer_norm(\n x, self.normalized_shape, self.weight, self.bias, self.eps\n )\n elif self.data_format == \"channels_first\":\n u = x.mean(1, keepdim=True)\n s = (x - u).pow(2).mean(1, keepdim=True)\n x = (x - u) / torch.sqrt(s + self.eps)\n x = self.weight[:, None, None] * x + self.bias[:, None, None]\n return x\n\n\nclass GRN(nn.Module):\n \"\"\"GRN (Global Response Normalization) layer\"\"\"\n\n def __init__(self, dim):\n super().__init__()\n self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))\n self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))\n\n def forward(self, x):\n Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)\n Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)\n return self.gamma * (x * Nx) + self.beta + x\n\n\nclass Block(nn.Module):\n \"\"\"ConvNeXtV2 Block.\n\n Args:\n dim (int): Number of input channels.\n drop_path (float): Stochastic depth rate. Default: 0.0\n \"\"\"\n\n def __init__(self, dim, drop_path=0.0, mult=4, use_checkpoint=False):\n super().__init__()\n self.dwconv = nn.Conv2d(\n dim, dim, kernel_size=7, padding=3, groups=dim\n ) # depthwise conv\n self.norm = LayerNorm(dim, eps=1e-6)\n self.pwconv1 = nn.Linear(\n dim, mult * dim\n ) # pointwise/1x1 convs, implemented with linear layers\n self.act = nn.GELU()\n self.grn = GRN(mult * dim)\n self.pwconv2 = nn.Linear(mult * dim, dim)\n self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n self.use_checkpoint = use_checkpoint\n\n def forward(self, x):\n input = x\n x = self.dwconv(x)\n x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)\n x = self.norm(x)\n x = self.pwconv1(x)\n x = self.act(x)\n x = self.grn(x)\n x = self.pwconv2(x)\n x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)\n\n x = input + self.drop_path(x)\n return x\n\n\nclass ConvNeXtV2(nn.Module):\n \"\"\"ConvNeXt V2\n\n Args:\n in_chans (int): Number of input image channels. Default: 3\n num_classes (int): Number of classes for classification head. Default: 1000\n depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]\n dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]\n drop_path_rate (float): Stochastic depth rate. Default: 0.\n head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.\n \"\"\"\n\n def __init__(\n self,\n in_chans=3,\n depths=[3, 3, 9, 3],\n dims=96,\n drop_path_rate=0.0,\n output_idx=[],\n use_checkpoint=False,\n ):\n super().__init__()\n self.num_layers = len(depths)\n self.depths = output_idx\n self.embed_dims = [\n int(dim) for i, dim in enumerate(dims) for _ in range(depths[i])\n ]\n self.embed_dim = dims[0]\n\n self.downsample_layers = (\n nn.ModuleList()\n ) # stem and 3 intermediate downsampling conv layers\n stem = nn.Sequential(\n nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),\n LayerNorm(dims[0], eps=1e-6, data_format=\"channels_first\"),\n )\n self.downsample_layers.append(stem)\n for i in range(3):\n downsample_layer = nn.Sequential(\n LayerNorm(dims[i], eps=1e-6, data_format=\"channels_first\"),\n nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),\n )\n self.downsample_layers.append(downsample_layer)\n\n self.stages = (\n nn.ModuleList()\n ) # 4 feature resolution stages, each consisting of multiple residual blocks\n self.out_norms = nn.ModuleList()\n dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]\n cur = 0\n for i in range(4):\n stage = nn.ModuleList(\n [\n Block(\n dim=dims[i],\n drop_path=dp_rates[cur + j],\n use_checkpoint=use_checkpoint,\n )\n for j in range(depths[i])\n ]\n )\n self.stages.append(stage)\n cur += depths[i]\n\n self.apply(self._init_weights)\n\n def _init_weights(self, m):\n if isinstance(m, (nn.Conv2d, nn.Linear)):\n trunc_normal_(m.weight, std=0.02)\n nn.init.constant_(m.bias, 0)\n\n def forward(self, x):\n outs = []\n for i in range(4):\n x = self.downsample_layers[i](x)\n for stage in self.stages[i]:\n x = stage(x)\n outs.append(x.permute(0, 2, 3, 1))\n cls_tokens = [x.mean(dim=(1, 2)).unsqueeze(1).contiguous() for x in outs]\n return outs, cls_tokens\n\n def get_params(self, lr, wd, ld, *args, **kwargs):\n encoder_p, encoder_lr = get_parameter_groups(self, lr, wd, ld)\n return encoder_p, encoder_lr\n\n def freeze(self) -> None:\n for module in self.modules():\n module.eval()\n for parameters in self.parameters():\n parameters.requires_grad = False\n\n @classmethod\n def build(cls, config):\n obj = globals()[config[\"model\"][\"encoder\"][\"name\"]](config)\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/dinov2.py", "content": "import contextlib\nimport logging\nimport math\nfrom functools import partial\nfrom typing import Callable, Sequence\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn.init import trunc_normal_\n\nfrom .metadinov2 import Attention, MemEffAttention, Mlp\nfrom .metadinov2 import NestedTensorBlock as Block\nfrom .metadinov2 import PatchEmbed, SwiGLUFFNFused\n\n_DINOV2_BASE_URL = \"https://dl.fbaipublicfiles.com/dinov2\"\nlogger = logging.getLogger(\"dinov2\")\n\n\ndef named_apply(\n fn: Callable, module: nn.Module, name=\"\", depth_first=True, include_root=False\n) -> nn.Module:\n if not depth_first and include_root:\n fn(module=module, name=name)\n for child_name, child_module in module.named_children():\n child_name = \".\".join((name, child_name)) if name else child_name\n named_apply(\n fn=fn,\n module=child_module,\n name=child_name,\n depth_first=depth_first,\n include_root=True,\n )\n if depth_first and include_root:\n fn(module=module, name=name)\n return module\n\n\ndef get_parameter_groups(model, lr, wd=1e-5, ld=0.9, skip_list=()):\n parameter_group_names = {}\n parameter_group_vars = {}\n skip = {}\n if skip_list is not None:\n skip = skip_list\n elif hasattr(model, \"no_weight_decay\"):\n skip = model.no_weight_decay()\n\n num_layers = model.n_blocks\n layer_scale = list(ld ** (num_layers - i) for i in range(num_layers))\n\n for name, param in model.named_parameters():\n if not param.requires_grad:\n continue\n\n if len(param.shape) == 1: # norm\n group_name = \"no_decay\"\n this_wd = 0.0\n # layer scale, bias beta?\n elif (\n name in skip\n or name.endswith(\".gamma\")\n or name.endswith(\".beta\")\n or name.endswith(\".bias\")\n ):\n group_name = \"no_decay\"\n this_wd = 0.0\n elif \"cls_token\" in name or \"pos_embed\" in name or \"mask_token\" in name:\n group_name = \"no_decay\"\n this_wd = 0.0\n else:\n group_name = \"decay\"\n this_wd = wd\n\n if name.startswith(\"blocks\"):\n layer_id = int(name.split(\".\")[1])\n elif name.startswith(\"patch_embed\"):\n layer_id = 0\n else:\n layer_id = 0\n\n group_name = f\"layer_{layer_id}_{group_name}\"\n\n if group_name not in parameter_group_names:\n scale = layer_scale[layer_id]\n cur_lr = lr * scale\n\n parameter_group_names[group_name] = {\n \"weight_decay\": this_wd,\n \"params\": [],\n \"lr_init\": cur_lr,\n \"lr_base\": lr,\n \"lr\": cur_lr,\n }\n parameter_group_vars[group_name] = {\n \"weight_decay\": this_wd,\n \"params\": [],\n \"lr_init\": cur_lr,\n \"lr_base\": lr,\n \"lr\": cur_lr,\n }\n parameter_group_vars[group_name][\"params\"].append(param)\n parameter_group_names[group_name][\"params\"].append(name)\n\n return list(parameter_group_vars.values()), [\n v[\"lr\"] for k, v in parameter_group_vars.items()\n ]\n\n\nclass BlockChunk(nn.ModuleList):\n def forward(self, x):\n for b in self:\n x = b(x)\n return x\n\n\nclass DinoVisionTransformer(nn.Module):\n def __init__(\n self,\n img_size=224,\n patch_size=16,\n in_chans=3,\n embed_dim=768,\n depth=12,\n num_heads=12,\n mlp_ratio=4.0,\n qkv_bias=True,\n ffn_bias=True,\n proj_bias=True,\n drop_path_rate=0.0,\n drop_path_uniform=False,\n init_values=None, # for layerscale: None or 0 => no layerscale\n embed_layer=PatchEmbed,\n act_layer=nn.GELU,\n block_fn=Block,\n ffn_layer=\"mlp\",\n block_chunks=1,\n output_idx=[5, 12, 18, 24],\n checkpoint: bool = False,\n num_register_tokens=0,\n interpolate_antialias=False,\n interpolate_offset=0.0,\n use_norm=False,\n frozen_stages=0,\n ):\n \"\"\"\n Args:\n img_size (int, tuple): input image size\n patch_size (int, tuple): patch size\n in_chans (int): number of input channels\n embed_dim (int): embedding dimension\n depth (int): depth of transformer\n num_heads (int): number of attention heads\n mlp_ratio (int): ratio of mlp hidden dim to embedding dim\n qkv_bias (bool): enable bias for qkv if True\n proj_bias (bool): enable bias for proj in attn if True\n ffn_bias (bool): enable bias for ffn if True\n drop_path_rate (float): stochastic depth rate\n drop_path_uniform (bool): apply uniform drop rate across blocks\n weight_init (str): weight init scheme\n init_values (float): layer-scale init values\n embed_layer (nn.Module): patch embedding layer\n act_layer (nn.Module): MLP activation layer\n block_fn (nn.Module): transformer block class\n ffn_layer (str): \"mlp\", \"swiglu\", \"swiglufused\" or \"identity\"\n block_chunks: (int) split block sequence into block_chunks units for FSDP wrap\n \"\"\"\n super().__init__()\n norm_layer = partial(nn.LayerNorm, eps=1e-6)\n\n self.num_features = self.embed_dim = (\n embed_dim # num_features for consistency with other models\n )\n self.frozen_stages = frozen_stages\n self.embed_dims = [embed_dim] * output_idx[-1]\n self.num_tokens = 1\n self.n_blocks = depth\n self.num_heads = num_heads\n self.patch_size = patch_size\n self.depths = output_idx\n self.checkpoint = checkpoint\n self.num_register_tokens = num_register_tokens\n self.interpolate_antialias = interpolate_antialias\n self.interpolate_offset = interpolate_offset\n\n self.patch_embed = embed_layer(\n img_size=img_size,\n patch_size=patch_size,\n in_chans=in_chans,\n embed_dim=embed_dim,\n )\n num_patches = self.patch_embed.num_patches\n\n self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n self.pos_embed = nn.Parameter(\n torch.zeros(1, num_patches + self.num_tokens, embed_dim)\n )\n assert num_register_tokens >= 0\n self.register_tokens = nn.Parameter(\n torch.zeros(1, max(1, num_register_tokens), embed_dim)\n )\n\n if drop_path_uniform is True:\n dpr = [drop_path_rate] * depth\n else:\n dpr = [\n x.item() for x in torch.linspace(0, drop_path_rate, depth)\n ] # stochastic depth decay rule\n\n if ffn_layer == \"mlp\":\n logger.info(\"using MLP layer as FFN\")\n ffn_layer = Mlp\n elif ffn_layer == \"swiglufused\" or ffn_layer == \"swiglu\":\n logger.info(\"using SwiGLU layer as FFN\")\n ffn_layer = SwiGLUFFNFused\n elif ffn_layer == \"identity\":\n logger.info(\"using Identity layer as FFN\")\n\n def f(*args, **kwargs):\n return nn.Identity()\n\n ffn_layer = f\n else:\n raise NotImplementedError\n\n blocks_list = [\n block_fn(\n dim=embed_dim,\n num_heads=num_heads,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias,\n proj_bias=proj_bias,\n ffn_bias=ffn_bias,\n drop_path=dpr[i],\n norm_layer=norm_layer,\n act_layer=act_layer,\n ffn_layer=ffn_layer,\n init_values=init_values,\n )\n for i in range(depth)\n ]\n if block_chunks > 0:\n self.chunked_blocks = True\n chunked_blocks = []\n chunksize = depth // block_chunks\n for i in range(0, depth, chunksize):\n # this is to keep the block index consistent if we chunk the block list\n chunked_blocks.append(\n [nn.Identity()] * i + blocks_list[i : i + chunksize]\n )\n self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])\n else:\n self.chunked_blocks = False\n self.blocks = nn.ModuleList(blocks_list)\n\n self.norm = nn.LayerNorm(embed_dim)\n self.use_norm = use_norm\n self.head = nn.Identity()\n self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))\n self.init_weights()\n\n def init_weights(self):\n trunc_normal_(self.pos_embed, std=0.02)\n nn.init.normal_(self.cls_token, std=1e-6)\n if self.num_register_tokens:\n nn.init.normal_(self.register_tokens, std=1e-6)\n named_apply(init_weights_vit_timm, self)\n\n def interpolate_pos_encoding(self, x, w, h):\n previous_dtype = x.dtype\n npatch = x.shape[1] - 1\n N = self.pos_embed.shape[1] - 1\n if npatch == N and w == h:\n return self.pos_embed\n pos_embed = self.pos_embed.float()\n class_pos_embed = pos_embed[:, 0]\n patch_pos_embed = pos_embed[:, 1:]\n dim = x.shape[-1]\n w0 = w // self.patch_size\n h0 = h // self.patch_size\n\n M = int(math.sqrt(N)) # Recover the number of patches in each dimension\n assert N == M * M\n kwargs = {}\n if self.interpolate_offset:\n # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8\n # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors\n sx = float(w0 + self.interpolate_offset) / M\n sy = float(h0 + self.interpolate_offset) / M\n kwargs[\"scale_factor\"] = (sx, sy)\n else:\n # Simply specify an output size instead of a scale factor\n kwargs[\"size\"] = (w0, h0)\n\n patch_pos_embed = nn.functional.interpolate(\n patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2),\n mode=\"bicubic\",\n antialias=self.interpolate_antialias,\n **kwargs,\n )\n assert (w0, h0) == patch_pos_embed.shape[-2:]\n\n patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)\n return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(\n previous_dtype\n )\n\n def prepare_tokens_with_masks(self, x, masks=None):\n B, nc, w, h = x.shape\n with torch.no_grad() if self.frozen_stages > -1 else contextlib.nullcontext():\n x = self.patch_embed(x)\n if masks is not None:\n masks = masks.bool().view(B, -1, 1)\n x = torch.where(masks, self.mask_token.to(x.dtype).unsqueeze(0), x)\n\n x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)\n x = x + self.interpolate_pos_encoding(x, w, h)\n\n if self.num_register_tokens:\n x = torch.cat(\n (x[:, :1], self.register_tokens.expand(x.shape[0], -1, -1), x[:, 1:]),\n dim=1,\n )\n return x\n\n def forward(self, x, masks=None):\n shapes = [val // self.patch_size for val in x.shape[-2:]]\n batch_size = x.shape[0]\n x = self.prepare_tokens_with_masks(x, masks)\n outputs = []\n for i, blk in enumerate(self.blocks):\n with (\n torch.no_grad() if i < self.frozen_stages else contextlib.nullcontext()\n ):\n x = blk(x)\n outputs.append(x)\n\n if self.use_norm:\n with (\n torch.no_grad()\n if self.frozen_stages >= len(self.blocks)\n else contextlib.nullcontext()\n ):\n outputs = [self.norm(out) for out in outputs]\n class_tokens = [out[:, :1] for out in outputs]\n outputs = [out[:, self.num_register_tokens + 1 :] for out in outputs]\n outputs = [out.reshape(batch_size, *shapes, -1) for out in outputs]\n\n return (outputs, class_tokens)\n\n def get_params(self, lr, wd, ld, *args, **kwargs):\n encoder_p, encoder_lr = get_parameter_groups(self, lr, wd, ld)\n return encoder_p, encoder_lr\n\n def freeze(self) -> None:\n for module in self.modules():\n module.eval()\n for parameters in self.parameters():\n parameters.requires_grad = False\n\n def train(self, mode=True):\n super().train(mode)\n if self.frozen_stages > -1:\n for p in self.patch_embed.parameters():\n p.requires_grad = False\n\n for i, blk in enumerate(self.blocks):\n if i < self.frozen_stages:\n blk.eval()\n for p in blk.parameters():\n p.requires_grad = False\n\n for p in self.norm.parameters():\n p.requires_grad = self.frozen_stages <= len(self.blocks) and self.use_norm\n\n self.cls_token.requires_grad = self.frozen_stages < 1\n self.pos_embed.requires_grad = self.frozen_stages < 1\n self.mask_token.requires_grad = False\n self.register_tokens.requires_grad = False\n\n\ndef init_weights_vit_timm(module: nn.Module, name: str = \"\"):\n \"\"\"ViT weight initialization, original timm impl (for reproducibility)\"\"\"\n if isinstance(module, nn.Linear):\n trunc_normal_(module.weight, std=0.02)\n if module.bias is not None:\n nn.init.zeros_(module.bias)\n\n\ndef vit_small(patch_size=16, num_register_tokens=0, export=False, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=384,\n depth=12,\n num_heads=6,\n mlp_ratio=4,\n num_register_tokens=num_register_tokens,\n block_fn=partial(Block, attn_class=Attention if export else MemEffAttention),\n **kwargs,\n )\n return model\n\n\ndef vit_base(patch_size=16, num_register_tokens=0, export=False, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=768,\n depth=12,\n num_heads=12,\n mlp_ratio=4,\n num_register_tokens=num_register_tokens,\n block_fn=partial(Block, attn_class=Attention if export else MemEffAttention),\n **kwargs,\n )\n return model\n\n\ndef vit_large(patch_size=16, num_register_tokens=0, export=False, **kwargs):\n model = DinoVisionTransformer(\n patch_size=patch_size,\n embed_dim=1024,\n depth=24,\n num_heads=16,\n mlp_ratio=4,\n num_register_tokens=num_register_tokens,\n block_fn=partial(Block, attn_class=Attention if export else MemEffAttention),\n **kwargs,\n )\n return model\n\n\ndef _make_dinov2_model_name(arch_name: str, patch_size: int) -> str:\n compact_arch_name = arch_name.replace(\"_\", \"\")[:4]\n return f\"dinov2_{compact_arch_name}{patch_size}\"\n\n\ndef _make_dinov2_model(\n *,\n arch_name: str = \"vit_large\",\n img_size: int = 518,\n patch_size: int = 14,\n init_values: float = 1.0,\n ffn_layer: str = \"mlp\",\n block_chunks: int = 0,\n pretrained: str = \"\",\n output_idx: Sequence[int] = [],\n num_register_tokens: int = 0,\n drop_path_rate: float = 0.0,\n use_norm: bool = False,\n export: bool = False,\n interpolate_offset: float = 0.0,\n frozen_stages: int = 0,\n **kwargs,\n):\n model_name = _make_dinov2_model_name(arch_name, patch_size)\n vit_kwargs = dict(\n img_size=img_size,\n patch_size=patch_size,\n init_values=init_values,\n ffn_layer=ffn_layer,\n block_chunks=block_chunks,\n output_idx=output_idx,\n drop_path_rate=drop_path_rate,\n num_register_tokens=num_register_tokens,\n use_norm=use_norm,\n export=export,\n interpolate_offset=interpolate_offset,\n frozen_stages=frozen_stages,\n )\n vit_kwargs.update(**kwargs)\n model = eval(arch_name)(**vit_kwargs)\n if pretrained == \"\":\n url = _DINOV2_BASE_URL + f\"/{model_name}/{model_name}\"\n if num_register_tokens > 0:\n url += \"_reg4\"\n url += \"_pretrain.pth\"\n state_dict = torch.hub.load_state_dict_from_url(\n url, map_location=\"cpu\", progress=False\n )\n info = model.load_state_dict(state_dict, strict=False)\n print(info)\n elif pretrained is not None:\n state_dict = torch.load(pretrained, map_location=\"cpu\")\n info = model.load_state_dict(state_dict, strict=False)\n print(f\"loading from {pretrained} with:\", info)\n else:\n print(\"Not loading pretrained weights for backbone\")\n return model\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom .attention import Attention, MemEffAttention\nfrom .block import NestedTensorBlock\nfrom .dino_head import DINOHead\nfrom .mlp import Mlp\nfrom .patch_embed import PatchEmbed\nfrom .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/attention.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py\n\nimport logging\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nlogger = logging.getLogger(\"dinov2\")\n\n\ntry:\n from xformers.ops import fmha, memory_efficient_attention, unbind\n\n XFORMERS_AVAILABLE = True\nexcept ImportError:\n logger.warning(\"xFormers not available\")\n XFORMERS_AVAILABLE = False\n\nXFORMERS_AVAILABLE = XFORMERS_AVAILABLE and torch.cuda.is_available()\n\n\nclass Attention(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int = 8,\n qkv_bias: bool = False,\n proj_bias: bool = True,\n attn_drop: float = 0.0,\n proj_drop: float = 0.0,\n ) -> None:\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.scale = head_dim**-0.5\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(dim, dim, bias=proj_bias)\n self.proj_drop = nn.Dropout(proj_drop)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, N, C = x.shape\n qkv = (\n self.qkv(x)\n .reshape(B, N, 3, self.num_heads, C // self.num_heads)\n .permute(2, 0, 3, 1, 4)\n )\n x = F.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2])\n x = x.transpose(1, 2).reshape(B, N, C)\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n\n\nclass MemEffAttention(Attention):\n def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:\n # new pytorch have good attn efficient, no need for xformers\n if not XFORMERS_AVAILABLE or x.device.type == \"cpu\":\n assert attn_bias is None, \"xFormers is required for nested tensors usage\"\n return super().forward(x)\n\n B, N, C = x.shape\n qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)\n\n q, k, v = unbind(qkv, 2)\n\n x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)\n x = x.reshape([B, N, C])\n\n x = self.proj(x)\n x = self.proj_drop(x)\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/block.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py\n\nimport logging\nfrom typing import Any, Callable, Dict, List, Tuple\n\nimport torch\nimport torch.nn as nn\n\nfrom .attention import Attention, MemEffAttention\nfrom .drop_path import DropPath\nfrom .layer_scale import LayerScale\nfrom .mlp import Mlp\n\nlogger = logging.getLogger(\"dinov2\")\n\n\ntry:\n from xformers.ops import fmha, index_select_cat, scaled_index_add\n\n XFORMERS_AVAILABLE = True\nexcept ImportError:\n logger.warning(\"xFormers not available\")\n XFORMERS_AVAILABLE = False\n\n\nclass Block(nn.Module):\n def __init__(\n self,\n dim: int,\n num_heads: int,\n mlp_ratio: float = 4.0,\n qkv_bias: bool = False,\n proj_bias: bool = True,\n ffn_bias: bool = True,\n drop: float = 0.0,\n attn_drop: float = 0.0,\n init_values=None,\n drop_path: float = 0.0,\n act_layer: Callable[..., nn.Module] = nn.GELU,\n norm_layer: Callable[..., nn.Module] = nn.LayerNorm,\n attn_class: Callable[..., nn.Module] = Attention,\n ffn_layer: Callable[..., nn.Module] = Mlp,\n ) -> None:\n super().__init__()\n # print(f\"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}\")\n self.norm1 = norm_layer(dim)\n self.attn = attn_class(\n dim,\n num_heads=num_heads,\n qkv_bias=qkv_bias,\n proj_bias=proj_bias,\n attn_drop=attn_drop,\n proj_drop=drop,\n )\n self.ls1 = (\n LayerScale(dim, init_values=init_values) if init_values else nn.Identity()\n )\n self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = ffn_layer(\n in_features=dim,\n hidden_features=mlp_hidden_dim,\n act_layer=act_layer,\n drop=drop,\n bias=ffn_bias,\n )\n self.ls2 = (\n LayerScale(dim, init_values=init_values) if init_values else nn.Identity()\n )\n self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n self.sample_drop_ratio = drop_path\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n def attn_residual_func(x: torch.Tensor) -> torch.Tensor:\n return self.ls1(self.attn(self.norm1(x)))\n\n def ffn_residual_func(x: torch.Tensor) -> torch.Tensor:\n return self.ls2(self.mlp(self.norm2(x)))\n\n if self.training and self.sample_drop_ratio > 0.1:\n # the overhead is compensated only for a drop path rate larger than 0.1\n x = drop_add_residual_stochastic_depth(\n x,\n residual_func=attn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n )\n x = drop_add_residual_stochastic_depth(\n x,\n residual_func=ffn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n )\n elif self.training and self.sample_drop_ratio > 0.0:\n x = x + self.drop_path1(attn_residual_func(x))\n x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2\n else:\n x = x + attn_residual_func(x)\n x = x + ffn_residual_func(x)\n return x\n\n\ndef drop_add_residual_stochastic_depth(\n x: torch.Tensor,\n residual_func: Callable[[torch.Tensor], torch.Tensor],\n sample_drop_ratio: float = 0.0,\n) -> torch.Tensor:\n # 1) extract subset using permutation\n b, n, d = x.shape\n sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)\n brange = (torch.randperm(b, device=x.device))[:sample_subset_size]\n x_subset = x[brange]\n\n # 2) apply residual_func to get residual\n residual = residual_func(x_subset)\n\n x_flat = x.flatten(1)\n residual = residual.flatten(1)\n\n residual_scale_factor = b / sample_subset_size\n\n # 3) add the residual\n x_plus_residual = torch.index_add(\n x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor\n )\n return x_plus_residual.view_as(x)\n\n\ndef get_branges_scales(x, sample_drop_ratio=0.0):\n b, n, d = x.shape\n sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)\n brange = (torch.randperm(b, device=x.device))[:sample_subset_size]\n residual_scale_factor = b / sample_subset_size\n return brange, residual_scale_factor\n\n\ndef add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):\n if scaling_vector is None:\n x_flat = x.flatten(1)\n residual = residual.flatten(1)\n x_plus_residual = torch.index_add(\n x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor\n )\n else:\n x_plus_residual = scaled_index_add(\n x,\n brange,\n residual.to(dtype=x.dtype),\n scaling=scaling_vector,\n alpha=residual_scale_factor,\n )\n return x_plus_residual\n\n\nattn_bias_cache: Dict[Tuple, Any] = {}\n\n\ndef get_attn_bias_and_cat(x_list, branges=None):\n \"\"\"\n this will perform the index select, cat the tensors, and provide the attn_bias from cache\n \"\"\"\n batch_sizes = (\n [b.shape[0] for b in branges]\n if branges is not None\n else [x.shape[0] for x in x_list]\n )\n all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))\n if all_shapes not in attn_bias_cache.keys():\n seqlens = []\n for b, x in zip(batch_sizes, x_list):\n for _ in range(b):\n seqlens.append(x.shape[1])\n attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)\n attn_bias._batch_sizes = batch_sizes\n attn_bias_cache[all_shapes] = attn_bias\n\n if branges is not None:\n cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(\n 1, -1, x_list[0].shape[-1]\n )\n else:\n tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)\n cat_tensors = torch.cat(tensors_bs1, dim=1)\n\n return attn_bias_cache[all_shapes], cat_tensors\n\n\ndef drop_add_residual_stochastic_depth_list(\n x_list: List[torch.Tensor],\n residual_func: Callable[[torch.Tensor, Any], torch.Tensor],\n sample_drop_ratio: float = 0.0,\n scaling_vector=None,\n) -> torch.Tensor:\n # 1) generate random set of indices for dropping samples in the batch\n branges_scales = [\n get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list\n ]\n branges = [s[0] for s in branges_scales]\n residual_scale_factors = [s[1] for s in branges_scales]\n\n # 2) get attention bias and index+concat the tensors\n attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)\n\n # 3) apply residual_func to get residual, and split the result\n residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore\n\n outputs = []\n for x, brange, residual, residual_scale_factor in zip(\n x_list, branges, residual_list, residual_scale_factors\n ):\n outputs.append(\n add_residual(\n x, brange, residual, residual_scale_factor, scaling_vector\n ).view_as(x)\n )\n return outputs\n\n\nclass NestedTensorBlock(Block):\n def forward_nested(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:\n \"\"\"\n x_list contains a list of tensors to nest together and run\n \"\"\"\n assert isinstance(self.attn, MemEffAttention)\n\n if self.training and self.sample_drop_ratio > 0.0:\n\n def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:\n return self.attn(self.norm1(x), attn_bias=attn_bias)\n\n def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:\n return self.mlp(self.norm2(x))\n\n x_list = drop_add_residual_stochastic_depth_list(\n x_list,\n residual_func=attn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n scaling_vector=(\n self.ls1.gamma if isinstance(self.ls1, LayerScale) else None\n ),\n )\n x_list = drop_add_residual_stochastic_depth_list(\n x_list,\n residual_func=ffn_residual_func,\n sample_drop_ratio=self.sample_drop_ratio,\n scaling_vector=(\n self.ls2.gamma if isinstance(self.ls1, LayerScale) else None\n ),\n )\n return x_list\n else:\n\n def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:\n return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))\n\n def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:\n return self.ls2(self.mlp(self.norm2(x)))\n\n attn_bias, x = get_attn_bias_and_cat(x_list)\n x = x + attn_residual_func(x, attn_bias=attn_bias)\n x = x + ffn_residual_func(x)\n return attn_bias.split(x)\n\n def forward(self, x_or_x_list):\n if isinstance(x_or_x_list, torch.Tensor):\n return super(NestedTensorBlock, self).forward(x_or_x_list)\n elif isinstance(x_or_x_list, list):\n assert (\n XFORMERS_AVAILABLE\n ), \"Please install xFormers for nested tensors usage\"\n return self.forward_nested(x_or_x_list)\n else:\n raise AssertionError\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/dino_head.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn.init import trunc_normal_\nfrom torch.nn.utils import weight_norm\n\n\nclass DINOHead(nn.Module):\n def __init__(\n self,\n in_dim,\n out_dim,\n use_bn=False,\n nlayers=3,\n hidden_dim=2048,\n bottleneck_dim=256,\n mlp_bias=True,\n ):\n super().__init__()\n nlayers = max(nlayers, 1)\n self.mlp = _build_mlp(\n nlayers,\n in_dim,\n bottleneck_dim,\n hidden_dim=hidden_dim,\n use_bn=use_bn,\n bias=mlp_bias,\n )\n self.apply(self._init_weights)\n self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))\n self.last_layer.weight_g.data.fill_(1)\n\n def _init_weights(self, m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=0.02)\n if isinstance(m, nn.Linear) and m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def forward(self, x):\n x = self.mlp(x)\n eps = 1e-6 if x.dtype == torch.float16 else 1e-12\n x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)\n x = self.last_layer(x)\n return x\n\n\ndef _build_mlp(\n nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True\n):\n if nlayers == 1:\n return nn.Linear(in_dim, bottleneck_dim, bias=bias)\n else:\n layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]\n if use_bn:\n layers.append(nn.BatchNorm1d(hidden_dim))\n layers.append(nn.GELU())\n for _ in range(nlayers - 2):\n layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))\n if use_bn:\n layers.append(nn.BatchNorm1d(hidden_dim))\n layers.append(nn.GELU())\n layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))\n return nn.Sequential(*layers)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/drop_path.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py\n\n\nimport torch.nn as nn\n\n\ndef drop_path(x, drop_prob: float = 0.0, training: bool = False):\n if drop_prob == 0.0 or not training:\n return x\n keep_prob = 1 - drop_prob\n shape = (x.shape[0],) + (1,) * (\n x.ndim - 1\n ) # work with diff dim tensors, not just 2D ConvNets\n random_tensor = x.new_empty(shape).bernoulli_(keep_prob)\n if keep_prob > 0.0:\n random_tensor.div_(keep_prob)\n output = x * random_tensor\n return output\n\n\nclass DropPath(nn.Module):\n \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\"\"\"\n\n def __init__(self, drop_prob=None):\n super(DropPath, self).__init__()\n self.drop_prob = drop_prob\n\n def forward(self, x):\n return drop_path(x, self.drop_prob, self.training)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/layer_scale.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110\n\nfrom typing import Union\n\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor\n\n\nclass LayerScale(nn.Module):\n def __init__(\n self,\n dim: int,\n init_values: Union[float, Tensor] = 1e-5,\n inplace: bool = False,\n ) -> None:\n super().__init__()\n self.inplace = inplace\n self.gamma = nn.Parameter(init_values * torch.ones(dim))\n\n def forward(self, x: Tensor) -> Tensor:\n return x.mul_(self.gamma) if self.inplace else x * self.gamma\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/mlp.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py\n\n\nfrom typing import Callable, Optional\n\nfrom torch import Tensor, nn\n\n\nclass Mlp(nn.Module):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = nn.GELU,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x: Tensor) -> Tensor:\n x = self.fc1(x)\n x = self.act(x)\n x = self.drop(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/patch_embed.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# References:\n# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py\n# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py\n\nfrom typing import Callable, Optional, Tuple, Union\n\nimport torch.nn as nn\nfrom torch import Tensor\n\n\ndef make_2tuple(x):\n if isinstance(x, tuple):\n assert len(x) == 2\n return x\n\n assert isinstance(x, int)\n return (x, x)\n\n\nclass PatchEmbed(nn.Module):\n \"\"\"\n 2D image to patch embedding: (B,C,H,W) -> (B,N,D)\n\n Args:\n img_size: Image size.\n patch_size: Patch token size.\n in_chans: Number of input image channels.\n embed_dim: Number of linear projection output channels.\n norm_layer: Normalization layer.\n \"\"\"\n\n def __init__(\n self,\n img_size: Union[int, Tuple[int, int]] = 224,\n patch_size: Union[int, Tuple[int, int]] = 16,\n in_chans: int = 3,\n embed_dim: int = 768,\n norm_layer: Optional[Callable] = None,\n flatten_embedding: bool = True,\n ) -> None:\n super().__init__()\n\n image_HW = make_2tuple(img_size)\n patch_HW = make_2tuple(patch_size)\n patch_grid_size = (\n image_HW[0] // patch_HW[0],\n image_HW[1] // patch_HW[1],\n )\n\n self.img_size = image_HW\n self.patch_size = patch_HW\n self.patches_resolution = patch_grid_size\n self.num_patches = patch_grid_size[0] * patch_grid_size[1]\n\n self.in_chans = in_chans\n self.embed_dim = embed_dim\n\n self.flatten_embedding = flatten_embedding\n\n self.proj = nn.Conv2d(\n in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW\n )\n self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()\n\n def forward(self, x: Tensor) -> Tensor:\n _, _, H, W = x.shape\n patch_H, patch_W = self.patch_size\n\n assert (\n H % patch_H == 0\n ), f\"Input image height {H} is not a multiple of patch height {patch_H}\"\n assert (\n W % patch_W == 0\n ), f\"Input image width {W} is not a multiple of patch width: {patch_W}\"\n\n x = self.proj(x) # B C H W\n H, W = x.size(2), x.size(3)\n x = x.flatten(2).transpose(1, 2) # B HW C\n x = self.norm(x)\n if not self.flatten_embedding:\n x = x.reshape(-1, H, W, self.embed_dim) # B H W C\n return x\n\n def flops(self) -> float:\n Ho, Wo = self.patches_resolution\n flops = (\n Ho\n * Wo\n * self.embed_dim\n * self.in_chans\n * (self.patch_size[0] * self.patch_size[1])\n )\n if self.norm is not None:\n flops += Ho * Wo * self.embed_dim\n return flops\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/backbones/metadinov2/swiglu_ffn.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Callable, Optional\n\nimport torch.nn.functional as F\nfrom torch import Tensor, nn\n\n\nclass SwiGLUFFN(nn.Module):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = None,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)\n self.w3 = nn.Linear(hidden_features, out_features, bias=bias)\n\n def forward(self, x: Tensor) -> Tensor:\n x12 = self.w12(x)\n x1, x2 = x12.chunk(2, dim=-1)\n hidden = F.silu(x1) * x2\n return self.w3(hidden)\n\n\ntry:\n from xformers.ops import SwiGLU\n\n XFORMERS_AVAILABLE = True\nexcept ImportError:\n SwiGLU = SwiGLUFFN\n XFORMERS_AVAILABLE = False\n\n\nclass SwiGLUFFNFused(SwiGLU):\n def __init__(\n self,\n in_features: int,\n hidden_features: Optional[int] = None,\n out_features: Optional[int] = None,\n act_layer: Callable[..., nn.Module] = None,\n drop: float = 0.0,\n bias: bool = True,\n ) -> None:\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8\n super().__init__(\n in_features=in_features,\n hidden_features=hidden_features,\n out_features=out_features,\n bias=bias,\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/encoder.py", "content": "import torch\nimport torch.nn as nn\n\nfrom unidepth.models.backbones import ConvNeXt, ConvNeXtV2, _make_dinov2_model\n\n\nclass ModelWrap(nn.Module):\n def __init__(self, model) -> None:\n super().__init__()\n self.backbone = model\n\n def forward(self, x, *args, **kwargs):\n features = []\n for layer in self.backbone.features:\n x = layer(x)\n features.append(x)\n return features\n\n\ndef convnextv2_base(config, **kwargs):\n model = ConvNeXtV2(\n depths=[3, 3, 27, 3],\n dims=[128, 256, 512, 1024],\n output_idx=config.get(\"output_idx\", [3, 6, 33, 36]),\n use_checkpoint=config.get(\"use_checkpoint\", False),\n **kwargs,\n )\n url = \"https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt\"\n state_dict = torch.hub.load_state_dict_from_url(\n url, map_location=\"cpu\", progress=False\n )[\"model\"]\n info = model.load_state_dict(state_dict, strict=False)\n print(info)\n return model\n\n\ndef convnextv2_large(config, **kwargs):\n model = ConvNeXtV2(\n depths=[3, 3, 27, 3],\n dims=[192, 384, 768, 1536],\n output_idx=config.get(\"output_idx\", [3, 6, 33, 36]),\n use_checkpoint=config.get(\"use_checkpoint\", False),\n **kwargs,\n )\n url = \"https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt\"\n state_dict = torch.hub.load_state_dict_from_url(\n url, map_location=\"cpu\", progress=False\n )[\"model\"]\n info = model.load_state_dict(state_dict, strict=False)\n print(info)\n return model\n\n\ndef convnextv2_large_mae(config, **kwargs):\n model = ConvNeXtV2(\n depths=[3, 3, 27, 3],\n dims=[192, 384, 768, 1536],\n output_idx=config.get(\"output_idx\", [3, 6, 33, 36]),\n use_checkpoint=config.get(\"use_checkpoint\", False),\n **kwargs,\n )\n url = \"https://dl.fbaipublicfiles.com/convnext/convnextv2/pt_only/convnextv2_large_1k_224_fcmae.pt\"\n state_dict = torch.hub.load_state_dict_from_url(\n url, map_location=\"cpu\", progress=False\n )[\"model\"]\n info = model.load_state_dict(state_dict, strict=False)\n print(info)\n return model\n\n\ndef convnextv2_huge(config, **kwargs):\n model = ConvNeXtV2(\n depths=[3, 3, 27, 3],\n dims=[352, 704, 1408, 2816],\n output_idx=config.get(\"output_idx\", [3, 6, 33, 36]),\n use_checkpoint=config.get(\"use_checkpoint\", False),\n **kwargs,\n )\n url = \"https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt\"\n state_dict = torch.hub.load_state_dict_from_url(\n url, map_location=\"cpu\", progress=False\n )[\"model\"]\n info = model.load_state_dict(state_dict, strict=False)\n print(info)\n return model\n\n\ndef convnextv2_huge_mae(config, **kwargs):\n model = ConvNeXtV2(\n depths=[3, 3, 27, 3],\n dims=[352, 704, 1408, 2816],\n output_idx=config.get(\"output_idx\", [3, 6, 33, 36]),\n use_checkpoint=config.get(\"use_checkpoint\", False),\n **kwargs,\n )\n url = \"https://dl.fbaipublicfiles.com/convnext/convnextv2/pt_only/convnextv2_huge_1k_224_fcmae.pt\"\n state_dict = torch.hub.load_state_dict_from_url(\n url, map_location=\"cpu\", progress=False\n )[\"model\"]\n info = model.load_state_dict(state_dict, strict=False)\n print(info)\n return model\n\n\ndef convnext_large_pt(config, **kwargs):\n model = ConvNeXt(\n depths=[3, 3, 27, 3],\n dims=[192, 384, 768, 1536],\n output_idx=config.get(\"output_idx\", [3, 6, 33, 36]),\n use_checkpoint=config.get(\"use_checkpoint\", False),\n **kwargs,\n )\n from huggingface_hub import hf_hub_download\n from huggingface_hub.utils import disable_progress_bars\n\n from unidepth.models.backbones.convnext import HF_URL, checkpoint_filter_fn\n\n disable_progress_bars()\n repo_id, filename = HF_URL[\"convnext_large_pt\"]\n state_dict = torch.load(hf_hub_download(repo_id=repo_id, filename=filename))\n state_dict = checkpoint_filter_fn(state_dict, model)\n info = model.load_state_dict(state_dict, strict=False)\n print(info)\n return model\n\n\ndef convnext_large(config, **kwargs):\n model = ConvNeXt(\n depths=[3, 3, 27, 3],\n dims=[192, 384, 768, 1536],\n output_idx=config.get(\"output_idx\", [3, 6, 33, 36]),\n use_checkpoint=config.get(\"use_checkpoint\", False),\n drop_path_rate=config.get(\"drop_path\", 0.0),\n **kwargs,\n )\n return model\n\n\ndef dinov2_vits14(config, pretrained: bool = True, **kwargs):\n \"\"\"\n DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.\n \"\"\"\n vit = _make_dinov2_model(\n arch_name=\"vit_small\",\n pretrained=config[\"pretrained\"],\n output_idx=config.get(\"output_idx\", [3, 6, 9, 12]),\n checkpoint=config.get(\"use_checkpoint\", False),\n drop_path_rate=config.get(\"drop_path\", 0.0),\n num_register_tokens=config.get(\"num_register_tokens\", 0),\n use_norm=config.get(\"use_norm\", False),\n export=config.get(\"export\", False),\n interpolate_offset=config.get(\"interpolate_offset\", 0.0),\n **kwargs,\n )\n return vit\n\n\ndef dinov2_vitb14(config, pretrained: bool = True, **kwargs):\n \"\"\"\n DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.\n \"\"\"\n vit = _make_dinov2_model(\n arch_name=\"vit_base\",\n pretrained=config[\"pretrained\"],\n output_idx=config.get(\"output_idx\", [3, 6, 9, 12]),\n checkpoint=config.get(\"use_checkpoint\", False),\n drop_path_rate=config.get(\"drop_path\", 0.0),\n num_register_tokens=config.get(\"num_register_tokens\", 0),\n use_norm=config.get(\"use_norm\", False),\n export=config.get(\"export\", False),\n interpolate_offset=config.get(\"interpolate_offset\", 0.0),\n **kwargs,\n )\n return vit\n\n\ndef dinov2_vitl14(config, pretrained: str = \"\", **kwargs):\n \"\"\"\n DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.\n \"\"\"\n vit = _make_dinov2_model(\n arch_name=\"vit_large\",\n pretrained=config[\"pretrained\"],\n output_idx=config.get(\"output_idx\", [5, 12, 18, 24]),\n checkpoint=config.get(\"use_checkpoint\", False),\n drop_path_rate=config.get(\"drop_path\", 0.0),\n num_register_tokens=config.get(\"num_register_tokens\", 0),\n use_norm=config.get(\"use_norm\", False),\n export=config.get(\"export\", False),\n interpolate_offset=config.get(\"interpolate_offset\", 0.0),\n **kwargs,\n )\n return vit\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv1/__init__.py", "content": "from .unidepthv1 import UniDepthV1\n\n__all__ = [\n \"UniDepthV1\",\n]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv1/decoder.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom typing import List, Tuple\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom timm.models.layers import trunc_normal_\n\nfrom unidepth.layers import (MLP, AttentionBlock, ConvUpsample, NystromBlock,\n PositionEmbeddingSine)\nfrom unidepth.utils.geometric import flat_interpolate, generate_rays\nfrom unidepth.utils.misc import max_stack\nfrom unidepth.utils.sht import rsh_cart_8\n\n\nclass ListAdapter(nn.Module):\n def __init__(self, input_dims: List[int], hidden_dim: int):\n super().__init__()\n self.input_adapters = nn.ModuleList([])\n self.num_chunks = len(input_dims)\n for input_dim in input_dims:\n self.input_adapters.append(\n nn.Sequential(\n nn.LayerNorm(input_dim), nn.Linear(input_dim, hidden_dim), nn.GELU()\n )\n )\n\n def forward(self, x: torch.Tensor, splits: torch.Tensor) -> torch.Tensor:\n xs = torch.split(x, splits.int().tolist(), dim=-1)\n xs = [adapter(x) for x, adapter in zip(xs, self.input_adapters)]\n return torch.cat(xs, dim=-1)\n\n\nclass CameraHead(nn.Module):\n def __init__(\n self,\n input_dim: int,\n hidden_dim: int,\n num_heads: int = 8,\n expansion: int = 4,\n depth: int = 4,\n dropout: float = 0.0,\n layer_scale: float = 1.0,\n **kwargs,\n ):\n super().__init__()\n\n self.aggregate = AttentionBlock(\n hidden_dim,\n num_heads=1,\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n )\n self.latents_pos = nn.Parameter(\n torch.randn(1, 4, hidden_dim), requires_grad=True\n )\n self.layers = nn.ModuleList([])\n self.in_features = MLP(hidden_dim, expansion=2, dropout=dropout)\n for _ in range(depth):\n blk = AttentionBlock(\n hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n )\n self.layers.append(blk)\n self.out = MLP(hidden_dim, expansion=2, dropout=0.0, output_dim=1)\n self.cls_project = nn.Sequential(\n nn.LayerNorm(input_dim),\n nn.Linear(input_dim, hidden_dim // 2),\n nn.GELU(),\n nn.Linear(hidden_dim // 2, hidden_dim),\n )\n\n def forward(self, features, cls_tokens, pos_embed) -> torch.Tensor:\n features = features.unbind(dim=-1)\n cls_tokens = self.cls_project(cls_tokens)\n features_stack = torch.cat(features, dim=1)\n features_stack = features_stack + pos_embed\n latents_pos = self.latents_pos.expand(cls_tokens.shape[0], -1, -1)\n features_stack = self.in_features(features_stack)\n features = torch.cat((features_stack, cls_tokens), dim=1)\n cls_tokens = self.aggregate(cls_tokens, context=features, pos_embed=latents_pos)\n for i, layer in enumerate(self.layers):\n cls_tokens = layer(cls_tokens, pos_embed=latents_pos)\n\n # project\n x = self.out(cls_tokens).squeeze(-1)\n camera_intrinsics = torch.zeros(\n x.shape[0], 3, 3, device=x.device, requires_grad=False\n )\n camera_intrinsics[:, 0, 0] = x[:, 0].exp()\n camera_intrinsics[:, 1, 1] = x[:, 1].exp()\n camera_intrinsics[:, 0, 2] = x[:, 2].sigmoid()\n camera_intrinsics[:, 1, 2] = x[:, 3].sigmoid()\n camera_intrinsics[:, 2, 2] = 1.0\n return camera_intrinsics\n\n def set_shapes(self, shapes: Tuple[int, int]):\n self.shapes = shapes\n\n\nclass DepthHead(nn.Module):\n def __init__(\n self,\n hidden_dim: int,\n num_heads: int = 8,\n expansion: int = 4,\n depths: int | list[int] = 4,\n camera_dim: int = 256,\n num_resolutions: int = 4,\n dropout: float = 0.0,\n layer_scale: float = 1.0,\n **kwargs,\n ) -> None:\n super().__init__()\n if isinstance(depths, int):\n depths = [depths] * 3\n assert len(depths) == 3\n\n self.project_rays16 = MLP(\n camera_dim, expansion=expansion, dropout=dropout, output_dim=hidden_dim\n )\n self.project_rays8 = MLP(\n camera_dim, expansion=expansion, dropout=dropout, output_dim=hidden_dim // 2\n )\n self.project_rays4 = MLP(\n camera_dim, expansion=expansion, dropout=dropout, output_dim=hidden_dim // 4\n )\n self.to_latents = MLP(hidden_dim, expansion=2, dropout=dropout)\n\n self.features_channel_cat = nn.Linear(hidden_dim * num_resolutions, hidden_dim)\n\n self.up8 = ConvUpsample(\n hidden_dim, expansion=expansion, layer_scale=layer_scale\n )\n self.up4 = ConvUpsample(\n hidden_dim // 2, expansion=expansion, layer_scale=layer_scale\n )\n self.up2 = ConvUpsample(\n hidden_dim // 4, expansion=expansion, layer_scale=layer_scale\n )\n\n self.layers_16 = nn.ModuleList([])\n self.layers_8 = nn.ModuleList([])\n self.layers_4 = nn.ModuleList([])\n self.aggregate_16 = AttentionBlock(\n hidden_dim,\n num_heads=1,\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n context_dim=hidden_dim,\n )\n self.prompt_camera = AttentionBlock(\n hidden_dim,\n num_heads=1,\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n context_dim=hidden_dim,\n )\n for i, (blk_lst, depth) in enumerate(\n zip([self.layers_16, self.layers_8, self.layers_4], depths)\n ):\n attn_cls = AttentionBlock if i == 0 else NystromBlock\n for _ in range(depth):\n blk_lst.append(\n attn_cls(\n hidden_dim // (2**i),\n num_heads=num_heads // (2**i),\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n )\n )\n\n self.out2 = nn.Conv2d(hidden_dim // 8, 1, 3, padding=1)\n self.out4 = nn.Conv2d(hidden_dim // 4, 1, 3, padding=1)\n self.out8 = nn.Conv2d(hidden_dim // 2, 1, 3, padding=1)\n\n def set_original_shapes(self, shapes: Tuple[int, int]):\n self.original_shapes = shapes\n\n def set_shapes(self, shapes: Tuple[int, int]):\n self.shapes = shapes\n\n def forward(\n self, features: torch.Tensor, rays_hr: torch.Tensor, pos_embed, level_embed\n ) -> torch.Tensor:\n features = features.unbind(dim=-1)\n shapes = self.shapes\n rays_hr = rays_hr.detach()\n\n # camera_embedding\n rays_embedding_16 = F.normalize(\n flat_interpolate(rays_hr, old=self.original_shapes, new=shapes), dim=-1\n )\n rays_embedding_8 = F.normalize(\n flat_interpolate(\n rays_hr, old=self.original_shapes, new=[x * 2 for x in shapes]\n ),\n dim=-1,\n )\n rays_embedding_4 = F.normalize(\n flat_interpolate(\n rays_hr, old=self.original_shapes, new=[x * 4 for x in shapes]\n ),\n dim=-1,\n )\n rays_embedding_16 = self.project_rays16(rsh_cart_8(rays_embedding_16))\n rays_embedding_8 = self.project_rays8(rsh_cart_8(rays_embedding_8))\n rays_embedding_4 = self.project_rays4(rsh_cart_8(rays_embedding_4))\n features_tokens = torch.cat(features, dim=1)\n features_tokens_pos = pos_embed + level_embed\n\n # Generate latents with init as pooled features\n features_channels = torch.cat(features, dim=-1)\n features_16 = self.features_channel_cat(features_channels)\n latents_16 = self.to_latents(\n flat_interpolate(features_16, old=self.shapes, new=shapes, antialias=False)\n )\n\n # Aggregate features: F -> D\n latents_16 = self.aggregate_16(\n latents_16, context=features_tokens, pos_embed_context=features_tokens_pos\n )\n\n # Aggregate camera: D- > D|E\n latents_16 = self.prompt_camera(latents_16, context=rays_embedding_16)\n\n # Block 16 - Out 8\n for layer in self.layers_16:\n latents_16 = layer(latents_16, pos_embed=rays_embedding_16)\n latents_8 = self.up8(\n rearrange(\n latents_16 + rays_embedding_16,\n \"b (h w) c -> b c h w\",\n h=shapes[0],\n w=shapes[1],\n ).contiguous()\n )\n out8 = self.out8(\n rearrange(\n latents_8, \"b (h w) c -> b c h w\", h=shapes[0] * 2, w=shapes[1] * 2\n )\n )\n\n # Block 8 - Out 4\n for layer in self.layers_8:\n latents_8 = layer(latents_8, pos_embed=rays_embedding_8)\n latents_4 = self.up4(\n rearrange(\n latents_8 + rays_embedding_8,\n \"b (h w) c -> b c h w\",\n h=shapes[0] * 2,\n w=shapes[1] * 2,\n ).contiguous()\n )\n out4 = self.out4(\n rearrange(\n latents_4, \"b (h w) c -> b c h w\", h=shapes[0] * 4, w=shapes[1] * 4\n )\n )\n\n # Block 4 - Out 2\n for layer in self.layers_4:\n latents_4 = layer(latents_4, pos_embed=rays_embedding_4)\n latents_2 = self.up2(\n rearrange(\n latents_4 + rays_embedding_4,\n \"b (h w) c -> b c h w\",\n h=shapes[0] * 4,\n w=shapes[1] * 4,\n ).contiguous()\n )\n out2 = self.out2(\n rearrange(\n latents_2, \"b (h w) c -> b c h w\", h=shapes[0] * 8, w=shapes[1] * 8\n )\n )\n\n # Depth features\n proj_latents_16 = rearrange(\n latents_16, \"b (h w) c -> b c h w\", h=shapes[0], w=shapes[1]\n ).contiguous()\n\n # MS Outputs\n out2 = out2.clamp(-10.0, 10.0).exp()\n out4 = out4.clamp(-10.0, 10.0).exp()\n out8 = out8.clamp(-10.0, 10.0).exp()\n\n return out8, out4, out2, proj_latents_16\n\n\nclass Decoder(nn.Module):\n def __init__(\n self,\n config,\n *args,\n **kwargs,\n ):\n super().__init__()\n self.build(config)\n self.apply(self._init_weights)\n self.test_fixed_camera = False\n self.skip_camera = False\n\n def _init_weights(self, m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=0.02)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Conv2d):\n trunc_normal_(m.weight, std=0.02)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.LayerNorm):\n nn.init.constant_(m.bias, 0)\n nn.init.constant_(m.weight, 1.0)\n\n def get_adapted_features(self, features_flat, splits):\n features_flat_cat = torch.cat(features_flat, dim=-1)\n features_projected = self.input_adapter(\n features_flat_cat, splits\n ) # list [b hw c] shapes\n features = torch.chunk(features_projected, len(splits), dim=-1)\n return features\n\n def run_camera(self, cls_tokens, features, pos_embed, original_shapes, rays):\n # get cls tokens projections\n cls_tokens_splits = torch.tensor(\n [x.shape[-1] for x in cls_tokens],\n device=features.device,\n requires_grad=False,\n dtype=features.dtype,\n )\n cls_tokens = torch.cat(cls_tokens, dim=-1)\n cls_tokens = self.token_adapter(cls_tokens, cls_tokens_splits)\n cls_tokens = torch.cat(\n torch.chunk(cls_tokens, len(cls_tokens_splits), dim=-1), dim=1\n )\n\n # camera layer\n intrinsics = self.camera_layer(\n features=features, cls_tokens=cls_tokens, pos_embed=pos_embed\n )\n intrinsics[:, 0, 0] = max(original_shapes) / 2 * intrinsics[:, 0, 0]\n intrinsics[:, 1, 1] = max(original_shapes) / 2 * intrinsics[:, 1, 1]\n intrinsics[:, 0, 2] = intrinsics[:, 0, 2] * original_shapes[1]\n intrinsics[:, 1, 2] = intrinsics[:, 1, 2] * original_shapes[0]\n if not self.test_fixed_camera:\n rays, _ = generate_rays(intrinsics, original_shapes, noisy=False)\n\n return intrinsics, rays\n\n def forward(self, inputs, image_metas) -> torch.Tensor:\n B, _, H, W = inputs[\"image\"].shape\n device = inputs[\"image\"].device\n\n # make stride happy?\n original_encoder_outputs = [x.contiguous() for x in inputs[\"encoder_outputs\"]]\n cls_tokens = [x.contiguous() for x in inputs[\"cls_tokens\"]]\n\n # collect features and tokens\n original_encoder_outputs = [\n max_stack(original_encoder_outputs[i:j])\n for i, j in self.slices_encoder_range\n ]\n # detach tokens for camera\n cls_tokens = [\n cls_tokens[-i - 1].detach() for i in range(len(self.slices_encoder_range))\n ]\n\n # get features in b n d format\n # level shapes, the shape per level, for swin like [[128, 128], [64, 64],...], for vit [[32,32]] -> mult times resolutions\n resolutions = [\n tuple(sorted([x.shape[1], x.shape[2]])) for x in original_encoder_outputs\n ]\n level_shapes = sorted(list(set(resolutions)))[::-1]\n\n if len(level_shapes) == 1:\n level_shapes = level_shapes * self.num_resolutions\n input_shapes = [\n level_shapes[i]\n for i, (start, end) in enumerate(self.slices_encoder)\n for _ in range(end - start)\n ]\n common_shape = level_shapes[-2]\n\n # input shapes repeat shapes for each level, times the amount of the layers:\n features_flat = [\n flat_interpolate(\n rearrange(x, \"b h w c -> b (h w) c\"), old=input_shape, new=common_shape\n )\n for x, input_shape in zip(original_encoder_outputs, input_shapes)\n ]\n features_splits = torch.tensor(\n [x.shape[-1] for x in features_flat],\n device=device,\n requires_grad=False,\n dtype=torch.float32,\n )\n\n # input adapter, then do mean of features in same blocks\n features = self.get_adapted_features(features_flat, features_splits)\n features = torch.stack(features, dim=-1)\n\n # positional embeddings, spatial and level\n level_embed = torch.cat(\n [\n self.level_embed_layer(self.level_embeds)[i : i + 1]\n .unsqueeze(0)\n .repeat(B, common_shape[0] * common_shape[1], 1)\n for i in range(self.num_resolutions)\n ],\n dim=1,\n )\n pos_embed = self.pos_embed(\n torch.zeros(\n B,\n 1,\n common_shape[0],\n common_shape[1],\n device=device,\n requires_grad=False,\n )\n )\n pos_embed = rearrange(pos_embed, \"b c h w -> b (h w) c\").repeat(\n 1, self.num_resolutions, 1\n )\n\n self.camera_layer.set_shapes(common_shape)\n intrinsics, rays = (\n self.run_camera(\n cls_tokens,\n features=features,\n pos_embed=pos_embed + level_embed,\n original_shapes=(H, W),\n rays=inputs.get(\"rays\", None),\n )\n if not self.skip_camera\n else (inputs[\"K\"], inputs[\"rays\"])\n )\n\n # run bulk of the model\n self.depth_layer.set_shapes(common_shape)\n self.depth_layer.set_original_shapes((H, W))\n out8, out4, out2, depth_features = self.depth_layer(\n features=features,\n rays_hr=rays,\n pos_embed=pos_embed,\n level_embed=level_embed,\n )\n\n return intrinsics, [out8, out4, out2], depth_features\n\n @torch.jit.ignore\n def no_weight_decay_keywords(self):\n return {\"latents_pos\", \"level_embeds\"}\n\n def build(self, config):\n depth = config[\"model\"][\"pixel_decoder\"][\"depths\"]\n input_dims = config[\"model\"][\"pixel_encoder\"][\"embed_dims\"]\n hidden_dim = config[\"model\"][\"pixel_decoder\"][\"hidden_dim\"]\n num_heads = config[\"model\"][\"num_heads\"]\n expansion = config[\"model\"][\"expansion\"]\n dropout = config[\"model\"][\"pixel_decoder\"][\"dropout\"]\n depths_encoder = config[\"model\"][\"pixel_encoder\"][\"depths\"]\n layer_scale = 1.0\n\n self.depth = depth\n self.dim = hidden_dim\n self.downsample = 4\n self.num_heads = num_heads\n self.num_resolutions = len(depths_encoder)\n self.depths_encoder = depths_encoder\n\n self.slices_encoder_single = list(\n zip([d - 1 for d in self.depths_encoder], self.depths_encoder)\n )\n self.slices_encoder_range = list(\n zip([0, *self.depths_encoder[:-1]], self.depths_encoder)\n )\n cls_token_input_dims = [input_dims[-i - 1] for i in range(len(depths_encoder))]\n\n input_dims = [input_dims[d - 1] for d in depths_encoder]\n self.slices_encoder = self.slices_encoder_single\n\n # adapt from encoder features, just project\n self.input_adapter = ListAdapter(input_dims, hidden_dim)\n self.token_adapter = ListAdapter(cls_token_input_dims, hidden_dim)\n\n # camera layer\n self.camera_layer = CameraHead(\n input_dim=hidden_dim,\n hidden_dim=hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n depth=2,\n dropout=dropout,\n layer_scale=layer_scale,\n )\n\n self.depth_layer = DepthHead(\n hidden_dim=hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n depths=depth,\n dropout=dropout,\n camera_dim=81,\n num_resolutions=self.num_resolutions,\n layer_scale=layer_scale,\n )\n\n # transformer part\n self.pos_embed = PositionEmbeddingSine(hidden_dim // 2, normalize=True)\n self.level_embeds = nn.Parameter(\n torch.randn(len(input_dims), hidden_dim), requires_grad=True\n )\n self.level_embed_layer = nn.Sequential(\n nn.Linear(hidden_dim, hidden_dim),\n nn.GELU(),\n nn.Linear(hidden_dim, hidden_dim),\n nn.LayerNorm(hidden_dim),\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv1/unidepthv1.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport importlib\nfrom copy import deepcopy\nfrom math import ceil\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms.functional as TF\nfrom einops import rearrange\nfrom huggingface_hub import PyTorchModelHubMixin\n\nfrom unidepth.models.unidepthv1.decoder import Decoder\nfrom unidepth.utils.constants import (IMAGENET_DATASET_MEAN,\n IMAGENET_DATASET_STD)\nfrom unidepth.utils.distributed import is_main_process\nfrom unidepth.utils.geometric import (generate_rays,\n spherical_zbuffer_to_euclidean)\nfrom unidepth.utils.misc import (get_params, match_gt, match_intrinsics,\n profile_method)\n\nVERBOSE = False\n\n\n# inference helpers\ndef _paddings(image_shape, network_shape):\n cur_h, cur_w = image_shape\n h, w = network_shape\n pad_top, pad_bottom = (h - cur_h) // 2, h - cur_h - (h - cur_h) // 2\n pad_left, pad_right = (w - cur_w) // 2, w - cur_w - (w - cur_w) // 2\n return pad_left, pad_right, pad_top, pad_bottom\n\n\ndef _shapes(image_shape, network_shape):\n h, w = image_shape\n input_ratio = w / h\n output_ratio = network_shape[1] / network_shape[0]\n if output_ratio > input_ratio:\n ratio = network_shape[0] / h\n elif output_ratio <= input_ratio:\n ratio = network_shape[1] / w\n return (ceil(h * ratio - 0.5), ceil(w * ratio - 0.5)), ratio\n\n\ndef _preprocess(rgbs, intrinsics, shapes, pads, ratio, output_shapes):\n (pad_left, pad_right, pad_top, pad_bottom) = pads\n rgbs = F.interpolate(\n rgbs, size=shapes, mode=\"bilinear\", align_corners=False, antialias=True\n )\n rgbs = F.pad(rgbs, (pad_left, pad_right, pad_top, pad_bottom), mode=\"constant\")\n if intrinsics is not None:\n intrinsics = intrinsics.clone()\n intrinsics[:, 0, 0] = intrinsics[:, 0, 0] * ratio\n intrinsics[:, 1, 1] = intrinsics[:, 1, 1] * ratio\n intrinsics[:, 0, 2] = intrinsics[:, 0, 2] * ratio + pad_left\n intrinsics[:, 1, 2] = intrinsics[:, 1, 2] * ratio + pad_top\n return rgbs, intrinsics\n return rgbs, None\n\n\ndef _postprocess(predictions, intrinsics, shapes, pads, ratio, original_shapes):\n (pad_left, pad_right, pad_top, pad_bottom) = pads\n # pred mean, trim paddings, and upsample to input dim\n predictions = sum(\n [\n F.interpolate(\n x.clone(),\n size=shapes,\n mode=\"bilinear\",\n align_corners=False,\n antialias=True,\n )\n for x in predictions\n ]\n ) / len(predictions)\n predictions = predictions[\n ..., pad_top : shapes[0] - pad_bottom, pad_left : shapes[1] - pad_right\n ]\n predictions = F.interpolate(\n predictions,\n size=original_shapes,\n mode=\"bilinear\",\n align_corners=False,\n antialias=True,\n )\n intrinsics[:, 0, 0] = intrinsics[:, 0, 0] / ratio\n intrinsics[:, 1, 1] = intrinsics[:, 1, 1] / ratio\n intrinsics[:, 0, 2] = (intrinsics[:, 0, 2] - pad_left) / ratio\n intrinsics[:, 1, 2] = (intrinsics[:, 1, 2] - pad_top) / ratio\n return predictions, intrinsics\n\n\nclass UniDepthV1(\n nn.Module,\n PyTorchModelHubMixin,\n library_name=\"UniDepth\",\n repo_url=\"https://github.com/lpiccinelli-eth/UniDepth\",\n tags=[\"monocular-metric-depth-estimation\"],\n):\n def __init__(\n self,\n config,\n eps: float = 1e-6,\n **kwargs,\n ):\n super().__init__()\n self.build(config)\n self.build_losses(config)\n self.eps = eps\n\n @profile_method(verbose=VERBOSE)\n def forward_train(self, inputs, image_metas):\n inputs, outputs = self.encode_decode(inputs, image_metas)\n losses = self.compute_losses(outputs, inputs, image_metas)\n return outputs, losses\n\n @profile_method(verbose=VERBOSE)\n def forward_test(self, inputs, image_metas):\n inputs, outputs = self.encode_decode(inputs, image_metas)\n depth_gt = inputs[\"depth\"]\n test_outputs = {}\n test_outputs[\"depth\"] = match_gt(\n outputs[\"depth\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"points\"] = match_gt(\n outputs[\"points\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"confidence\"] = match_gt(\n outputs[\"confidence\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"rays\"] = match_gt(\n outputs[\"rays\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"rays\"] = outputs[\"rays\"] / torch.norm(\n outputs[\"rays\"], dim=1, keepdim=True\n ).clip(min=1e-5)\n test_outputs[\"intrinsics\"] = match_intrinsics(\n outputs[\"intrinsics\"],\n inputs[\"image\"],\n depth_gt,\n padding1=inputs[\"paddings\"],\n padding2=None,\n )\n return test_outputs\n\n def forward(self, inputs, image_metas):\n if self.training:\n return self.forward_train(inputs, image_metas)\n else:\n return self.forward_test(inputs, image_metas)\n\n def encode_decode(self, inputs, image_metas):\n rgbs = inputs[\"image\"]\n B, _, H, W = rgbs.shape\n cameras = inputs[\"camera\"]\n\n # shortcut eval should avoid errors\n if len(image_metas) and \"paddings\" in image_metas[0]:\n inputs[\"paddings\"] = torch.tensor(\n [image_meta[\"paddings\"] for image_meta in image_metas],\n device=self.device,\n )[\n ..., [0, 2, 1, 3]\n ] # lrtb\n inputs[\"depth_paddings\"] = torch.tensor(\n [image_meta[\"depth_paddings\"] for image_meta in image_metas],\n device=self.device,\n )\n if (\n self.training\n ): # at inference we do not have image paddings on top of depth ones (we have not \"crop\" on gt in ContextCrop)\n inputs[\"depth_paddings\"] = inputs[\"depth_paddings\"] + inputs[\"paddings\"]\n\n # Get camera rays for supervision, all in unit sphere\n if inputs.get(\"camera\", None) is not None:\n inputs[\"rays\"] = rearrange(\n inputs[\"camera\"].get_rays(shapes=(B, H, W)), \"b c h w -> b (h w) c\"\n )\n\n # Encode\n encoder_outputs, cls_tokens = self.pixel_encoder(rgbs)\n if \"dino\" in self.pixel_encoder.__class__.__name__.lower():\n encoder_outputs = [\n (x + y.unsqueeze(1)).contiguous()\n for x, y in zip(encoder_outputs, cls_tokens)\n ]\n inputs[\"encoder_outputs\"] = encoder_outputs\n inputs[\"cls_tokens\"] = cls_tokens\n\n # Decode\n pred_intrinsics, predictions, depth_features = self.pixel_decoder(inputs, {})\n predictions = sum(\n [\n F.interpolate(\n x.clone(),\n size=(H, W),\n mode=\"bilinear\",\n align_corners=False,\n antialias=True,\n )\n for x in predictions\n ]\n ) / len(predictions)\n\n # Final 3D points backprojection\n pred_rays, pred_angles = generate_rays(pred_intrinsics, (H, W), noisy=False)\n\n # You may want to use inputs[\"angles\"] if available?\n pred_angles = rearrange(pred_angles, \"b (h w) c -> b c h w\", h=H, w=W)\n\n points_3d = torch.cat((pred_angles, predictions), dim=1)\n points_3d = spherical_zbuffer_to_euclidean(\n points_3d.permute(0, 2, 3, 1)\n ).permute(0, 3, 1, 2)\n\n # Output data, use for loss computation\n outputs = {\n \"angles\": pred_angles,\n \"rays\": pred_rays,\n \"intrinsics\": pred_intrinsics,\n \"points\": points_3d,\n \"depth\": predictions[:, -1:],\n \"cond_features\": depth_features,\n }\n self.pixel_decoder.test_fixed_camera = False\n outputs[\"rays\"] = rearrange(outputs[\"rays\"], \"b (h w) c -> b c h w\", h=H, w=W)\n if \"rays\" in inputs:\n inputs[\"rays\"] = rearrange(inputs[\"rays\"], \"b (h w) c -> b c h w\", h=H, w=W)\n return inputs, outputs\n\n def compute_losses(self, outputs, inputs, image_metas):\n B, _, H, W = inputs[\"image\"].shape\n losses = {\"opt\": {}, \"stat\": {}}\n if (\n not self.training\n ): # only compute losses during training, avoid issues for mismatch size of pred and GT\n return losses\n losses_to_be_computed = list(self.losses.keys())\n\n # depth loss\n si = torch.tensor(\n [x.get(\"si\", False) for x in image_metas], device=self.device\n ).reshape(B)\n loss = self.losses[\"depth\"]\n depth_losses = loss(\n outputs[\"depth\"],\n target=inputs[\"depth\"],\n mask=inputs[\"depth_mask\"].clone(),\n si=si,\n )\n losses[\"opt\"][loss.name] = loss.weight * depth_losses.mean()\n losses_to_be_computed.remove(\"depth\")\n\n # camera loss, here we apply to rays for simplicity\n # in the original training was on angles\n # however, we saw no difference (see supplementary)\n loss = self.losses[\"camera\"]\n camera_losses = loss(outputs[\"rays\"], target=inputs[\"rays\"])\n losses[\"opt\"][loss.name] = loss.weight * camera_losses.mean()\n losses_to_be_computed.remove(\"camera\")\n\n # invariance loss\n flips = torch.tensor(\n [x.get(\"flip\", False) for x in image_metas], device=self.device\n ).reshape(B)\n loss = self.losses[\"invariance\"]\n invariance_losses = loss(\n outputs[\"cond_features\"],\n intrinsics=inputs[\"camera\"].K,\n mask=inputs[\"depth_mask\"],\n flips=flips,\n )\n losses[\"opt\"][loss.name] = loss.weight * invariance_losses.mean()\n losses_to_be_computed.remove(\"invariance\")\n\n # remaining losses, we expect no more losses to be computed\n assert (\n not losses_to_be_computed\n ), f\"Losses {losses_to_be_computed} not computed, revise `compute_loss` method\"\n\n return losses\n\n @torch.no_grad()\n def infer(self, rgbs: torch.Tensor, intrinsics=None, skip_camera=False):\n if rgbs.ndim == 3:\n rgbs = rgbs.unsqueeze(0)\n if intrinsics is not None and intrinsics.ndim == 2:\n intrinsics = intrinsics.unsqueeze(0)\n B, _, H, W = rgbs.shape\n\n rgbs = rgbs.to(self.device)\n if intrinsics is not None:\n intrinsics = intrinsics.to(self.device)\n\n # process image and intrinsiscs (if any) to match network input (slow?)\n if rgbs.max() > 5 or rgbs.dtype == torch.uint8:\n rgbs = rgbs.to(torch.float32).div(255)\n if rgbs.min() >= 0.0 and rgbs.max() <= 1.0:\n rgbs = TF.normalize(\n rgbs,\n mean=IMAGENET_DATASET_MEAN,\n std=IMAGENET_DATASET_STD,\n )\n\n (h, w), ratio = _shapes((H, W), self.image_shape)\n pad_left, pad_right, pad_top, pad_bottom = _paddings((h, w), self.image_shape)\n rgbs, gt_intrinsics = _preprocess(\n rgbs,\n intrinsics,\n (h, w),\n (pad_left, pad_right, pad_top, pad_bottom),\n ratio,\n self.image_shape,\n )\n\n # run encoder\n encoder_outputs, cls_tokens = self.pixel_encoder(rgbs)\n if \"dino\" in self.pixel_encoder.__class__.__name__.lower():\n encoder_outputs = [\n (x + y.unsqueeze(1)).contiguous()\n for x, y in zip(encoder_outputs, cls_tokens)\n ]\n\n # get data for decoder and adapt to given camera\n inputs = {}\n inputs[\"encoder_outputs\"] = encoder_outputs\n inputs[\"cls_tokens\"] = cls_tokens\n inputs[\"image\"] = rgbs\n if gt_intrinsics is not None:\n rays, angles = generate_rays(\n gt_intrinsics, self.image_shape, noisy=self.training\n )\n inputs[\"rays\"] = rays\n inputs[\"angles\"] = angles\n inputs[\"K\"] = gt_intrinsics\n self.pixel_decoder.test_fixed_camera = True\n self.pixel_decoder.skip_camera = skip_camera\n\n # decode all\n pred_intrinsics, predictions, _ = self.pixel_decoder(inputs, {})\n\n # undo the reshaping and get original image size (slow)\n predictions, pred_intrinsics = _postprocess(\n predictions,\n pred_intrinsics,\n self.image_shape,\n (pad_left, pad_right, pad_top, pad_bottom),\n ratio,\n (H, W),\n )\n\n # final 3D points backprojection\n intrinsics = gt_intrinsics if gt_intrinsics is not None else pred_intrinsics\n angles = generate_rays(intrinsics, (H, W), noisy=False)[-1]\n angles = rearrange(angles, \"b (h w) c -> b c h w\", h=H, w=W)\n points_3d = torch.cat((angles, predictions), dim=1)\n points_3d = spherical_zbuffer_to_euclidean(\n points_3d.permute(0, 2, 3, 1)\n ).permute(0, 3, 1, 2)\n\n # output data\n outputs = {\n \"intrinsics\": pred_intrinsics,\n \"points\": points_3d,\n \"depth\": predictions[:, -1:],\n }\n self.pixel_decoder.test_fixed_camera = False\n self.pixel_decoder.skip_camera = False\n return outputs\n\n def load_pretrained(self, model_file):\n device = (\n torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n )\n dict_model = torch.load(model_file, map_location=device)\n\n if \"model\" in dict_model:\n dict_model = dict_model[\"model\"]\n new_state_dict = deepcopy(\n {k.replace(\"module.\", \"\"): v for k, v in dict_model.items()}\n )\n\n info = self.load_state_dict(new_state_dict, strict=False)\n if is_main_process():\n print(\n f\"Loaded from {model_file} for {self.__class__.__name__} results in:\",\n info,\n )\n\n def get_params(self, config):\n if hasattr(self.pixel_encoder, \"get_params\"):\n encoder_p, encoder_lr = self.pixel_encoder.get_params(\n config[\"model\"][\"pixel_encoder\"][\"lr\"],\n config[\"training\"][\"wd\"],\n config[\"training\"][\"ld\"],\n )\n else:\n encoder_p, encoder_lr = get_params(\n self.pixel_encoder,\n config[\"model\"][\"pixel_encoder\"][\"lr\"],\n config[\"training\"][\"wd\"],\n )\n decoder_p, decoder_lr = get_params(\n self.pixel_decoder, config[\"training\"][\"lr\"], config[\"training\"][\"wd\"]\n )\n return [*encoder_p, *decoder_p]\n\n @property\n def device(self):\n return next(self.parameters()).device\n\n def build(self, config):\n mod = importlib.import_module(\"unidepth.models.encoder\")\n pixel_encoder_factory = getattr(mod, config[\"model\"][\"pixel_encoder\"][\"name\"])\n pixel_encoder_config = {\n **config[\"training\"],\n **config[\"data\"],\n **config[\"model\"][\"pixel_encoder\"],\n \"interpolate_offset\": 0.1,\n }\n pixel_encoder = pixel_encoder_factory(pixel_encoder_config)\n\n config[\"model\"][\"pixel_encoder\"][\"patch_size\"] = (\n 14 if \"dino\" in config[\"model\"][\"pixel_encoder\"][\"name\"] else 16\n )\n pixel_encoder_embed_dims = (\n pixel_encoder.embed_dims\n if hasattr(pixel_encoder, \"embed_dims\")\n else [getattr(pixel_encoder, \"embed_dim\") * 2**i for i in range(4)]\n )\n config[\"model\"][\"pixel_encoder\"][\"embed_dim\"] = getattr(\n pixel_encoder, \"embed_dim\"\n )\n config[\"model\"][\"pixel_encoder\"][\"embed_dims\"] = pixel_encoder_embed_dims\n config[\"model\"][\"pixel_encoder\"][\"depths\"] = pixel_encoder.depths\n\n self.pixel_encoder = pixel_encoder\n self.pixel_decoder = Decoder(config)\n self.image_shape = config[\"data\"][\"image_shape\"]\n\n def build_losses(self, config):\n self.losses = {}\n for loss_name, loss_config in config[\"training\"].get(\"losses\", {}).items():\n mod = importlib.import_module(\"unidepth.ops.losses\")\n loss_factory = getattr(mod, loss_config[\"name\"])\n self.losses[loss_name] = loss_factory.build(loss_config)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv2/__init__.py", "content": "from .unidepthv2 import UniDepthV2\nfrom .unidepthv2_old import UniDepthV2old\n\n__all__ = [\n \"UniDepthV2\",\n \"UniDepthV2old\",\n]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv2/decoder.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom timm.models.layers import trunc_normal_\n\nfrom unidepth.layers import (MLP, AttentionBlock, AttentionLayer,\n PositionEmbeddingSine, ResUpsampleBil)\nfrom unidepth.utils.coordinate import coords_grid\nfrom unidepth.utils.geometric import flat_interpolate\nfrom unidepth.utils.positional_embedding import generate_fourier_features\n\n\ndef orthonormal_init(num_tokens, dims):\n\n pe = torch.randn(num_tokens, dims)\n\n # Apply Gram-Schmidt process to make the matrix orthonormal\n for i in range(num_tokens):\n for j in range(i):\n # Subtract the projection of current row onto previous row\n pe[i] -= torch.dot(pe[i], pe[j]) * pe[j]\n\n # Normalize the current row\n pe[i] = F.normalize(pe[i], p=2, dim=0)\n\n return pe\n\n\nclass ListAdapter(nn.Module):\n def __init__(self, input_dims: list[int], hidden_dim: int):\n super().__init__()\n self.input_adapters = nn.ModuleList([])\n self.num_chunks = len(input_dims)\n for input_dim in input_dims:\n self.input_adapters.append(nn.Linear(input_dim, hidden_dim))\n\n def forward(self, xs: torch.Tensor) -> list[torch.Tensor]:\n outs = [self.input_adapters[i](x) for i, x in enumerate(xs)]\n return outs\n\n\nclass CameraHead(nn.Module):\n def __init__(\n self,\n hidden_dim: int,\n num_heads: int = 8,\n expansion: int = 4,\n dropout: float = 0.0,\n layer_scale: float = 1.0,\n **kwargs,\n ):\n super().__init__()\n self.num_params = 4\n\n self.aggregate1 = AttentionBlock(\n hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n use_bias=False,\n )\n self.aggregate2 = AttentionBlock(\n hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n use_bias=False,\n )\n self.latents_pos = nn.Parameter(\n torch.randn(1, self.num_params, hidden_dim), requires_grad=True\n )\n self.project = MLP(\n hidden_dim, expansion=1, dropout=dropout, output_dim=hidden_dim\n )\n self.out_pinhole = MLP(hidden_dim, expansion=1, dropout=dropout, output_dim=1)\n\n def fill_intrinsics(self, x):\n fx, fy, cx, cy = x.unbind(dim=-1)\n fx = torch.exp(fx)\n fy = torch.exp(fy)\n cx = torch.sigmoid(cx)\n cy = torch.sigmoid(cy)\n diagonal = (self.shapes[0] ** 2 + self.shapes[1] ** 2) ** 0.5\n correction_tensor = torch.tensor(\n [0.7 * diagonal, 0.7 * diagonal, self.shapes[1], self.shapes[0]],\n device=x.device,\n dtype=x.dtype,\n )\n intrinsics = torch.stack([fx, fy, cx, cy], dim=1)\n intrinsics = correction_tensor.unsqueeze(0) * intrinsics\n return intrinsics\n\n def forward(self, features, cls_tokens, pos_embed) -> torch.Tensor:\n features = features.unbind(dim=-1)\n tokens = self.project(cls_tokens)\n\n latents_pos = self.latents_pos.expand(cls_tokens.shape[0], -1, -1)\n tokens = self.aggregate1(tokens, pos_embed=latents_pos)\n tokens = self.aggregate2(tokens, pos_embed=latents_pos)\n\n x = self.out_pinhole(tokens.clone()).squeeze(-1)\n camera_intrinsics = self.fill_intrinsics(x)\n return camera_intrinsics\n\n def set_shapes(self, shapes: tuple[int, int]):\n self.shapes = shapes\n\n\nclass DepthHead(nn.Module):\n def __init__(\n self,\n hidden_dim: int,\n num_heads: int = 8,\n expansion: int = 4,\n depths: int | list[int] = 4,\n camera_dim: int = 256,\n dropout: float = 0.0,\n kernel_size: int = 7,\n layer_scale: float = 1.0,\n out_dim: int = 1,\n use_norm=False,\n num_prompt_blocks=1,\n **kwargs,\n ) -> None:\n super().__init__()\n\n self.camera_dim = camera_dim\n self.out_dim = out_dim\n self.hidden_dim = hidden_dim\n\n self.ups = nn.ModuleList([])\n self.depth_mlp = nn.ModuleList([])\n self.process_features = nn.ModuleList([])\n self.project_features = nn.ModuleList([])\n self.prompt_camera = nn.ModuleList([])\n mult = 2\n self.to_latents = nn.Linear(hidden_dim, hidden_dim)\n\n for _ in range(4):\n self.prompt_camera.append(\n AttentionLayer(\n num_blocks=num_prompt_blocks,\n dim=hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n layer_scale=-1.0,\n context_dim=hidden_dim,\n use_bias=False,\n )\n )\n\n for i, depth in enumerate(depths):\n current_dim = min(hidden_dim, mult * hidden_dim // int(2**i))\n next_dim = mult * hidden_dim // int(2 ** (i + 1))\n output_dim = max(next_dim, out_dim)\n self.process_features.append(\n nn.ConvTranspose2d(\n hidden_dim,\n current_dim,\n kernel_size=max(1, 2 * i),\n stride=max(1, 2 * i),\n padding=0,\n )\n )\n\n self.ups.append(\n ResUpsampleBil(\n current_dim,\n output_dim=output_dim,\n expansion=expansion,\n layer_scale=layer_scale,\n kernel_size=kernel_size,\n num_layers=depth,\n use_norm=use_norm,\n )\n )\n\n depth_mlp = nn.Identity()\n if i == len(depths) - 1:\n depth_mlp = nn.Sequential(\n nn.LayerNorm(next_dim), nn.Linear(next_dim, output_dim)\n )\n\n self.depth_mlp.append(depth_mlp)\n\n self.confidence_mlp = nn.Sequential(\n nn.LayerNorm(next_dim), nn.Linear(next_dim, output_dim)\n )\n\n self.to_depth_lr = nn.Conv2d(\n output_dim,\n output_dim // 2,\n kernel_size=3,\n padding=1,\n padding_mode=\"reflect\",\n )\n self.to_confidence_lr = nn.Conv2d(\n output_dim,\n output_dim // 2,\n kernel_size=3,\n padding=1,\n padding_mode=\"reflect\",\n )\n self.to_depth_hr = nn.Sequential(\n nn.Conv2d(\n output_dim // 2, 32, kernel_size=3, padding=1, padding_mode=\"reflect\"\n ),\n nn.LeakyReLU(),\n nn.Conv2d(32, 1, kernel_size=1),\n )\n self.to_confidence_hr = nn.Sequential(\n nn.Conv2d(\n output_dim // 2, 32, kernel_size=3, padding=1, padding_mode=\"reflect\"\n ),\n nn.LeakyReLU(),\n nn.Conv2d(32, 1, kernel_size=1),\n )\n\n def set_original_shapes(self, shapes: tuple[int, int]):\n self.original_shapes = shapes\n\n def set_shapes(self, shapes: tuple[int, int]):\n self.shapes = shapes\n\n def embed_rays(self, rays):\n rays_embedding = flat_interpolate(\n rays, old=self.original_shapes, new=self.shapes, antialias=True\n )\n rays_embedding = rays_embedding / torch.norm(\n rays_embedding, dim=-1, keepdim=True\n ).clip(min=1e-4)\n x, y, z = rays_embedding[..., 0], rays_embedding[..., 1], rays_embedding[..., 2]\n polar = torch.acos(z)\n x_clipped = x.abs().clip(min=1e-3) * (2 * (x >= 0).int() - 1)\n azimuth = torch.atan2(y, x_clipped)\n rays_embedding = torch.stack([polar, azimuth], dim=-1)\n rays_embedding = generate_fourier_features(\n rays_embedding,\n dim=self.hidden_dim,\n max_freq=max(self.shapes) // 2,\n use_log=True,\n cat_orig=False,\n )\n return rays_embedding\n\n def condition(self, feat, rays_embeddings):\n conditioned_features = [\n prompter(rearrange(feature, \"b h w c -> b (h w) c\"), rays_embeddings)\n for prompter, feature in zip(self.prompt_camera, feat)\n ]\n return conditioned_features\n\n def process(self, features_list, rays_embeddings):\n conditioned_features = self.condition(features_list, rays_embeddings)\n init_latents = self.to_latents(conditioned_features[0])\n init_latents = rearrange(\n init_latents, \"b (h w) c -> b c h w\", h=self.shapes[0], w=self.shapes[1]\n ).contiguous()\n conditioned_features = [\n rearrange(\n x, \"b (h w) c -> b c h w\", h=self.shapes[0], w=self.shapes[1]\n ).contiguous()\n for x in conditioned_features\n ]\n latents = init_latents\n\n out_features = []\n for i, up in enumerate(self.ups):\n latents = latents + self.process_features[i](conditioned_features[i + 1])\n latents = up(latents)\n out_features.append(latents)\n\n return out_features, init_latents\n\n def depth_proj(self, out_features):\n h_out, w_out = out_features[-1].shape[-2:]\n # aggregate output and project to depth\n for i, (layer, features) in enumerate(zip(self.depth_mlp, out_features)):\n out_depth_features = layer(features.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)\n out_depth_features = F.interpolate(\n out_depth_features,\n size=(h_out, w_out),\n mode=\"bilinear\",\n align_corners=True,\n )\n if i == len(self.depth_mlp) - 1:\n logdepth = out_depth_features\n\n logdepth = self.to_depth_lr(logdepth)\n logdepth = F.interpolate(\n logdepth, size=self.original_shapes, mode=\"bilinear\", align_corners=True\n )\n logdepth = self.to_depth_hr(logdepth)\n return logdepth\n\n def confidence_proj(self, out_features):\n highres_features = out_features[-1].permute(0, 2, 3, 1)\n confidence = self.confidence_mlp(highres_features).permute(0, 3, 1, 2)\n confidence = self.to_confidence_lr(confidence)\n confidence = F.interpolate(\n confidence, size=self.original_shapes, mode=\"bilinear\", align_corners=True\n )\n confidence = self.to_confidence_hr(confidence)\n return confidence\n\n def decode(self, out_features):\n logdepth = self.depth_proj(out_features)\n confidence = self.confidence_proj(out_features)\n return logdepth, confidence\n\n def forward(\n self,\n features: list[torch.Tensor],\n rays_hr: torch.Tensor,\n pos_embed,\n level_embed,\n ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n B = features[0].shape[0]\n\n rays_embeddings = self.embed_rays(rays_hr)\n features, proj_latents_16 = self.process(features, rays_embeddings)\n logdepth, logconf = self.decode(features)\n\n return logdepth, logconf, proj_latents_16\n\n\nclass Decoder(nn.Module):\n def __init__(\n self,\n config,\n ):\n super().__init__()\n self.build(config)\n self.apply(self._init_weights)\n self.test_gt_camera = False\n\n def _init_weights(self, m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=0.02)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Conv2d):\n trunc_normal_(m.weight, std=0.02)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.LayerNorm):\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n if m.weight is not None:\n nn.init.constant_(m.weight, 1.0)\n\n def run_camera(self, cls_tokens, features, pos_embed, original_shapes, rays_gt):\n H, W = original_shapes\n\n # camera layer\n intrinsics = self.camera_layer(\n features=features, cls_tokens=cls_tokens, pos_embed=pos_embed\n )\n B, N = intrinsics.shape\n device = intrinsics.device\n dtype = intrinsics.dtype\n\n id_coords = coords_grid(B, H, W, device=features.device, homogeneous=True)\n intrinsics_matrix_inverse = torch.eye(3, device=device, dtype=dtype).repeat(\n B, 1, 1\n )\n intrinsics_matrix_inverse[:, 0, 0] = 1.0 / intrinsics[:, 0]\n intrinsics_matrix_inverse[:, 1, 1] = 1.0 / intrinsics[:, 1]\n intrinsics_matrix_inverse[:, 0, 2] = -intrinsics[:, 2] / intrinsics[:, 0]\n intrinsics_matrix_inverse[:, 1, 2] = -intrinsics[:, 3] / intrinsics[:, 1]\n\n intrinsics_matrix = torch.eye(3, device=device, dtype=dtype).repeat(B, 1, 1)\n intrinsics_matrix[:, 0, 0] = intrinsics[:, 0]\n intrinsics_matrix[:, 1, 1] = intrinsics[:, 1]\n intrinsics_matrix[:, 0, 2] = intrinsics[:, 2]\n intrinsics_matrix[:, 1, 2] = intrinsics[:, 3]\n\n rays_pred = intrinsics_matrix_inverse @ id_coords.reshape(B, 3, -1)\n rays_pred = rays_pred.reshape(B, 3, H, W)\n rays_pred = rays_pred / torch.norm(rays_pred, dim=1, keepdim=True).clamp(\n min=1e-5\n )\n\n ### LEGACY CODE FOR TRAINING\n # if self.training and rays_gt is not None: \n # prob = -1.0 # 0.8 * (1 - tanh(self.steps / 100000)) + 0.2\n # where_use_gt_rays = torch.rand(B, 1, 1, device=device, dtype=dtype) < prob\n # where_use_gt_rays = where_use_gt_rays.int()\n # rays = rays_gt * where_use_gt_rays + rays_pred * (1 - where_use_gt_rays)\n\n rays = rays_pred if rays_gt is None else rays_gt\n rays = rearrange(rays, \"b c h w -> b (h w) c\")\n \n return intrinsics_matrix, rays\n\n def forward(\n self,\n inputs: dict[str, torch.Tensor],\n image_metas: list[dict[str, torch.Tensor]],\n ) -> dict[str, torch.Tensor]:\n B, C, H, W = inputs[\"image\"].shape\n device = inputs[\"image\"].device\n dtype = inputs[\"features\"][0].dtype\n\n # get features in b n d format\n common_shape = inputs[\"features\"][0].shape[1:3]\n\n # input shapes repeat shapes for each level, times the amount of the layers:\n features = self.input_adapter(inputs[\"features\"])\n\n # positional embeddings, spatial and level\n level_embed = self.level_embeds.repeat(\n B, common_shape[0] * common_shape[1], 1, 1\n )\n level_embed = rearrange(level_embed, \"b n l d -> b (n l) d\")\n dummy_tensor = torch.zeros(\n B, 1, common_shape[0], common_shape[1], device=device, requires_grad=False\n )\n pos_embed = self.pos_embed(dummy_tensor)\n pos_embed = rearrange(pos_embed, \"b c h w -> b (h w) c\").repeat(\n 1, self.num_resolutions, 1\n )\n\n # get cls tokens projections\n camera_tokens = inputs[\"tokens\"]\n camera_tokens = self.camera_token_adapter(camera_tokens)\n self.camera_layer.set_shapes((H, W))\n\n intrinsics, rays = self.run_camera(\n torch.cat(camera_tokens, dim=1),\n features=torch.stack(features, dim=-1).detach(),\n pos_embed=(pos_embed + level_embed).detach(),\n original_shapes=(H, W),\n rays_gt=inputs.get(\"rays\", None),\n )\n\n # run bulk of the model\n self.depth_layer.set_shapes(common_shape)\n self.depth_layer.set_original_shapes((H, W))\n logdepth, logconfidence, depth_features = self.depth_layer(\n features=features,\n rays_hr=rays,\n pos_embed=pos_embed,\n level_embed=level_embed,\n )\n\n return {\n \"radius\": torch.exp(logdepth.clip(min=-8.0, max=8.0) + 2.0),\n \"depth_features\": depth_features,\n \"confidence\": torch.exp(logconfidence.clip(min=-8.0, max=8.0)),\n \"intrinsics\": intrinsics,\n \"rays\": rays,\n }\n\n @torch.jit.ignore\n def no_weight_decay_keywords(self):\n return {\"latents_pos\", \"level_embeds\"}\n\n def build(self, config):\n input_dims = config[\"model\"][\"pixel_encoder\"][\"embed_dims\"]\n hidden_dim = config[\"model\"][\"pixel_decoder\"][\"hidden_dim\"]\n expansion = config[\"model\"][\"expansion\"]\n num_heads = config[\"model\"][\"num_heads\"]\n dropout = config[\"model\"][\"pixel_decoder\"][\"dropout\"]\n depths_encoder = config[\"model\"][\"pixel_encoder\"][\"depths\"]\n layer_scale = config[\"model\"][\"layer_scale\"]\n\n depth = config[\"model\"][\"pixel_decoder\"][\"depths\"]\n self.downsample = 4\n depths_encoder = config[\"model\"][\"pixel_encoder\"][\"depths\"]\n\n self.num_resolutions = len(depths_encoder)\n self.test_fixed_camera = False\n\n out_dim = config[\"model\"][\"pixel_decoder\"][\"out_dim\"]\n kernel_size = config[\"model\"][\"pixel_decoder\"].get(\"kernel_size\", 7)\n\n self.slices_encoder = list(zip([d - 1 for d in depths_encoder], depths_encoder))\n input_dims = [input_dims[d - 1] for d in depths_encoder]\n\n # # adapt from encoder features, just project\n camera_dims = input_dims\n self.input_adapter = ListAdapter(input_dims, hidden_dim)\n self.camera_token_adapter = ListAdapter(camera_dims, hidden_dim)\n\n # # camera layer\n self.camera_layer = CameraHead(\n hidden_dim=hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n layer_scale=layer_scale,\n )\n\n self.depth_layer = DepthHead(\n hidden_dim=hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n depths=depth,\n dropout=dropout,\n camera_dim=96,\n num_resolutions=self.num_resolutions,\n layer_scale=layer_scale,\n out_dim=out_dim,\n kernel_size=kernel_size,\n num_prompt_blocks=1,\n use_norm=False,\n )\n self.pos_embed = PositionEmbeddingSine(hidden_dim // 2, normalize=True)\n self.level_embeds = nn.Parameter(\n orthonormal_init(len(input_dims), hidden_dim).reshape(\n 1, 1, len(input_dims), hidden_dim\n ),\n requires_grad=False,\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv2/decoder_old.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom timm.models.layers import trunc_normal_\n\nfrom unidepth.layers import (MLP, AttentionBlock, ConvUpsampleShuffleResidual,\n NystromBlock, PositionEmbeddingSine)\nfrom unidepth.utils.geometric import flat_interpolate, generate_rays\nfrom unidepth.utils.positional_embedding import generate_fourier_features\n\n\nclass ListAdapter(nn.Module):\n def __init__(self, input_dims: list[int], hidden_dim: int):\n super().__init__()\n self.input_adapters = nn.ModuleList([])\n self.num_chunks = len(input_dims)\n self.checkpoint = True\n for input_dim in input_dims:\n self.input_adapters.append(\n nn.Sequential(\n nn.LayerNorm(input_dim), nn.Linear(input_dim, hidden_dim), nn.GELU()\n )\n )\n\n def forward(self, x: torch.Tensor, splits: torch.Tensor) -> torch.Tensor:\n xs = torch.split(x, splits.int().tolist(), dim=-1)\n xs = [adapter(x) for x, adapter in zip(xs, self.input_adapters)]\n return torch.cat(xs, dim=-1)\n\n\nclass CameraHead(nn.Module):\n def __init__(\n self,\n hidden_dim: int,\n num_heads: int = 8,\n expansion: int = 4,\n dropout: float = 0.0,\n **kwargs,\n ):\n super().__init__()\n self.aggregate1 = AttentionBlock(\n hidden_dim, num_heads=1, expansion=expansion, dropout=dropout\n )\n self.aggregate2 = AttentionBlock(\n hidden_dim, num_heads=1, expansion=expansion, dropout=dropout\n )\n self.latents_pos = nn.Parameter(\n torch.randn(1, 4, hidden_dim), requires_grad=True\n )\n self.in_features = MLP(hidden_dim, expansion=2, dropout=dropout)\n self.project_cls = MLP(hidden_dim, dropout=dropout)\n self.out = MLP(hidden_dim, expansion=2, dropout=0.0, output_dim=1)\n\n def fill_intrinsics(self, x):\n camera_intrinsics = torch.zeros(\n x.shape[0], 3, 3, device=x.device, requires_grad=False\n )\n camera_intrinsics[:, 0, 0] = x[:, 0].exp()\n camera_intrinsics[:, 1, 1] = x[:, 1].exp()\n camera_intrinsics[:, 0, 2] = x[:, 2].sigmoid()\n camera_intrinsics[:, 1, 2] = x[:, 3].sigmoid()\n camera_intrinsics[:, 2, 2] = 1.0\n return camera_intrinsics\n\n def forward(self, features, cls_tokens, pos_embed) -> torch.Tensor:\n features = features.unbind(dim=-1)\n cls_tokens = self.project_cls(cls_tokens)\n latents_pos = self.latents_pos.expand(cls_tokens.shape[0], -1, -1)\n features = self.in_features(torch.cat(features, dim=1) + pos_embed)\n features = torch.cat((features, cls_tokens), dim=1)\n cls_tokens = self.aggregate1(\n cls_tokens, context=features, pos_embed=latents_pos\n )\n cls_tokens = self.aggregate2(\n cls_tokens, context=features, pos_embed=latents_pos\n )\n\n # project to intrinsics\n x = self.out(cls_tokens).squeeze(-1)\n camera_intrinsics = self.fill_intrinsics(x)\n\n return camera_intrinsics\n\n def set_shapes(self, shapes: tuple[int, int]):\n self.shapes = shapes\n\n\nclass GlobalHead(nn.Module):\n def __init__(\n self,\n hidden_dim: int,\n camera_dim: int,\n expansion: int = 4,\n dropout: float = 0.0,\n **kwargs,\n ):\n super().__init__()\n self.camera_dim = camera_dim\n self.in_features = nn.Linear(hidden_dim, hidden_dim)\n self.project_rays = nn.Linear(camera_dim + 3, hidden_dim)\n self.aggregate1 = AttentionBlock(\n hidden_dim, num_heads=1, expansion=expansion, dropout=dropout\n )\n self.aggregate2 = AttentionBlock(\n hidden_dim, num_heads=1, expansion=expansion, dropout=dropout\n )\n self.project_cls = MLP(hidden_dim, dropout=dropout)\n self.out = MLP(hidden_dim, expansion=2, dropout=0.0, output_dim=1)\n\n def embed_rays(self, rays, shapes):\n rays_embedding = flat_interpolate(rays, old=self.original_shapes, new=shapes)\n rays_embedding = F.normalize(rays_embedding, dim=-1)\n rays_embedding = generate_fourier_features(\n rays_embedding,\n dim=self.camera_dim,\n max_freq=max(shapes) // 2,\n use_log=True,\n cat_orig=True,\n )\n return rays_embedding\n\n def set_original_shapes(self, shapes: tuple[int, int]):\n self.original_shapes = shapes\n\n def set_shapes(self, shapes: tuple[int, int]):\n self.shapes = shapes\n\n def get_scaleshift(self, x):\n scale, shift = torch.chunk(x, 2, dim=1)\n scale = scale.exp().reshape(-1, 1, 1, 1)\n shift = shift.reshape(-1, 1, 1, 1)\n return scale, shift\n\n def forward(self, features, cls_tokens, rays) -> torch.Tensor:\n features = features.unbind(dim=-1)\n cls_tokens = self.project_cls(cls_tokens)\n rays_embedding = self.project_rays(self.embed_rays(rays, self.shapes))\n rays_embedding = rays_embedding.repeat(1, len(features), 1)\n features = self.in_features(torch.cat(features, dim=1) + rays_embedding)\n features = torch.cat((features, cls_tokens), dim=1)\n cls_tokens = self.aggregate1(cls_tokens, context=features)\n cls_tokens = self.aggregate2(cls_tokens, context=features)\n x = self.out(cls_tokens).squeeze(-1)\n scale, shift = self.get_scaleshift(x)\n return scale, shift\n\n\nclass DepthHead(nn.Module):\n def __init__(\n self,\n hidden_dim: int,\n num_heads: int = 8,\n expansion: int = 4,\n depths: int | list[int] = 4,\n checkpoint: bool = True,\n camera_dim: int = 256,\n num_resolutions: int = 4,\n dropout: float = 0.0,\n **kwargs,\n ) -> None:\n super().__init__()\n self.checkpoint = checkpoint\n self.camera_dim = camera_dim\n self.skip_depth = False\n\n self.to_latents = MLP(hidden_dim, expansion=2, dropout=dropout)\n self.features_channel_cat = nn.Linear(hidden_dim * num_resolutions, hidden_dim)\n self.aggregate_16 = AttentionBlock(\n hidden_dim,\n num_heads=1,\n expansion=expansion,\n dropout=dropout,\n context_dim=hidden_dim,\n )\n self.prompt_camera = AttentionBlock(\n hidden_dim,\n num_heads=1,\n expansion=expansion,\n dropout=dropout,\n context_dim=hidden_dim,\n )\n\n self.rays_layers = nn.ModuleList([])\n self.ups = nn.ModuleList([])\n self.process_layers = nn.ModuleList([])\n self.norms, self.out_layers = nn.ModuleList([]), nn.ModuleList([])\n self.depth_mlp, self.confidence_mlp = nn.ModuleList([]), nn.ModuleList([])\n for i, depth in enumerate(depths):\n blk_lst = nn.ModuleList([])\n for _ in range(depth):\n blk_lst.append(\n NystromBlock(\n hidden_dim // int(2**i),\n num_heads=num_heads // int(2**i),\n expansion=expansion,\n dropout=dropout,\n )\n )\n self.process_layers.append(blk_lst)\n self.rays_layers.append(nn.Linear(camera_dim + 3, hidden_dim // int(2**i)))\n self.ups.append(\n ConvUpsampleShuffleResidual(\n hidden_dim // int(2**i),\n expansion=expansion,\n kernel_size=7,\n num_layers=2,\n )\n )\n self.depth_mlp.append(\n MLP(\n input_dim=hidden_dim // int(2 ** (i + 1)),\n output_dim=16,\n expansion=1,\n )\n )\n self.confidence_mlp.append(\n MLP(\n input_dim=hidden_dim // int(2 ** (i + 1)),\n output_dim=16,\n expansion=1,\n )\n )\n self.to_depth = nn.Conv2d(\n 16 * len(depths), 1, 7, padding=3, padding_mode=\"reflect\"\n )\n self.to_confidence = nn.Conv2d(\n 16 * len(depths), 1, 7, padding=3, padding_mode=\"reflect\"\n )\n\n def set_original_shapes(self, shapes: tuple[int, int]):\n self.original_shapes = shapes\n\n def set_shapes(self, shapes: tuple[int, int]):\n self.shapes = shapes\n\n def embed_rays(self, rays, shapes):\n rays_embedding = flat_interpolate(rays, old=self.original_shapes, new=shapes)\n rays_embedding = F.normalize(rays_embedding, dim=-1)\n rays_embedding = generate_fourier_features(\n rays_embedding,\n dim=self.camera_dim,\n max_freq=max(shapes) // 2,\n use_log=True,\n cat_orig=True,\n )\n return rays_embedding\n\n def project_rays(self, rays, shapes):\n embedded_rays = []\n for i, layer in enumerate(self.rays_layers):\n embedded_rays.append(\n layer(self.embed_rays(rays, [(2**i) * x for x in shapes]))\n )\n return embedded_rays\n\n def decode_depth(self, latents_16, rays, shapes):\n latents = latents_16\n out_features, depths, confidences = [], [], []\n for i, (up, layers, rays_embedding) in enumerate(\n zip(self.ups, self.process_layers, rays)\n ):\n for layer in layers:\n latents = layer(latents, pos_embed=rays_embedding)\n latents = up(\n rearrange(\n latents + rays_embedding,\n \"b (h w) c -> b c h w\",\n h=shapes[0] * int(2**i),\n w=shapes[1] * int(2**i),\n ).contiguous()\n )\n out = rearrange(\n latents,\n \"b (h w) c -> b h w c\",\n h=shapes[0] * int(2 ** (1 + i)),\n w=shapes[1] * int(2 ** (1 + i)),\n )\n out_features.append(out)\n\n # aggregate output and project to depth\n for i, (layer, features) in enumerate(\n zip(self.depth_mlp[::-1], out_features[::-1])\n ):\n out_depth_features = layer(features).permute(0, 3, 1, 2)\n out_depth_features = F.interpolate(\n out_depth_features, size=self.original_shapes, mode=\"bilinear\"\n )\n depths.append(out_depth_features)\n logdepth = self.to_depth(torch.cat(depths, dim=1))\n\n # aggregate output and project to confidences\n for i, (layer, features) in enumerate(\n zip(self.confidence_mlp[::-1], out_features[::-1])\n ):\n out_conf_features = layer(features).permute(0, 3, 1, 2)\n out_conf_features = F.interpolate(\n out_conf_features, size=self.original_shapes, mode=\"bilinear\"\n )\n confidences.append(out_conf_features)\n confidence = self.to_confidence(torch.cat(confidences, dim=1))\n\n # apply sigmoid ot get conf in [0, 1]\n confidence = torch.sigmoid(confidence)\n\n return logdepth, confidence\n\n def init_latents(self, features, shapes):\n # Generate latents with init as pooled features\n features_channels = torch.cat(features, dim=-1)\n features_16 = self.features_channel_cat(features_channels)\n latents_16 = features_16 + self.to_latents(\n flat_interpolate(features_16, old=self.shapes, new=shapes, antialias=False)\n )\n return latents_16\n\n def forward(\n self, features: torch.Tensor, rays_hr: torch.Tensor, pos_embed, level_embed\n ) -> torch.Tensor:\n B = features.shape[0]\n features = features.unbind(dim=-1)\n shapes = self.shapes\n\n # camera_embedding\n rays_embeddings = self.project_rays(rays_hr, shapes)\n\n # Init latents\n init_latents_16 = self.init_latents(features, shapes)\n\n # Aggregate features: F -> D\n latents_16 = self.aggregate_16(\n init_latents_16,\n context=torch.cat(features, dim=1),\n pos_embed_context=pos_embed + level_embed,\n )\n\n # Aggregate camera: D -> D|E\n latents_16 = self.prompt_camera(latents_16, context=rays_embeddings[0])\n\n # Decode depth\n logdepth, confidence = self.decode_depth(latents_16, rays_embeddings, shapes)\n\n return logdepth, confidence, latents_16\n\n\nclass Decoder(nn.Module):\n def __init__(\n self,\n config,\n ):\n super().__init__()\n self.build(config)\n self.apply(self._init_weights)\n\n def _init_weights(self, m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=0.02)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Conv2d):\n trunc_normal_(m.weight, std=0.02)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.LayerNorm):\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n if m.weight is not None:\n nn.init.constant_(m.weight, 1.0)\n\n def get_adapted_features(self, features_flat, splits):\n features_flat_cat = torch.cat(features_flat, dim=-1)\n features_projected = self.input_adapter(\n features_flat_cat, splits\n ) # list [b hw c] shapes\n features = torch.chunk(features_projected, splits.shape[0], dim=-1)\n return features\n\n def run_camera(self, cls_tokens, features, pos_embed, original_shapes, rays_gt):\n # get cls tokens projections\n cls_tokens_splits = torch.tensor(\n [x.shape[-1] for x in cls_tokens],\n device=features.device,\n requires_grad=False,\n dtype=features.dtype,\n )\n cls_tokens = torch.cat(cls_tokens, dim=-1)\n cls_tokens = self.camera_token_adapter(cls_tokens, cls_tokens_splits)\n cls_tokens = torch.cat(\n torch.chunk(cls_tokens, cls_tokens_splits.shape[0], dim=-1), dim=1\n )\n\n # camera layer\n intrinsics = self.camera_layer(\n features=features, cls_tokens=cls_tokens, pos_embed=pos_embed\n )\n intrinsics[:, 0, 0] = max(original_shapes) / 2 * intrinsics[:, 0, 0]\n intrinsics[:, 1, 1] = max(original_shapes) / 2 * intrinsics[:, 1, 1]\n intrinsics[:, 0, 2] = intrinsics[:, 0, 2] * original_shapes[1]\n intrinsics[:, 1, 2] = intrinsics[:, 1, 2] * original_shapes[0]\n\n rays = (\n rays_gt\n if rays_gt is not None\n else generate_rays(intrinsics, original_shapes)[0]\n )\n return intrinsics, rays\n\n def run_global(self, cls_tokens, features, rays):\n # get cls tokens projections\n cls_tokens_splits = torch.tensor(\n [x.shape[-1] for x in cls_tokens],\n device=features.device,\n requires_grad=False,\n dtype=torch.float32,\n )\n cls_tokens = torch.cat(cls_tokens, dim=-1)\n cls_tokens = self.global_token_adapter(cls_tokens, cls_tokens_splits)\n cls_tokens = torch.cat(\n torch.chunk(cls_tokens, cls_tokens_splits.shape[0], dim=-1), dim=1\n )\n\n scale, shift = self.global_layer(\n features=features, rays=rays, cls_tokens=cls_tokens\n )\n\n return scale, shift\n\n def forward(self, inputs, image_metas) -> torch.Tensor:\n B, C, H, W = inputs[\"image\"].shape\n device = inputs[\"image\"].device\n dtype = inputs[\"image\"].dtype\n\n # get features in b n d format\n # level shapes, the shape per level, for swin like [[128, 128], [64, 64],...], for vit [[32,32]] -> mult times resolutions\n level_shapes = sorted(\n list(set([tuple([x.shape[1], x.shape[2]]) for x in inputs[\"features\"]]))\n )[::-1]\n if len(level_shapes) == 1:\n level_shapes = level_shapes * self.num_resolutions\n input_shapes = [\n level_shapes[i]\n for i, (start, end) in enumerate(self.slices_encoder)\n for _ in range(end - start)\n ]\n common_shape = level_shapes[-2]\n\n # input shapes repeat shapes for each level, times the amount of the layers:\n features_flat = [\n flat_interpolate(\n rearrange(x, \"b h w c -> b (h w) c\"), old=input_shape, new=common_shape\n )\n for x, input_shape in zip(inputs[\"features\"], input_shapes)\n ]\n features_splits = torch.tensor(\n [x.shape[-1] for x in features_flat],\n device=device,\n requires_grad=False,\n dtype=torch.float32,\n )\n features = self.get_adapted_features(features_flat, features_splits)\n features = torch.stack(features, dim=-1)\n\n # positional embeddings, spatial and level\n level_embed = torch.cat(\n [\n self.level_embed_layer(self.level_embeds)[i : i + 1]\n .unsqueeze(0)\n .repeat(B, common_shape[0] * common_shape[1], 1)\n for i in range(self.num_resolutions)\n ],\n dim=1,\n )\n dummy_tensor = torch.zeros(\n B, 1, common_shape[0], common_shape[1], device=device, requires_grad=False\n )\n pos_embed = self.pos_embed(dummy_tensor)\n pos_embed = rearrange(pos_embed, \"b c h w -> b (h w) c\").repeat(\n 1, self.num_resolutions, 1\n )\n\n self.camera_layer.set_shapes(common_shape)\n intrinsics, rays = self.run_camera(\n inputs[\"camera_tokens\"],\n features=features,\n pos_embed=pos_embed + level_embed,\n original_shapes=(H, W),\n rays_gt=inputs.get(\"rays\"),\n )\n\n self.global_layer.set_shapes(common_shape)\n self.global_layer.set_original_shapes((H, W))\n scale, shift = self.run_global(\n inputs[\"global_tokens\"], features=features, rays=rays\n )\n\n # run bulk of the model\n self.depth_layer.set_shapes(common_shape)\n self.depth_layer.set_original_shapes((H, W))\n logdepth, confidence, depth_features = self.depth_layer(\n features=features,\n rays_hr=rays,\n pos_embed=pos_embed,\n level_embed=level_embed,\n )\n logdepth = logdepth.to(torch.float32, non_blocking=True)\n\n # norm in log space, why performs better?\n shapes = [int(x) for x in logdepth.shape[-2:]]\n depth_normalized = F.layer_norm(logdepth, shapes).exp()\n\n depth = (\n depth_normalized + shift\n ) * scale # shift is scale invariant if we do (x + mu) * sigma\n depth = F.softplus(depth, beta=10.0).to(dtype, non_blocking=True)\n\n outputs = {\n \"depth\": depth,\n \"confidence\": confidence,\n \"depth_features\": depth_features,\n \"K\": intrinsics,\n }\n return outputs\n\n @torch.jit.ignore\n def no_weight_decay_keywords(self):\n return {\"latents_pos\", \"level_embeds\"}\n\n def build(self, config):\n input_dims = config[\"model\"][\"pixel_encoder\"][\"embed_dims\"]\n hidden_dim = config[\"model\"][\"pixel_decoder\"][\"hidden_dim\"]\n expansion = config[\"model\"][\"expansion\"]\n num_heads = config[\"model\"][\"num_heads\"]\n dropout = config[\"model\"][\"pixel_decoder\"][\"dropout\"]\n depths_encoder = config[\"model\"][\"pixel_encoder\"][\"depths\"]\n depth = config[\"model\"][\"pixel_decoder\"][\"depths\"]\n depths_encoder = config[\"model\"][\"pixel_encoder\"][\"depths\"]\n self.downsample = 4\n self.num_resolutions = len(depths_encoder)\n\n self.slices_encoder = list(zip([d - 1 for d in depths_encoder], depths_encoder))\n cls_token_input_dims = [input_dims[i] for i in [-1, -2, -3, -4]]\n input_dims = [input_dims[d - 1] for d in depths_encoder]\n\n # # camera layer\n self.camera_layer = CameraHead(\n hidden_dim=hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n )\n\n # # scale shift layer\n self.global_layer = GlobalHead(\n hidden_dim=hidden_dim,\n camera_dim=96,\n num_heads=num_heads,\n expansion=expansion,\n dropout=dropout,\n )\n\n # # adapt from encoder features, just project\n self.input_adapter = ListAdapter(input_dims, hidden_dim)\n self.camera_token_adapter = ListAdapter(cls_token_input_dims, hidden_dim)\n self.global_token_adapter = ListAdapter(cls_token_input_dims[:2], hidden_dim)\n\n self.depth_layer = DepthHead(\n hidden_dim=hidden_dim,\n num_heads=num_heads,\n expansion=expansion,\n depths=depth,\n dropout=dropout,\n camera_dim=96,\n num_resolutions=self.num_resolutions,\n )\n\n self.pos_embed = PositionEmbeddingSine(hidden_dim // 2, normalize=True)\n self.level_embeds = nn.Parameter(\n torch.randn(len(input_dims), hidden_dim), requires_grad=True\n )\n self.level_embed_layer = nn.Sequential(\n nn.Linear(hidden_dim, hidden_dim),\n nn.GELU(),\n nn.Linear(hidden_dim, hidden_dim),\n nn.LayerNorm(hidden_dim),\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv2/export.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport argparse\nimport json\nimport os\nfrom math import ceil\n\nimport huggingface_hub\nimport torch.nn.functional as F\nimport torch.onnx\n\nfrom unidepth.models.unidepthv2 import UniDepthV2\n\n\nclass UniDepthV2ONNX(UniDepthV2):\n def __init__(\n self,\n config,\n eps: float = 1e-6,\n **kwargs,\n ):\n super().__init__(config, eps)\n\n def forward(self, rgbs):\n B, _, H, W = rgbs.shape\n features, tokens = self.pixel_encoder(rgbs)\n\n inputs = {}\n inputs[\"image\"] = rgbs\n inputs[\"features\"] = [\n self.stacking_fn(features[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n inputs[\"tokens\"] = [\n self.stacking_fn(tokens[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n outputs = self.pixel_decoder(inputs, [])\n outputs[\"rays\"] = outputs[\"rays\"].permute(0, 2, 1).reshape(B, 3, H, W)\n pts_3d = outputs[\"rays\"] * outputs[\"radius\"]\n\n return pts_3d, outputs[\"confidence\"], outputs[\"intrinsics\"]\n\n\nclass UniDepthV2ONNXcam(UniDepthV2):\n def __init__(\n self,\n config,\n eps: float = 1e-6,\n **kwargs,\n ):\n super().__init__(config, eps)\n\n def forward(self, rgbs, rays):\n B, _, H, W = rgbs.shape\n features, tokens = self.pixel_encoder(rgbs)\n\n inputs = {}\n inputs[\"image\"] = rgbs\n inputs[\"rays\"] = rays\n inputs[\"features\"] = [\n self.stacking_fn(features[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n inputs[\"tokens\"] = [\n self.stacking_fn(tokens[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n outputs = self.pixel_decoder(inputs, [])\n outputs[\"rays\"] = outputs[\"rays\"].permute(0, 2, 1).reshape(B, 3, H, W)\n pts_3d = outputs[\"rays\"] * outputs[\"radius\"]\n\n return pts_3d, outputs[\"confidence\"], outputs[\"intrinsics\"]\n\n\n\ndef export(model, path, shape=(462, 630), with_camera=False):\n model.eval()\n image = torch.rand(1, 3, *shape)\n dynamic_axes_in = {\"rgbs\": {0: \"batch\"}}\n inputs = [image]\n if with_camera:\n rays = torch.rand(1, 3, *shape)\n inputs.append(rays)\n dynamic_axes_in[\"rays\"] = {0: \"batch\"}\n\n dynamic_axes_out = {\n \"pts_3d\": {0: \"batch\"},\n \"confidence\": {0: \"batch\"},\n \"intrinsics\": {0: \"batch\"},\n }\n torch.onnx.export(\n model,\n tuple(inputs),\n path,\n input_names=list(dynamic_axes_in.keys()),\n output_names=list(dynamic_axes_out.keys()),\n opset_version=14,\n dynamic_axes={**dynamic_axes_in, **dynamic_axes_out},\n )\n print(f\"Model exported to {path}\")\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser(description=\"Export UniDepthV2 model to ONNX\")\n parser.add_argument(\n \"--version\", type=str, default=\"v2\", choices=[\"v2\"], help=\"UniDepth version\"\n )\n parser.add_argument(\n \"--backbone\",\n type=str,\n default=\"vitl\",\n choices=[\"vits\", \"vitb\", \"vitl\"],\n help=\"Backbone model\",\n )\n parser.add_argument(\n \"--shape\",\n type=int,\n nargs=2,\n default=(462, 630),\n help=\"Input shape. No dyamic shape supported!\",\n )\n parser.add_argument(\n \"--output-path\", type=str, default=\"unidepthv2.onnx\", help=\"Output ONNX file\"\n )\n parser.add_argument(\n \"--with-camera\",\n action=\"store_true\",\n help=\"Export model that expects GT camera as unprojected rays at inference\",\n )\n args = parser.parse_args()\n\n version = args.version\n backbone = args.backbone\n shape = args.shape\n output_path = args.output_path\n with_camera = args.with_camera\n\n # force shape to be multiple of 14\n shape_rounded = [14 * ceil(x // 14 - 0.5) for x in shape]\n if list(shape) != list(shape_rounded):\n print(f\"Shape {shape} is not multiple of 14. Rounding to {shape_rounded}\")\n shape = shape_rounded\n\n # assumes command is from root of repo\n with open(os.path.join(\"configs\", f\"config_{version}_{backbone}14.json\")) as f:\n config = json.load(f)\n\n # tell DINO not to use efficient attention: not exportable\n config[\"training\"][\"export\"] = True\n\n model = UniDepthV2ONNX(config) if not with_camera else UniDepthV2ONNXcam(config)\n path = huggingface_hub.hf_hub_download(\n repo_id=f\"lpiccinelli/unidepth-{version}-{backbone}14\",\n filename=f\"pytorch_model.bin\",\n repo_type=\"model\",\n )\n info = model.load_state_dict(torch.load(path), strict=False)\n print(f\"UniDepth_{version}_{backbone} is loaded with:\")\n print(f\"\\t missing keys: {info.missing_keys}\")\n print(f\"\\t additional keys: {info.unexpected_keys}\")\n\n export(\n model=model,\n path=os.path.join(os.environ.get(\"TMPDIR\", \".\"), output_path),\n shape=shape,\n with_camera=with_camera,\n )\n\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv2/unidepthv2.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport importlib\nfrom copy import deepcopy\nfrom math import ceil\nimport warnings\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms.v2.functional as TF\nfrom einops import rearrange\nfrom huggingface_hub import PyTorchModelHubMixin\n\nfrom unidepth.models.unidepthv2.decoder import Decoder\nfrom unidepth.utils.camera import BatchCamera, Camera, Pinhole\nfrom unidepth.utils.constants import (IMAGENET_DATASET_MEAN,\n IMAGENET_DATASET_STD)\nfrom unidepth.utils.distributed import is_main_process\nfrom unidepth.utils.misc import (first_stack, get_params, last_stack, match_gt,\n match_intrinsics, max_stack, mean_stack,\n softmax_stack)\n\nSTACKING_FNS = {\n \"max\": max_stack,\n \"mean\": mean_stack,\n \"first\": first_stack,\n \"last\": last_stack,\n \"softmax\": softmax_stack,\n}\n\n\ndef get_paddings(original_shape, aspect_ratio_range):\n # Original dimensions\n H_ori, W_ori = original_shape\n orig_aspect_ratio = W_ori / H_ori\n\n # Determine the closest aspect ratio within the range\n min_ratio, max_ratio = aspect_ratio_range\n target_aspect_ratio = min(max_ratio, max(min_ratio, orig_aspect_ratio))\n\n if orig_aspect_ratio > target_aspect_ratio: # Too wide\n W_new = W_ori\n H_new = int(W_ori / target_aspect_ratio)\n pad_top = (H_new - H_ori) // 2\n pad_bottom = H_new - H_ori - pad_top\n pad_left, pad_right = 0, 0\n else: # Too tall\n H_new = H_ori\n W_new = int(H_ori * target_aspect_ratio)\n pad_left = (W_new - W_ori) // 2\n pad_right = W_new - W_ori - pad_left\n pad_top, pad_bottom = 0, 0\n\n return (pad_left, pad_right, pad_top, pad_bottom), (H_new, W_new)\n\n\ndef get_resize_factor(original_shape, pixels_range, shape_multiplier=14):\n # Original dimensions\n H_ori, W_ori = original_shape\n n_pixels_ori = W_ori * H_ori\n\n # Determine the closest number of pixels within the range\n min_pixels, max_pixels = pixels_range\n target_pixels = min(max_pixels, max(min_pixels, n_pixels_ori))\n\n # Calculate the resize factor\n resize_factor = (target_pixels / n_pixels_ori) ** 0.5\n new_width = int(W_ori * resize_factor)\n new_height = int(H_ori * resize_factor)\n new_height = ceil(new_height / shape_multiplier) * shape_multiplier\n new_width = ceil(new_width / shape_multiplier) * shape_multiplier\n\n return resize_factor, (new_height, new_width)\n\n\ndef _postprocess(tensor, shapes, paddings, interpolation_mode=\"bilinear\"):\n # interpolate to original size\n tensor = F.interpolate(\n tensor, size=shapes, mode=interpolation_mode, align_corners=False\n )\n\n # remove paddings\n pad1_l, pad1_r, pad1_t, pad1_b = paddings\n tensor = tensor[..., pad1_t : shapes[0] - pad1_b, pad1_l : shapes[1] - pad1_r]\n return tensor\n\n\ndef _postprocess_intrinsics(K, resize_factors, paddings):\n batch_size = K.shape[0]\n K_new = K.clone()\n\n for i in range(batch_size):\n scale = resize_factors[i]\n pad_l, _, pad_t, _ = paddings[i]\n\n K_new[i, 0, 0] /= scale # fx\n K_new[i, 1, 1] /= scale # fy\n K_new[i, 0, 2] /= scale # cx\n K_new[i, 1, 2] /= scale # cy\n\n K_new[i, 0, 2] -= pad_l # cx\n K_new[i, 1, 2] -= pad_t # cy\n\n return K_new\n\n\nclass UniDepthV2(\n nn.Module,\n PyTorchModelHubMixin,\n library_name=\"UniDepth\",\n repo_url=\"https://github.com/lpiccinelli-eth/UniDepth\",\n tags=[\"monocular-metric-depth-estimation\"],\n):\n def __init__(\n self,\n config,\n eps: float = 1e-6,\n **kwargs,\n ):\n super().__init__()\n self.eps = eps\n self.build(config)\n self.build_losses(config)\n\n def forward_train(self, inputs, image_metas):\n inputs, outputs = self.encode_decode(inputs, image_metas)\n losses = self.compute_losses(outputs, inputs, image_metas)\n return outputs, losses\n\n def forward_test(self, inputs, image_metas):\n inputs, outputs = self.encode_decode(inputs, image_metas)\n depth_gt = inputs[\"depth\"]\n test_outputs = {}\n test_outputs[\"depth\"] = match_gt(\n outputs[\"depth\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"points\"] = match_gt(\n outputs[\"points\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"confidence\"] = match_gt(\n outputs[\"confidence\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"rays\"] = match_gt(\n outputs[\"rays\"], depth_gt, padding1=inputs[\"paddings\"], padding2=None\n )\n test_outputs[\"rays\"] = outputs[\"rays\"] / torch.norm(\n outputs[\"rays\"], dim=1, keepdim=True\n ).clip(min=1e-5)\n test_outputs[\"intrinsics\"] = match_intrinsics(\n outputs[\"intrinsics\"],\n inputs[\"image\"],\n depth_gt,\n padding1=inputs[\"paddings\"],\n padding2=None,\n )\n return test_outputs\n\n def forward(self, inputs, image_metas):\n if self.training:\n return self.forward_train(inputs, image_metas)\n else:\n return self.forward_test(inputs, image_metas)\n\n def compute_losses(self, outputs, inputs, image_metas):\n B, _, H, W = inputs[\"image\"].shape\n losses = {\"opt\": {}, \"stat\": {}}\n losses_to_be_computed = list(self.losses.keys())\n\n # depth loss\n si = torch.tensor(\n [x.get(\"si\", False) for x in image_metas], device=self.device\n ).reshape(B)\n loss = self.losses[\"depth\"]\n depth_losses = loss(\n outputs[\"depth\"],\n target=inputs[\"depth\"],\n mask=inputs[\"depth_mask\"].clone(),\n si=si,\n )\n losses[\"opt\"][loss.name] = loss.weight * depth_losses.mean()\n losses_to_be_computed.remove(\"depth\")\n\n # camera loss, here we apply to rays for simplicity\n # in the original training was on angles\n # however, we saw no difference (see supplementary)\n loss = self.losses[\"camera\"]\n camera_losses = loss(outputs[\"rays\"], target=inputs[\"rays\"])\n losses[\"opt\"][loss.name] = loss.weight * camera_losses.mean()\n losses_to_be_computed.remove(\"camera\")\n\n # invariance loss on output depth\n flips = torch.tensor(\n [x.get(\"flip\", False) for x in image_metas], device=self.device\n ).reshape(B)\n loss = self.losses[\"invariance\"]\n invariance_losses = loss(\n outputs[\"depth\"],\n intrinsics=inputs[\"camera\"].K,\n mask=inputs[\"depth_mask\"],\n flips=flips,\n downsample_ratio=1,\n )\n losses[\"opt\"][loss.name] = loss.weight * invariance_losses.mean()\n losses_to_be_computed.remove(\"invariance\")\n\n # edge guided ssi\n loss = self.losses[\"ssi\"]\n ssi_losses = loss(\n outputs[\"depth\"],\n target=inputs[\"depth\"],\n mask=inputs[\"depth_mask\"].clone(),\n image=inputs[\"image\"],\n validity_mask=inputs[\"validity_mask\"],\n )\n losses[\"opt\"][loss.name] = loss.weight * ssi_losses.mean()\n losses_to_be_computed.remove(\"ssi\")\n\n # remaining losses, we expect no more losses to be computed\n loss = self.losses[\"confidence\"]\n conf_losses = loss(\n outputs[\"confidence\"].log(),\n target_gt=inputs[\"depth\"],\n target_pred=outputs[\"depth\"],\n mask=inputs[\"depth_mask\"].clone(),\n )\n losses[\"opt\"][loss.name + \"_conf\"] = loss.weight * conf_losses.mean()\n losses_to_be_computed.remove(\"confidence\")\n\n assert (\n not losses_to_be_computed\n ), f\"Losses {losses_to_be_computed} not computed, revise `compute_loss` method\"\n\n return losses\n\n @torch.no_grad()\n @torch.autocast(device_type=\"cuda\", enabled=True, dtype=torch.float16)\n def infer(\n self,\n rgb: torch.Tensor,\n camera: torch.Tensor | Camera | None = None,\n normalize=True,\n ):\n ratio_bounds = self.shape_constraints[\"ratio_bounds\"]\n pixels_bounds = [\n self.shape_constraints[\"pixels_min\"],\n self.shape_constraints[\"pixels_max\"],\n ]\n if hasattr(self, \"resolution_level\"):\n assert (\n self.resolution_level >= 0 and self.resolution_level < 10\n ), \"resolution_level should be in [0, 10)\"\n pixels_range = pixels_bounds[1] - pixels_bounds[0]\n interval = pixels_range / 10\n new_lowbound = self.resolution_level * interval + pixels_bounds[0]\n new_upbound = (self.resolution_level + 1) * interval + pixels_bounds[0]\n pixels_bounds = (new_lowbound, new_upbound)\n else:\n warnings.warn(\"!! self.resolution_level not set, using default bounds !!\")\n\n # houskeeping on cpu/cuda and batchify\n if rgb.ndim == 3:\n rgb = rgb.unsqueeze(0)\n if camera is not None:\n if isinstance(camera, torch.Tensor):\n assert (\n camera.shape[-1] == 3 and camera.shape[-2] == 3\n ), \"camera tensor should be of shape (..., 3, 3): assume pinhole\"\n camera = Pinhole(K=camera)\n camera = BatchCamera.from_camera(camera)\n camera = camera.to(self.device)\n B, _, H, W = rgb.shape\n\n rgb = rgb.to(self.device)\n if camera is not None:\n camera = camera.to(self.device)\n\n # preprocess\n paddings, (padded_H, padded_W) = get_paddings((H, W), ratio_bounds)\n (pad_left, pad_right, pad_top, pad_bottom) = paddings\n resize_factor, (new_H, new_W) = get_resize_factor(\n (padded_H, padded_W), pixels_bounds\n )\n # -> rgb preprocess (input std-ized and resized)\n if normalize:\n rgb = TF.normalize(\n rgb.float() / 255.0,\n mean=IMAGENET_DATASET_MEAN,\n std=IMAGENET_DATASET_STD,\n )\n rgb = F.pad(rgb, (pad_left, pad_right, pad_top, pad_bottom), value=0.0)\n rgb = F.interpolate(\n rgb, size=(new_H, new_W), mode=\"bilinear\", align_corners=False\n )\n # -> camera preprocess\n if camera is not None:\n camera = camera.crop(\n left=-pad_left, top=-pad_top, right=-pad_right, bottom=-pad_bottom\n )\n camera = camera.resize(resize_factor)\n\n # run model\n _, model_outputs = self.encode_decode(\n inputs={\"image\": rgb, \"camera\": camera}, image_metas=[]\n )\n\n # collect outputs\n out = {}\n out[\"confidence\"] = _postprocess(\n model_outputs[\"confidence\"],\n (padded_H, padded_W),\n paddings=paddings,\n interpolation_mode=self.interpolation_mode,\n )\n points = _postprocess(\n model_outputs[\"points\"],\n (padded_H, padded_W),\n paddings=paddings,\n interpolation_mode=self.interpolation_mode,\n )\n rays = _postprocess(\n model_outputs[\"rays\"],\n (padded_H, padded_W),\n paddings=paddings,\n interpolation_mode=self.interpolation_mode,\n )\n out[\"intrinsics\"] = _postprocess_intrinsics(\n model_outputs[\"intrinsics\"], [resize_factor] * B, [paddings] * B\n )\n\n out[\"radius\"] = points.norm(dim=1, keepdim=True)\n out[\"depth\"] = points[:, -1:]\n out[\"points\"] = points\n out[\"rays\"] = rays / torch.norm(rays, dim=1, keepdim=True).clip(min=1e-5)\n out[\"depth_features\"] = model_outputs[\"depth_features\"]\n return out\n\n def encode_decode(self, inputs, image_metas=[]):\n B, _, H, W = inputs[\"image\"].shape\n\n # shortcut eval should avoid errors\n if len(image_metas) and \"paddings\" in image_metas[0]:\n inputs[\"paddings\"] = torch.tensor(\n [image_meta[\"paddings\"] for image_meta in image_metas],\n device=self.device,\n )[\n ..., [0, 2, 1, 3]\n ] # lrtb\n inputs[\"depth_paddings\"] = torch.tensor(\n [image_meta[\"depth_paddings\"] for image_meta in image_metas],\n device=self.device,\n )\n if (\n self.training\n ): # at inference we do not have image paddings on top of depth ones (we have not \"crop\" on gt in ContextCrop)\n inputs[\"depth_paddings\"] = inputs[\"depth_paddings\"] + inputs[\"paddings\"]\n\n if inputs.get(\"camera\", None) is not None:\n inputs[\"rays\"] = inputs[\"camera\"].get_rays(shapes=(B, H, W))\n\n features, tokens = self.pixel_encoder(inputs[\"image\"])\n inputs[\"features\"] = [\n self.stacking_fn(features[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n inputs[\"tokens\"] = [\n self.stacking_fn(tokens[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n\n outputs = self.pixel_decoder(inputs, image_metas)\n outputs[\"rays\"] = rearrange(outputs[\"rays\"], \"b (h w) c -> b c h w\", h=H, w=W)\n pts_3d = outputs[\"rays\"] * outputs[\"radius\"]\n outputs.update({\"points\": pts_3d, \"depth\": pts_3d[:, -1:]})\n\n return inputs, outputs\n\n def load_pretrained(self, model_file):\n device = (\n torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n )\n dict_model = torch.load(model_file, map_location=device, weights_only=False)\n if \"model\" in dict_model:\n dict_model = dict_model[\"model\"]\n dict_model = {k.replace(\"module.\", \"\"): v for k, v in dict_model.items()}\n info = self.load_state_dict(dict_model, strict=False)\n if is_main_process():\n print(\n f\"Loaded from {model_file} for {self.__class__.__name__} results in:\",\n info,\n )\n\n def get_params(self, config):\n if hasattr(self.pixel_encoder, \"get_params\"):\n encoder_p, encoder_lr = self.pixel_encoder.get_params(\n config[\"model\"][\"pixel_encoder\"][\"lr\"],\n config[\"training\"][\"wd\"],\n config[\"training\"][\"ld\"],\n )\n else:\n encoder_p, encoder_lr = get_params(\n self.pixel_encoder,\n config[\"model\"][\"pixel_encoder\"][\"lr\"],\n config[\"training\"][\"wd\"],\n )\n decoder_p, decoder_lr = get_params(\n self.pixel_decoder, config[\"training\"][\"lr\"], config[\"training\"][\"wd\"]\n )\n return [*encoder_p, *decoder_p]\n\n @property\n def device(self):\n return next(self.parameters()).device\n\n def build(self, config):\n mod = importlib.import_module(\"unidepth.models.encoder\")\n pixel_encoder_factory = getattr(mod, config[\"model\"][\"pixel_encoder\"][\"name\"])\n pixel_encoder_config = {\n **config[\"training\"],\n **config[\"model\"][\"pixel_encoder\"],\n **config[\"data\"],\n }\n pixel_encoder = pixel_encoder_factory(pixel_encoder_config)\n\n config[\"model\"][\"pixel_encoder\"][\"patch_size\"] = (\n 14 if \"dino\" in config[\"model\"][\"pixel_encoder\"][\"name\"] else 16\n )\n pixel_encoder_embed_dims = (\n pixel_encoder.embed_dims\n if hasattr(pixel_encoder, \"embed_dims\")\n else [getattr(pixel_encoder, \"embed_dim\") * 2**i for i in range(4)]\n )\n config[\"model\"][\"pixel_encoder\"][\"embed_dim\"] = getattr(\n pixel_encoder, \"embed_dim\"\n )\n config[\"model\"][\"pixel_encoder\"][\"embed_dims\"] = pixel_encoder_embed_dims\n config[\"model\"][\"pixel_encoder\"][\"depths\"] = pixel_encoder.depths\n config[\"model\"][\"pixel_encoder\"][\"cls_token_embed_dims\"] = getattr(\n pixel_encoder, \"cls_token_embed_dims\", pixel_encoder_embed_dims\n )\n\n pixel_decoder = Decoder(config)\n\n self.pixel_encoder = pixel_encoder\n self.pixel_decoder = pixel_decoder\n\n self.slices_encoder_range = list(\n zip([0, *self.pixel_encoder.depths[:-1]], self.pixel_encoder.depths)\n )\n\n stacking_fn = config[\"model\"][\"pixel_encoder\"][\"stacking_fn\"]\n assert (\n stacking_fn in STACKING_FNS\n ), f\"Stacking function {stacking_fn} not found in {STACKING_FNS.keys()}\"\n self.stacking_fn = STACKING_FNS[stacking_fn]\n self.shape_constraints = config[\"data\"][\"augmentations\"][\"shape_constraints\"]\n self.interpolation_mode = \"bilinear\"\n\n def build_losses(self, config):\n self.losses = {}\n for loss_name, loss_config in config[\"training\"][\"losses\"].items():\n mod = importlib.import_module(\"unidepth.ops.losses\")\n loss_factory = getattr(mod, loss_config[\"name\"])\n self.losses[loss_name] = loss_factory.build(loss_config)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/models/unidepthv2/unidepthv2_old.py", "content": "import importlib\nimport warnings\nfrom copy import deepcopy\nfrom math import ceil\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms.functional as TF\nfrom einops import rearrange\nfrom huggingface_hub import PyTorchModelHubMixin\n\nfrom unidepth.models.unidepthv2.decoder_old import Decoder\nfrom unidepth.utils.constants import (IMAGENET_DATASET_MEAN,\n IMAGENET_DATASET_STD)\nfrom unidepth.utils.distributed import is_main_process\nfrom unidepth.utils.geometric import (generate_rays,\n spherical_zbuffer_to_euclidean)\nfrom unidepth.utils.misc import (first_stack, last_stack, max_stack,\n mean_stack, softmax_stack)\n\nSTACKING_FNS = {\n \"max\": max_stack,\n \"mean\": mean_stack,\n \"first\": first_stack,\n \"last\": last_stack,\n \"softmax\": softmax_stack,\n}\nRESOLUTION_LEVELS = 10\n\n\n# inference helpers\ndef _check_ratio(image_ratio, ratio_bounds):\n ratio_bounds = sorted(ratio_bounds)\n if ratio_bounds is not None and (\n image_ratio < ratio_bounds[0] or image_ratio > ratio_bounds[1]\n ):\n warnings.warn(\n f\"Input image ratio ({image_ratio:.3f}) is out of training \"\n f\"distribution: {ratio_bounds}. This may lead to unexpected results. \"\n f\"Consider resizing/padding the image to match the training distribution.\"\n )\n\n\ndef _check_resolution(shape_constraints, resolution_level):\n if resolution_level is None:\n warnings.warn(\n \"Resolution level is not set. Using max resolution. \"\n \"You can tradeoff resolution for speed by setting a number in [0,10]. \"\n \"This can be achieved by setting model's `resolution_level` attribute.\"\n )\n resolution_level = RESOLUTION_LEVELS\n pixel_bounds = sorted(shape_constraints[\"pixels_bounds_ori\"])\n pixel_range = pixel_bounds[-1] - pixel_bounds[0]\n clipped_resolution_level = min(max(resolution_level, 0), RESOLUTION_LEVELS)\n if clipped_resolution_level != resolution_level:\n warnings.warn(\n f\"Resolution level {resolution_level} is out of bounds ([0,{RESOLUTION_LEVELS}]). \"\n f\"Clipping to {clipped_resolution_level}.\"\n )\n shape_constraints[\"pixels_bounds\"] = [\n pixel_bounds[0]\n + ceil(pixel_range * clipped_resolution_level / RESOLUTION_LEVELS),\n pixel_bounds[0]\n + ceil(pixel_range * clipped_resolution_level / RESOLUTION_LEVELS),\n ]\n return shape_constraints\n\n\ndef _get_closes_num_pixels(image_shape, pixels_bounds):\n h, w = image_shape\n num_pixels = h * w\n pixels_bounds = sorted(pixels_bounds)\n num_pixels = max(min(num_pixels, pixels_bounds[1]), pixels_bounds[0])\n return num_pixels\n\n\ndef _shapes(image_shape, shape_constraints):\n h, w = image_shape\n image_ratio = w / h\n _check_ratio(image_ratio, shape_constraints[\"ratio_bounds\"])\n num_pixels = _get_closes_num_pixels(\n (h / shape_constraints[\"patch_size\"], w / shape_constraints[\"patch_size\"]),\n shape_constraints[\"pixels_bounds\"],\n )\n h = ceil((num_pixels / image_ratio) ** 0.5 - 0.5)\n w = ceil(h * image_ratio - 0.5)\n ratio = h / image_shape[0] * shape_constraints[\"patch_size\"]\n return (\n h * shape_constraints[\"patch_size\"],\n w * shape_constraints[\"patch_size\"],\n ), ratio\n\n\ndef _preprocess(rgbs, intrinsics, shapes, ratio):\n rgbs = F.interpolate(rgbs, size=shapes, mode=\"bilinear\", antialias=True)\n if intrinsics is not None:\n intrinsics = intrinsics.clone()\n intrinsics[:, 0, 0] = intrinsics[:, 0, 0] * ratio\n intrinsics[:, 1, 1] = intrinsics[:, 1, 1] * ratio\n intrinsics[:, 0, 2] = intrinsics[:, 0, 2] * ratio\n intrinsics[:, 1, 2] = intrinsics[:, 1, 2] * ratio\n return rgbs, intrinsics\n return rgbs, None\n\n\ndef _postprocess(outs, ratio, original_shapes, mode=\"nearest-exact\"):\n outs[\"depth\"] = F.interpolate(outs[\"depth\"], size=original_shapes, mode=mode)\n outs[\"confidence\"] = F.interpolate(\n outs[\"confidence\"], size=original_shapes, mode=\"bilinear\", antialias=True\n )\n outs[\"K\"][:, 0, 0] = outs[\"K\"][:, 0, 0] / ratio\n outs[\"K\"][:, 1, 1] = outs[\"K\"][:, 1, 1] / ratio\n outs[\"K\"][:, 0, 2] = outs[\"K\"][:, 0, 2] / ratio\n outs[\"K\"][:, 1, 2] = outs[\"K\"][:, 1, 2] / ratio\n return outs\n\n\nclass UniDepthV2old(\n nn.Module,\n PyTorchModelHubMixin,\n library_name=\"UniDepth\",\n repo_url=\"https://github.com/lpiccinelli-eth/UniDepth\",\n tags=[\"monocular-metric-depth-estimation\"],\n):\n def __init__(\n self,\n config,\n **kwargs,\n ):\n super().__init__()\n self.build(config)\n\n def forward(self, inputs, image_metas):\n H, W = inputs[\"depth\"].shape[-2:]\n\n if \"K\" in inputs:\n rays, angles = generate_rays(inputs[\"K\"], (H, W))\n inputs[\"rays\"] = rays\n inputs[\"angles\"] = angles\n\n features, tokens = self.pixel_encoder(inputs[f\"image\"])\n\n cls_tokens = [x.contiguous() for x in tokens]\n features = [\n self.stacking_fn(features[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n tokens = [\n self.stacking_fn(tokens[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n global_tokens = [cls_tokens[i] for i in [-2, -1]]\n camera_tokens = [cls_tokens[i] for i in [-3, -2, -1]] + [tokens[-2]]\n\n inputs[\"features\"] = features\n inputs[\"tokens\"] = tokens\n inputs[\"global_tokens\"] = global_tokens\n inputs[\"camera_tokens\"] = camera_tokens\n\n outs = self.pixel_decoder(inputs, image_metas)\n\n angles = rearrange(\n generate_rays(outs[\"K\"], (H, W), noisy=False)[-1],\n \"b (h w) c -> b c h w\",\n h=H,\n w=W,\n )\n predictions = F.interpolate(\n outs[\"depth\"],\n size=(H, W),\n mode=\"bilinear\",\n align_corners=False,\n antialias=True,\n )\n confidence = F.interpolate(\n outs[\"confidence\"],\n size=(H, W),\n mode=\"bilinear\",\n align_corners=False,\n antialias=True,\n )\n predictions_3d = torch.cat((angles, predictions), dim=1)\n predictions_3d = spherical_zbuffer_to_euclidean(\n predictions_3d.permute(0, 2, 3, 1)\n ).permute(0, 3, 1, 2)\n\n outputs = {\n \"K\": outs[\"K\"],\n \"depth\": predictions,\n \"confidence\": confidence,\n \"points\": predictions_3d,\n \"depth_features\": outs[\"depth_features\"],\n }\n return outputs\n\n @torch.no_grad()\n def infer(self, rgbs: torch.Tensor, intrinsics=None):\n shape_constraints = self.shape_constraints\n if rgbs.ndim == 3:\n rgbs = rgbs.unsqueeze(0)\n if intrinsics is not None and intrinsics.ndim == 2:\n intrinsics = intrinsics.unsqueeze(0)\n B, _, H, W = rgbs.shape\n\n rgbs = rgbs.to(self.device)\n if intrinsics is not None:\n intrinsics = intrinsics.to(self.device)\n\n # process image and intrinsiscs (if any) to match network input (slow?)\n if rgbs.max() > 5 or rgbs.dtype == torch.uint8:\n rgbs = rgbs.to(torch.float32).div(255)\n if rgbs.min() >= 0.0 and rgbs.max() <= 1.0:\n rgbs = TF.normalize(\n rgbs,\n mean=IMAGENET_DATASET_MEAN,\n std=IMAGENET_DATASET_STD,\n )\n\n # check resolution constraints: tradeoff resolution and speed\n shape_constraints = _check_resolution(shape_constraints, self.resolution_level)\n\n # get image shape\n (h, w), ratio = _shapes((H, W), shape_constraints)\n rgbs, gt_intrinsics = _preprocess(\n rgbs,\n intrinsics,\n (h, w),\n ratio,\n )\n\n # run encoder\n features, tokens = self.pixel_encoder(rgbs)\n\n cls_tokens = [x.contiguous() for x in tokens]\n features = [\n self.stacking_fn(features[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n tokens = [\n self.stacking_fn(tokens[i:j]).contiguous()\n for i, j in self.slices_encoder_range\n ]\n global_tokens = [cls_tokens[i] for i in [-2, -1]]\n camera_tokens = [cls_tokens[i] for i in [-3, -2, -1]] + [tokens[-2]]\n\n # get data fro decoder and adapt to given camera\n inputs = {}\n inputs[\"features\"] = features\n inputs[\"tokens\"] = tokens\n inputs[\"global_tokens\"] = global_tokens\n inputs[\"camera_tokens\"] = camera_tokens\n inputs[\"image\"] = rgbs\n if gt_intrinsics is not None:\n rays, angles = generate_rays(gt_intrinsics, (h, w))\n inputs[\"rays\"] = rays\n inputs[\"angles\"] = angles\n inputs[\"K\"] = gt_intrinsics\n\n outs = self.pixel_decoder(inputs, {})\n # undo the reshaping and get original image size (slow)\n outs = _postprocess(outs, ratio, (H, W), mode=self.interpolation_mode)\n pred_intrinsics = outs[\"K\"]\n depth = outs[\"depth\"]\n confidence = outs[\"confidence\"]\n\n # final 3D points backprojection\n intrinsics = intrinsics if intrinsics is not None else pred_intrinsics\n angles = generate_rays(intrinsics, (H, W))[-1]\n angles = rearrange(angles, \"b (h w) c -> b c h w\", h=H, w=W)\n points_3d = torch.cat((angles, depth), dim=1)\n points_3d = spherical_zbuffer_to_euclidean(\n points_3d.permute(0, 2, 3, 1)\n ).permute(0, 3, 1, 2)\n\n outputs = {\n \"intrinsics\": pred_intrinsics,\n \"points\": points_3d,\n \"depth\": depth,\n \"confidence\": confidence,\n }\n return outputs\n\n def load_pretrained(self, model_file):\n device = (\n torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n )\n dict_model = torch.load(model_file, map_location=device)\n if \"model\" in dict_model:\n dict_model = dict_model[\"model\"]\n dict_model = deepcopy(\n {k.replace(\"module.\", \"\"): v for k, v in dict_model.items()}\n )\n\n info = self.load_state_dict(dict_model, strict=False)\n if is_main_process():\n print(\n f\"Loaded from {model_file} for {self.__class__.__name__} results in:\",\n info,\n )\n\n @property\n def device(self):\n return next(self.parameters()).device\n\n def build(self, config):\n mod = importlib.import_module(\"unidepth.models.encoder\")\n pixel_encoder_factory = getattr(mod, config[\"model\"][\"pixel_encoder\"][\"name\"])\n pixel_encoder_config = {\n **config[\"training\"],\n **config[\"model\"][\"pixel_encoder\"],\n **config[\"data\"],\n }\n pixel_encoder = pixel_encoder_factory(pixel_encoder_config)\n\n config[\"model\"][\"pixel_encoder\"][\"patch_size\"] = (\n 14 if \"dino\" in config[\"model\"][\"pixel_encoder\"][\"name\"] else 16\n )\n pixel_encoder_embed_dims = (\n pixel_encoder.embed_dims\n if hasattr(pixel_encoder, \"embed_dims\")\n else [getattr(pixel_encoder, \"embed_dim\") * 2**i for i in range(4)]\n )\n config[\"model\"][\"pixel_encoder\"][\"embed_dim\"] = getattr(\n pixel_encoder, \"embed_dim\"\n )\n config[\"model\"][\"pixel_encoder\"][\"embed_dims\"] = pixel_encoder_embed_dims\n config[\"model\"][\"pixel_encoder\"][\"depths\"] = pixel_encoder.depths\n\n pixel_decoder = Decoder(config)\n\n self.pixel_encoder = pixel_encoder\n self.pixel_decoder = pixel_decoder\n stacking_fn = config[\"model\"][\"pixel_encoder\"][\"stacking_fn\"]\n assert (\n stacking_fn in STACKING_FNS\n ), f\"Stacking function {stacking_fn} not found in {STACKING_FNS.keys()}\"\n self.stacking_fn = STACKING_FNS[stacking_fn]\n\n self.slices_encoder_range = list(\n zip([0, *pixel_encoder.depths[:-1]], pixel_encoder.depths)\n )\n self.shape_constraints = config[\"data\"][\"shape_constraints\"]\n self.shape_constraints[\"pixels_bounds_ori\"] = self.shape_constraints.get(\n \"pixels_bounds\", [1400, 2400]\n )\n self.interpolation_mode = \"bilinear\"\n self.eps = 1e-6\n self.resolution_level = None\n\n def build_losses(self, config):\n self.losses = {}\n for loss_name, loss_config in config[\"training\"][\"losses\"].items():\n mod = importlib.import_module(\"unidepth.ops.losses\")\n loss_factory = getattr(mod, loss_config[\"name\"])\n self.losses[loss_name] = loss_factory.build(loss_config)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/__init__.py", "content": "from .losses import (ARel, Confidence, Dummy, EdgeGuidedLocalSSI, LocalSSI,\n Regression, SelfDistill, SILog, TeacherDistill)\nfrom .scheduler import CosineScheduler, PlainCosineScheduler\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/__init__.py", "content": "from .functions import ExtractPatchesFunction\nfrom .modules import RandomPatchExtractor\n\n__all__ = [\"ExtractPatchesFunction\", \"RandomPatchExtractor\"]" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/compile.sh", "content": "#!/usr/bin/env bash\n\nif [ -z \"$TORCH_CUDA_ARCH_LIST\" ]; then\n export TORCH_CUDA_ARCH_LIST=\"7.5 8.0 8.6+PTX\"\nfi\n\npython setup.py build install\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/functions/__init__.py", "content": "from .extract_patches import ExtractPatchesFunction\n\n__all__ = [\"ExtractPatchesFunction\"]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/functions/extract_patches.py", "content": "import RandomPatchExtraction\nimport torch\nfrom torch.autograd import Function\n\n\nclass ExtractPatchesFunction(Function):\n @staticmethod\n def forward(ctx, input, centers, h, w):\n # Save variables for backward pass. inputs for shapes\n ctx.save_for_backward(input, centers)\n\n return RandomPatchExtraction.extract_patches_forward(input, centers, h, w)\n\n @staticmethod\n def backward(ctx, grad_output):\n input, centers = ctx.saved_tensors\n\n (grad_input,) = RandomPatchExtraction.extract_patches_backward(\n grad_output, centers, input.shape[2], input.shape[3]\n )\n # breakpoint()\n\n # Return gradients with respect to inputs only\n return grad_input, None, None, None\n\n\n# Test\nif __name__ == \"__main__\":\n B, C, H, W = 1, 1, 10, 10\n N = 2\n h, w = 3, 3\n input = torch.arange(\n B * C * H * W, device=\"cuda\", dtype=torch.float32, requires_grad=True\n ).view(B, C, H, W)\n centers = torch.tensor([[[4, 4], [6, 6]]], device=\"cuda\", dtype=torch.int32)\n output = ExtractPatchesFunction.apply(input, centers, h, w)\n output.mean().backward()\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/modules/__init__.py", "content": "from .patch_extractor import RandomPatchExtractor\n\n__all__ = [\"RandomPatchExtractor\"]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/modules/patch_extractor.py", "content": "from __future__ import absolute_import, division, print_function\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\n\nfrom ..functions import ExtractPatchesFunction\n\n\nclass RandomPatchExtractor(nn.Module):\n def __init__(\n self,\n ):\n super().__init__()\n\n def forward(\n self, tensor: torch.Tensor, centers: torch.Tensor, patch_size: tuple[int, int]\n ):\n device = tensor.device\n dtype = tensor.dtype\n patch_width, patch_height = patch_size\n pad_width = patch_width // 2\n pad_height = patch_height // 2\n dtype = tensor.dtype\n\n # Pad input to avoid out-of-bounds\n tensor_padded = F.pad(\n tensor,\n (pad_width, pad_width, pad_height, pad_height),\n mode=\"constant\",\n value=0.0,\n )\n\n # Adjust edge coordinates to account for padding\n centers_padded = centers + torch.tensor(\n [pad_height, pad_width], dtype=dtype, device=device\n ).reshape(1, 1, 2)\n\n output = ExtractPatchesFunction.apply(\n tensor_padded.float(), centers_padded.int(), patch_height, patch_width\n )\n return output.to(dtype)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/setup.py", "content": "import glob\nimport os\n\nimport torch\nfrom setuptools import find_packages, setup\nfrom torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension\n\nrequirements = [\"torch\", \"torchvision\"]\n\n\ndef get_extensions():\n this_dir = os.path.dirname(os.path.abspath(__file__))\n extensions_dir = os.path.join(this_dir, \"src\")\n\n main_file = glob.glob(os.path.join(extensions_dir, \"*.cpp\"))\n source_cpu = glob.glob(os.path.join(extensions_dir, \"cpu\", \"*.cpp\"))\n source_cuda = glob.glob(os.path.join(extensions_dir, \"cuda\", \"*.cu\"))\n\n sources = main_file + source_cpu\n extension = CppExtension\n extra_compile_args = {\"cxx\": [\"-O2\"]}\n define_macros = []\n\n if torch.cuda.is_available() and CUDA_HOME is not None:\n extension = CUDAExtension\n sources += source_cuda\n define_macros += [(\"WITH_CUDA\", None)]\n extra_compile_args[\"nvcc\"] = [\n \"-O2\",\n ]\n else:\n raise NotImplementedError(\"Cuda is not available\")\n\n sources = list(set([os.path.join(extensions_dir, s) for s in sources]))\n include_dirs = [extensions_dir]\n ext_modules = [\n extension(\n \"RandomPatchExtraction\",\n sources,\n include_dirs=include_dirs,\n define_macros=define_macros,\n extra_compile_args=extra_compile_args,\n )\n ]\n\n return ext_modules\n\n\nsetup(\n name=\"RandomPatchExtraction\",\n version=\"0.1\",\n author=\"Luigi Piccinelli\",\n ext_modules=get_extensions(),\n packages=find_packages(\n exclude=(\n \"configs\",\n \"tests\",\n )\n ),\n cmdclass={\"build_ext\": torch.utils.cpp_extension.BuildExtension},\n)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/src/cpu/extract_patches_cpu.cpp", "content": "#include \n\n#include \n#include \n\ntorch::Tensor extract_patches_cpu_forward(\n const torch::Tensor &input,\n const torch::Tensor ¢ers,\n int h,\n int w\n) {\n AT_ERROR(\"Not implement on cpu\");\n}\n\nstd::vector extract_patches_cpu_backward(\n const torch::Tensor &grad_patches,\n const torch::Tensor &coords,\n int H,\n int W\n) {\n AT_ERROR(\"Not implement on cpu\");\n}\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/src/cpu/extract_patches_cpu.h", "content": "#pragma once\n#include \n#include \n\n\ntorch::Tensor extract_patches_cpu_forward(\n const torch::Tensor &input,\n const torch::Tensor ¢ers,\n int h,\n int w\n);\n\nstd::vector extract_patches_cpu_backward(\n const torch::Tensor &grad_patches,\n const torch::Tensor &coords,\n int H,\n int W\n);\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/src/cuda/extract_patches_cuda.h", "content": "#ifndef EXTRACT_PATCHES_CUDA_H\n#define EXTRACT_PATCHES_CUDA_H\n\n#include \n#include \n#include \n#include \n\n// Function prototypes for the CUDA functions\ntorch::Tensor extract_patches_cuda_forward(\n const torch::Tensor &input,\n const torch::Tensor ¢ers,\n int h,\n int w\n);\n\nstd::vector extract_patches_cuda_backward(\n const torch::Tensor &grad_output,\n const torch::Tensor ¢ers,\n int H,\n int W\n);\n\n\n#endif // EXTRACT_PATCHES_CUDA_H\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/src/cuda/extract_patches_kernel.cu", "content": "#include \n#include \n\n#include \"cuda/extract_patches_kernel.cuh\"\n#include \"cuda/extract_patches_cuda.h\"\n\n\n// Need to templetize these two to get fp16 working, but problems in compilation...\ntorch::Tensor extract_patches_cuda_forward(\n const torch::Tensor &input,\n const torch::Tensor ¢ers,\n int h,\n int w\n) {\n \n int B = input.size(0);\n int C = input.size(1);\n int H = input.size(2);\n int W = input.size(3);\n int N = centers.size(1);\n\n auto output = torch::zeros({B, C, N, h, w}, input.options());\n\n const int threads = C;\n const dim3 blocks(B, N);\n\n extract_patches_cuda_forward_kernel<<>>(\n input.data_ptr(),\n output.data_ptr(),\n centers.data_ptr(),\n B, C, H, W,\n N, h, w);\n\n return {output};\n}\n\nstd::vector extract_patches_cuda_backward(\n const torch::Tensor &grad_output,\n const torch::Tensor ¢ers,\n int H,\n int W\n) {\n \n int B = grad_output.size(0);\n int C = grad_output.size(1);\n int N = centers.size(1);\n int h = grad_output.size(3);\n int w = grad_output.size(4);\n\n auto grad_input = torch::zeros({B, C, H, W}, grad_output.options());\n\n const int threads = C;\n const dim3 blocks(B, N);\n\n extract_patches_cuda_backward_kernel<<>>(\n grad_output.data_ptr(),\n grad_input.data_ptr(),\n centers.data_ptr(),\n B, C, H, W,\n N, h, w);\n\n return {grad_input};\n}\n\ntemplate \n__global__ void extract_patches_cuda_forward_kernel(\n const T* __restrict__ input,\n T* __restrict__ output,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w) {\n \n // Calculate thread indices\n int batch_idx = blockIdx.x;\n int patch_idx = blockIdx.y;\n int channel_idx = threadIdx.x;\n\n // Extract center coordinates\n int center_y = centers[(batch_idx * N + patch_idx) * 2];\n int center_x = centers[(batch_idx * N + patch_idx) * 2 + 1];\n\n // Calculate half patch size\n int half_h = h / 2;\n int half_w = w / 2;\n\n // Extract patch\n for (int i = 0; i < h; ++i) {\n for (int j = 0; j < w; ++j) {\n int y = center_y - half_h + i;\n int x = center_x - half_w + j;\n output[batch_idx * C * N * h * w + patch_idx * C * h * w + channel_idx * h * w + i * w + j] = \n input[batch_idx * C * H * W + channel_idx * H * W + y * W + x];\n }\n }\n}\n\ntemplate __global__ void extract_patches_cuda_forward_kernel(\n const float* __restrict__ input,\n float* __restrict__ output,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w);\n\ntemplate __global__ void extract_patches_cuda_forward_kernel<__half>(\n const __half* __restrict__ input,\n __half* __restrict__ output,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w);\n\ntemplate \n__global__ void extract_patches_cuda_backward_kernel(\n const T* __restrict__ grad_output,\n T* __restrict__ grad_input,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w) {\n \n // Calculate thread indices\n int batch_idx = blockIdx.x;\n int patch_idx = blockIdx.y;\n int channel_idx = threadIdx.x;\n\n // Extract center coordinates\n int center_y = centers[(batch_idx * N + patch_idx) * 2];\n int center_x = centers[(batch_idx * N + patch_idx) * 2 + 1];\n\n // Calculate half patch size\n int half_h = h / 2;\n int half_w = w / 2;\n\n // Compute gradients with respect to input tensor using chain rule\n for (int i = 0; i < h; ++i) {\n for (int j = 0; j < w; ++j) {\n int y = center_y - half_h + i;\n int x = center_x - half_w + j;\n \n atomicAdd(\n &grad_input[batch_idx * C * H * W + channel_idx * H * W + y * W + x],\n grad_output[batch_idx * C * N * h * w + patch_idx * C * h * w + channel_idx * h * w + i * w + j]\n );\n }\n }\n}\n\ntemplate __global__ void extract_patches_cuda_backward_kernel(\n const float* __restrict__ grad_output,\n float* __restrict__ grad_input,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w);\n\ntemplate __global__ void extract_patches_cuda_backward_kernel<__half>(\n const __half* __restrict__ grad_output,\n __half* __restrict__ grad_input,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w);" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/src/cuda/extract_patches_kernel.cuh", "content": "#ifndef EXTRACT_PATCHES_KERNEL_CUH\n#define EXTRACT_PATCHES_KERNEL_CUH\n\n#include \n#include \n#include \n#include \n#include // should contain __half\n\n// Declare the forward CUDA kernel function\ntemplate __global__ void extract_patches_cuda_forward_kernel(\n const T* __restrict__ input,\n T* __restrict__ output,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w);\n\n// Declare the backward CUDA kernel function\ntemplate __global__ void extract_patches_cuda_backward_kernel(\n const T* __restrict__ grad_output,\n T* __restrict__ grad_input,\n const int* __restrict__ centers,\n int B, int C, int H, int W,\n int N, int h, int w);\n\n#endif // EXTRACT_PATCHES_KERNEL_CUH\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/src/extract_patches.cpp", "content": "\n#include \"extract_patches.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"extract_patches_forward\", &extract_patches_forward, \"Extract patches forward (CUDA)\");\n m.def(\"extract_patches_backward\", &extract_patches_backward, \"Extract patches backward (CUDA)\");\n}\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/src/extract_patches.h", "content": "#pragma once\n\n#include \"cpu/extract_patches_cpu.h\"\n\n#ifdef WITH_CUDA\n#include \"cuda/extract_patches_cuda.h\"\n#endif\n\n#include \n\n#include \n#include \n\ntorch::Tensor extract_patches_forward(\n const torch::Tensor &images,\n const torch::Tensor &coords,\n int patch_height,\n int patch_width)\n{\n if (images.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return extract_patches_cuda_forward(images, coords, patch_height, patch_width);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n AT_ERROR(\"Not implemented on the CPU\");\n}\n\nstd::vector extract_patches_backward(\n const torch::Tensor &grad_patches,\n const torch::Tensor &coords,\n int H,\n int W)\n{\n if (grad_patches.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return extract_patches_cuda_backward(grad_patches, coords, H, W);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n AT_ERROR(\"Not implemented on the CPU\");\n}" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/extract_patches/test.py", "content": "import RandomPatchExtraction\nimport torch\n\n\ndef extract_patches(input, centers, patch_size):\n h, w = patch_size\n output = RandomPatchExtraction.extract_patches_forward(input, centers, h, w)\n breakpoint()\n return output\n\n\n# Example usage\nif __name__ == \"__main__\":\n B, C, H, W = 1, 1, 10, 10\n N = 2\n h, w = 3, 3\n input = torch.arange(\n B * C * H * W, device=\"cuda\", dtype=torch.float32, requires_grad=True\n ).view(B, C, H, W)\n centers = torch.tensor([[[4, 4], [6, 6]]], device=\"cuda\", dtype=torch.int32)\n patches = extract_patches(input, centers, (h, w))\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/__init__.py", "content": "from .functions.knn import knn_gather, knn_points\n\n__all__ = [\n \"knn_points\",\n \"knn_gather\",\n]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/compile.sh", "content": "#!/usr/bin/env bash\n\nexport TORCH_CUDA_ARCH_LIST=\"6.1 7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX\" \n# export FORCE_CUDA=1 #if you do not actually have cuda, workaround\npython setup.py build install" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/functions/__init__.py", "content": "from .knn import knn_gather, knn_points\n\n__all__ = [\n \"knn_points\",\n \"knn_gather\",\n]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/functions/knn.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the BSD-style license found in the\n# LICENSE file in the root directory of this source tree.\n\n# pyre-unsafe\n\nfrom collections import namedtuple\nfrom typing import Union\n\nimport torch\nfrom KNN import knn_points_backward, knn_points_idx\nfrom torch.autograd import Function\nfrom torch.autograd.function import once_differentiable\n\n_KNN = namedtuple(\"KNN\", \"dists idx knn\")\n\n\nclass _knn_points(Function):\n \"\"\"\n Torch autograd Function wrapper for KNN C++/CUDA implementations.\n \"\"\"\n\n @staticmethod\n # pyre-fixme[14]: `forward` overrides method defined in `Function` inconsistently.\n def forward(\n ctx,\n p1,\n p2,\n lengths1,\n lengths2,\n K,\n version,\n norm: int = 2,\n return_sorted: bool = True,\n ):\n \"\"\"\n K-Nearest neighbors on point clouds.\n\n Args:\n p1: Tensor of shape (N, P1, D) giving a batch of N point clouds, each\n containing up to P1 points of dimension D.\n p2: Tensor of shape (N, P2, D) giving a batch of N point clouds, each\n containing up to P2 points of dimension D.\n lengths1: LongTensor of shape (N,) of values in the range [0, P1], giving the\n length of each pointcloud in p1. Or None to indicate that every cloud has\n length P1.\n lengths2: LongTensor of shape (N,) of values in the range [0, P2], giving the\n length of each pointcloud in p2. Or None to indicate that every cloud has\n length P2.\n K: Integer giving the number of nearest neighbors to return.\n version: Which KNN implementation to use in the backend. If version=-1,\n the correct implementation is selected based on the shapes of the inputs.\n norm: (int) indicating the norm. Only supports 1 (for L1) and 2 (for L2).\n return_sorted: (bool) whether to return the nearest neighbors sorted in\n ascending order of distance.\n\n Returns:\n p1_dists: Tensor of shape (N, P1, K) giving the squared distances to\n the nearest neighbors. This is padded with zeros both where a cloud in p2\n has fewer than K points and where a cloud in p1 has fewer than P1 points.\n\n p1_idx: LongTensor of shape (N, P1, K) giving the indices of the\n K nearest neighbors from points in p1 to points in p2.\n Concretely, if `p1_idx[n, i, k] = j` then `p2[n, j]` is the k-th nearest\n neighbors to `p1[n, i]` in `p2[n]`. This is padded with zeros both where a cloud\n in p2 has fewer than K points and where a cloud in p1 has fewer than P1 points.\n \"\"\"\n if not ((norm == 1) or (norm == 2)):\n raise ValueError(\"Support for 1 or 2 norm.\")\n\n idx, dists = knn_points_idx(p1, p2, lengths1, lengths2, norm, K, version)\n\n # sort KNN in ascending order if K > 1\n if K > 1 and return_sorted:\n if lengths2.min() < K:\n P1 = p1.shape[1]\n mask = lengths2[:, None] <= torch.arange(K, device=dists.device)[None]\n # mask has shape [N, K], true where dists irrelevant\n mask = mask[:, None].expand(-1, P1, -1)\n # mask has shape [N, P1, K], true where dists irrelevant\n dists[mask] = float(\"inf\")\n dists, sort_idx = dists.sort(dim=2)\n dists[mask] = 0\n else:\n dists, sort_idx = dists.sort(dim=2)\n idx = idx.gather(2, sort_idx)\n\n ctx.save_for_backward(p1, p2, lengths1, lengths2, idx)\n ctx.mark_non_differentiable(idx)\n ctx.norm = norm\n return dists, idx\n\n @staticmethod\n @once_differentiable\n def backward(ctx, grad_dists, grad_idx):\n p1, p2, lengths1, lengths2, idx = ctx.saved_tensors\n norm = ctx.norm\n # TODO(gkioxari) Change cast to floats once we add support for doubles.\n if not (grad_dists.dtype == torch.float32):\n grad_dists = grad_dists.float()\n if not (p1.dtype == torch.float32):\n p1 = p1.float()\n if not (p2.dtype == torch.float32):\n p2 = p2.float()\n grad_p1, grad_p2 = knn_points_backward(\n p1, p2, lengths1, lengths2, idx, norm, grad_dists\n )\n return grad_p1, grad_p2, None, None, None, None, None, None\n\n\ndef knn_points(\n p1: torch.Tensor,\n p2: torch.Tensor,\n lengths1: Union[torch.Tensor, None] = None,\n lengths2: Union[torch.Tensor, None] = None,\n norm: int = 2,\n K: int = 1,\n version: int = -1,\n return_nn: bool = False,\n return_sorted: bool = True,\n) -> _KNN:\n \"\"\"\n K-Nearest neighbors on point clouds.\n\n Args:\n p1: Tensor of shape (N, P1, D) giving a batch of N point clouds, each\n containing up to P1 points of dimension D.\n p2: Tensor of shape (N, P2, D) giving a batch of N point clouds, each\n containing up to P2 points of dimension D.\n lengths1: LongTensor of shape (N,) of values in the range [0, P1], giving the\n length of each pointcloud in p1. Or None to indicate that every cloud has\n length P1.\n lengths2: LongTensor of shape (N,) of values in the range [0, P2], giving the\n length of each pointcloud in p2. Or None to indicate that every cloud has\n length P2.\n norm: Integer indicating the norm of the distance. Supports only 1 for L1, 2 for L2.\n K: Integer giving the number of nearest neighbors to return.\n version: Which KNN implementation to use in the backend. If version=-1,\n the correct implementation is selected based on the shapes of the inputs.\n return_nn: If set to True returns the K nearest neighbors in p2 for each point in p1.\n return_sorted: (bool) whether to return the nearest neighbors sorted in\n ascending order of distance.\n\n Returns:\n dists: Tensor of shape (N, P1, K) giving the squared distances to\n the nearest neighbors. This is padded with zeros both where a cloud in p2\n has fewer than K points and where a cloud in p1 has fewer than P1 points.\n\n idx: LongTensor of shape (N, P1, K) giving the indices of the\n K nearest neighbors from points in p1 to points in p2.\n Concretely, if `p1_idx[n, i, k] = j` then `p2[n, j]` is the k-th nearest\n neighbors to `p1[n, i]` in `p2[n]`. This is padded with zeros both where a cloud\n in p2 has fewer than K points and where a cloud in p1 has fewer than P1\n points.\n\n nn: Tensor of shape (N, P1, K, D) giving the K nearest neighbors in p2 for\n each point in p1. Concretely, `p2_nn[n, i, k]` gives the k-th nearest neighbor\n for `p1[n, i]`. Returned if `return_nn` is True.\n The nearest neighbors are collected using `knn_gather`\n\n .. code-block::\n\n p2_nn = knn_gather(p2, p1_idx, lengths2)\n\n which is a helper function that allows indexing any tensor of shape (N, P2, U) with\n the indices `p1_idx` returned by `knn_points`. The output is a tensor\n of shape (N, P1, K, U).\n\n \"\"\"\n if p1.shape[0] != p2.shape[0]:\n raise ValueError(\"pts1 and pts2 must have the same batch dimension.\")\n if p1.shape[2] != p2.shape[2]:\n raise ValueError(\"pts1 and pts2 must have the same point dimension.\")\n\n p1 = p1.contiguous()\n p2 = p2.contiguous()\n\n P1 = p1.shape[1]\n P2 = p2.shape[1]\n\n if lengths1 is None:\n lengths1 = torch.full((p1.shape[0],), P1, dtype=torch.int64, device=p1.device)\n if lengths2 is None:\n lengths2 = torch.full((p1.shape[0],), P2, dtype=torch.int64, device=p1.device)\n\n p1_dists, p1_idx = _knn_points.apply(\n p1, p2, lengths1, lengths2, K, version, norm, return_sorted\n )\n\n p2_nn = None\n if return_nn:\n p2_nn = knn_gather(p2, p1_idx, lengths2)\n\n return _KNN(dists=p1_dists, idx=p1_idx, knn=p2_nn if return_nn else None)\n\n\ndef knn_gather(\n x: torch.Tensor, idx: torch.Tensor, lengths: Union[torch.Tensor, None] = None\n):\n \"\"\"\n A helper function for knn that allows indexing a tensor x with the indices `idx`\n returned by `knn_points`.\n\n For example, if `dists, idx = knn_points(p, x, lengths_p, lengths, K)`\n where p is a tensor of shape (N, L, D) and x a tensor of shape (N, M, D),\n then one can compute the K nearest neighbors of p with `p_nn = knn_gather(x, idx, lengths)`.\n It can also be applied for any tensor x of shape (N, M, U) where U != D.\n\n Args:\n x: Tensor of shape (N, M, U) containing U-dimensional features to\n be gathered.\n idx: LongTensor of shape (N, L, K) giving the indices returned by `knn_points`.\n lengths: LongTensor of shape (N,) of values in the range [0, M], giving the\n length of each example in the batch in x. Or None to indicate that every\n example has length M.\n Returns:\n x_out: Tensor of shape (N, L, K, U) resulting from gathering the elements of x\n with idx, s.t. `x_out[n, l, k] = x[n, idx[n, l, k]]`.\n If `k > lengths[n]` then `x_out[n, l, k]` is filled with 0.0.\n \"\"\"\n N, M, U = x.shape\n _N, L, K = idx.shape\n\n if N != _N:\n raise ValueError(\"x and idx must have same batch dimension.\")\n\n if lengths is None:\n lengths = torch.full((x.shape[0],), M, dtype=torch.int64, device=x.device)\n\n idx_expanded = idx[:, :, :, None].expand(-1, -1, -1, U)\n # idx_expanded has shape [N, L, K, U]\n\n x_out = x[:, :, None].expand(-1, -1, K, -1).gather(1, idx_expanded)\n # p2_nn has shape [N, L, K, U]\n\n needs_mask = lengths.min() < K\n if needs_mask:\n # mask has shape [N, K], true where idx is irrelevant because\n # there is less number of points in p2 than K\n mask = lengths[:, None] <= torch.arange(K, device=x.device)[None]\n\n # expand mask to shape [N, L, K, U]\n mask = mask[:, None].expand(-1, L, -1)\n mask = mask[:, :, :, None].expand(-1, -1, -1, U)\n x_out[mask] = 0.0\n\n return x_out\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/setup.py", "content": "import glob\nimport os\n\nimport torch\nfrom setuptools import find_packages, setup\nfrom torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension\n\nrequirements = [\"torch\", \"torchvision\"]\n\n\ndef get_extensions():\n this_dir = os.path.dirname(os.path.abspath(__file__))\n extensions_dir = os.path.join(this_dir, \"src\")\n\n main_file = glob.glob(os.path.join(extensions_dir, \"*.cpp\"))\n source_cpu = glob.glob(os.path.join(extensions_dir, \"*.cpp\"))\n source_cuda = glob.glob(os.path.join(extensions_dir, \"*.cu\"))\n\n sources = main_file + source_cpu\n extension = CppExtension\n extra_compile_args = {\"cxx\": [\"-O3\"]}\n define_macros = []\n\n if torch.cuda.is_available() and CUDA_HOME is not None:\n extension = CUDAExtension\n sources += source_cuda\n define_macros += [(\"WITH_CUDA\", None)]\n extra_compile_args[\"nvcc\"] = [\n \"-O3\",\n ]\n else:\n raise NotImplementedError(\"Cuda is not available\")\n\n sources = list(set([os.path.join(extensions_dir, s) for s in sources]))\n include_dirs = [extensions_dir]\n ext_modules = [\n extension(\n \"KNN\",\n sources,\n include_dirs=include_dirs,\n define_macros=define_macros,\n extra_compile_args=extra_compile_args,\n )\n ]\n\n return ext_modules\n\n\nsetup(\n name=\"KNN\",\n version=\"0.1\",\n author=\"Luigi Piccinelli\",\n ext_modules=get_extensions(),\n packages=find_packages(\n exclude=(\n \"configs\",\n \"tests\",\n )\n ),\n cmdclass={\"build_ext\": torch.utils.cpp_extension.BuildExtension},\n)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/knn.cu", "content": "/*\n * Copyright (c) Meta Platforms, Inc. and affiliates.\n * All rights reserved.\n *\n * This source code is licensed under the BSD-style license found in the\n * LICENSE file in the root directory of this source tree.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"utils/dispatch.cuh\"\n#include \"utils/mink.cuh\"\n\n// A chunk of work is blocksize-many points of P1.\n// The number of potential chunks to do is N*(1+(P1-1)/blocksize)\n// call (1+(P1-1)/blocksize) chunks_per_cloud\n// These chunks are divided among the gridSize-many blocks.\n// In block b, we work on chunks b, b+gridSize, b+2*gridSize etc .\n// In chunk i, we work on cloud i/chunks_per_cloud on points starting from\n// blocksize*(i%chunks_per_cloud).\n\ntemplate \n__global__ void KNearestNeighborKernelV0(\n const scalar_t* __restrict__ points1,\n const scalar_t* __restrict__ points2,\n const int64_t* __restrict__ lengths1,\n const int64_t* __restrict__ lengths2,\n scalar_t* __restrict__ dists,\n int64_t* __restrict__ idxs,\n const size_t N,\n const size_t P1,\n const size_t P2,\n const size_t D,\n const size_t K,\n const size_t norm) {\n // Store both dists and indices for knn in global memory.\n const int64_t chunks_per_cloud = (1 + (P1 - 1) / blockDim.x);\n const int64_t chunks_to_do = N * chunks_per_cloud;\n for (int64_t chunk = blockIdx.x; chunk < chunks_to_do; chunk += gridDim.x) {\n const int64_t n = chunk / chunks_per_cloud;\n const int64_t start_point = blockDim.x * (chunk % chunks_per_cloud);\n int64_t p1 = start_point + threadIdx.x;\n if (p1 >= lengths1[n])\n continue;\n int offset = n * P1 * K + p1 * K;\n int64_t length2 = lengths2[n];\n MinK mink(dists + offset, idxs + offset, K);\n for (int p2 = 0; p2 < length2; ++p2) {\n // Find the distance between points1[n, p1] and points[n, p2]\n scalar_t dist = 0;\n for (int d = 0; d < D; ++d) {\n scalar_t coord1 = points1[n * P1 * D + p1 * D + d];\n scalar_t coord2 = points2[n * P2 * D + p2 * D + d];\n scalar_t diff = coord1 - coord2;\n scalar_t norm_diff = (norm == 2) ? (diff * diff) : abs(diff);\n dist += norm_diff;\n }\n mink.add(dist, p2);\n }\n }\n}\n\ntemplate \n__global__ void KNearestNeighborKernelV1(\n const scalar_t* __restrict__ points1,\n const scalar_t* __restrict__ points2,\n const int64_t* __restrict__ lengths1,\n const int64_t* __restrict__ lengths2,\n scalar_t* __restrict__ dists,\n int64_t* __restrict__ idxs,\n const size_t N,\n const size_t P1,\n const size_t P2,\n const size_t K,\n const size_t norm) {\n // Same idea as the previous version, but hoist D into a template argument\n // so we can cache the current point in a thread-local array. We still store\n // the current best K dists and indices in global memory, so this should work\n // for very large K and fairly large D.\n scalar_t cur_point[D];\n const int64_t chunks_per_cloud = (1 + (P1 - 1) / blockDim.x);\n const int64_t chunks_to_do = N * chunks_per_cloud;\n for (int64_t chunk = blockIdx.x; chunk < chunks_to_do; chunk += gridDim.x) {\n const int64_t n = chunk / chunks_per_cloud;\n const int64_t start_point = blockDim.x * (chunk % chunks_per_cloud);\n int64_t p1 = start_point + threadIdx.x;\n if (p1 >= lengths1[n])\n continue;\n for (int d = 0; d < D; ++d) {\n cur_point[d] = points1[n * P1 * D + p1 * D + d];\n }\n int offset = n * P1 * K + p1 * K;\n int64_t length2 = lengths2[n];\n MinK mink(dists + offset, idxs + offset, K);\n for (int p2 = 0; p2 < length2; ++p2) {\n // Find the distance between cur_point and points[n, p2]\n scalar_t dist = 0;\n for (int d = 0; d < D; ++d) {\n scalar_t diff = cur_point[d] - points2[n * P2 * D + p2 * D + d];\n scalar_t norm_diff = (norm == 2) ? (diff * diff) : abs(diff);\n dist += norm_diff;\n }\n mink.add(dist, p2);\n }\n }\n}\n\n// This is a shim functor to allow us to dispatch using DispatchKernel1D\ntemplate \nstruct KNearestNeighborV1Functor {\n static void run(\n size_t blocks,\n size_t threads,\n const scalar_t* __restrict__ points1,\n const scalar_t* __restrict__ points2,\n const int64_t* __restrict__ lengths1,\n const int64_t* __restrict__ lengths2,\n scalar_t* __restrict__ dists,\n int64_t* __restrict__ idxs,\n const size_t N,\n const size_t P1,\n const size_t P2,\n const size_t K,\n const size_t norm) {\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n KNearestNeighborKernelV1<<>>(\n points1, points2, lengths1, lengths2, dists, idxs, N, P1, P2, K, norm);\n }\n};\n\ntemplate \n__global__ void KNearestNeighborKernelV2(\n const scalar_t* __restrict__ points1,\n const scalar_t* __restrict__ points2,\n const int64_t* __restrict__ lengths1,\n const int64_t* __restrict__ lengths2,\n scalar_t* __restrict__ dists,\n int64_t* __restrict__ idxs,\n const int64_t N,\n const int64_t P1,\n const int64_t P2,\n const size_t norm) {\n // Same general implementation as V2, but also hoist K into a template arg.\n scalar_t cur_point[D];\n scalar_t min_dists[K];\n int min_idxs[K];\n const int64_t chunks_per_cloud = (1 + (P1 - 1) / blockDim.x);\n const int64_t chunks_to_do = N * chunks_per_cloud;\n for (int64_t chunk = blockIdx.x; chunk < chunks_to_do; chunk += gridDim.x) {\n const int64_t n = chunk / chunks_per_cloud;\n const int64_t start_point = blockDim.x * (chunk % chunks_per_cloud);\n int64_t p1 = start_point + threadIdx.x;\n if (p1 >= lengths1[n])\n continue;\n for (int d = 0; d < D; ++d) {\n cur_point[d] = points1[n * P1 * D + p1 * D + d];\n }\n int64_t length2 = lengths2[n];\n MinK mink(min_dists, min_idxs, K);\n for (int p2 = 0; p2 < length2; ++p2) {\n scalar_t dist = 0;\n for (int d = 0; d < D; ++d) {\n int offset = n * P2 * D + p2 * D + d;\n scalar_t diff = cur_point[d] - points2[offset];\n scalar_t norm_diff = (norm == 2) ? (diff * diff) : abs(diff);\n dist += norm_diff;\n }\n mink.add(dist, p2);\n }\n for (int k = 0; k < mink.size(); ++k) {\n idxs[n * P1 * K + p1 * K + k] = min_idxs[k];\n dists[n * P1 * K + p1 * K + k] = min_dists[k];\n }\n }\n}\n\n// This is a shim so we can dispatch using DispatchKernel2D\ntemplate \nstruct KNearestNeighborKernelV2Functor {\n static void run(\n size_t blocks,\n size_t threads,\n const scalar_t* __restrict__ points1,\n const scalar_t* __restrict__ points2,\n const int64_t* __restrict__ lengths1,\n const int64_t* __restrict__ lengths2,\n scalar_t* __restrict__ dists,\n int64_t* __restrict__ idxs,\n const int64_t N,\n const int64_t P1,\n const int64_t P2,\n const size_t norm) {\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n KNearestNeighborKernelV2<<>>(\n points1, points2, lengths1, lengths2, dists, idxs, N, P1, P2, norm);\n }\n};\n\ntemplate \n__global__ void KNearestNeighborKernelV3(\n const scalar_t* __restrict__ points1,\n const scalar_t* __restrict__ points2,\n const int64_t* __restrict__ lengths1,\n const int64_t* __restrict__ lengths2,\n scalar_t* __restrict__ dists,\n int64_t* __restrict__ idxs,\n const size_t N,\n const size_t P1,\n const size_t P2,\n const size_t norm) {\n // Same idea as V2, but use register indexing for thread-local arrays.\n // Enabling sorting for this version leads to huge slowdowns; I suspect\n // that it forces min_dists into local memory rather than registers.\n // As a result this version is always unsorted.\n scalar_t cur_point[D];\n scalar_t min_dists[K];\n int min_idxs[K];\n const int64_t chunks_per_cloud = (1 + (P1 - 1) / blockDim.x);\n const int64_t chunks_to_do = N * chunks_per_cloud;\n for (int64_t chunk = blockIdx.x; chunk < chunks_to_do; chunk += gridDim.x) {\n const int64_t n = chunk / chunks_per_cloud;\n const int64_t start_point = blockDim.x * (chunk % chunks_per_cloud);\n int64_t p1 = start_point + threadIdx.x;\n if (p1 >= lengths1[n])\n continue;\n for (int d = 0; d < D; ++d) {\n cur_point[d] = points1[n * P1 * D + p1 * D + d];\n }\n int64_t length2 = lengths2[n];\n RegisterMinK mink(min_dists, min_idxs);\n for (int p2 = 0; p2 < length2; ++p2) {\n scalar_t dist = 0;\n for (int d = 0; d < D; ++d) {\n int offset = n * P2 * D + p2 * D + d;\n scalar_t diff = cur_point[d] - points2[offset];\n scalar_t norm_diff = (norm == 2) ? (diff * diff) : abs(diff);\n dist += norm_diff;\n }\n mink.add(dist, p2);\n }\n for (int k = 0; k < mink.size(); ++k) {\n idxs[n * P1 * K + p1 * K + k] = min_idxs[k];\n dists[n * P1 * K + p1 * K + k] = min_dists[k];\n }\n }\n}\n\n// This is a shim so we can dispatch using DispatchKernel2D\ntemplate \nstruct KNearestNeighborKernelV3Functor {\n static void run(\n size_t blocks,\n size_t threads,\n const scalar_t* __restrict__ points1,\n const scalar_t* __restrict__ points2,\n const int64_t* __restrict__ lengths1,\n const int64_t* __restrict__ lengths2,\n scalar_t* __restrict__ dists,\n int64_t* __restrict__ idxs,\n const size_t N,\n const size_t P1,\n const size_t P2,\n const size_t norm) {\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n KNearestNeighborKernelV3<<>>(\n points1, points2, lengths1, lengths2, dists, idxs, N, P1, P2, norm);\n }\n};\n\nconstexpr int V1_MIN_D = 1;\nconstexpr int V1_MAX_D = 32;\n\nconstexpr int V2_MIN_D = 1;\nconstexpr int V2_MAX_D = 8;\nconstexpr int V2_MIN_K = 1;\nconstexpr int V2_MAX_K = 32;\n\nconstexpr int V3_MIN_D = 1;\nconstexpr int V3_MAX_D = 8;\nconstexpr int V3_MIN_K = 1;\nconstexpr int V3_MAX_K = 4;\n\nbool InBounds(const int64_t min, const int64_t x, const int64_t max) {\n return min <= x && x <= max;\n}\n\nbool KnnCheckVersion(int version, const int64_t D, const int64_t K) {\n if (version == 0) {\n return true;\n } else if (version == 1) {\n return InBounds(V1_MIN_D, D, V1_MAX_D);\n } else if (version == 2) {\n return InBounds(V2_MIN_D, D, V2_MAX_D) && InBounds(V2_MIN_K, K, V2_MAX_K);\n } else if (version == 3) {\n return InBounds(V3_MIN_D, D, V3_MAX_D) && InBounds(V3_MIN_K, K, V3_MAX_K);\n }\n return false;\n}\n\nint ChooseVersion(const int64_t D, const int64_t K) {\n for (int version = 3; version >= 1; version--) {\n if (KnnCheckVersion(version, D, K)) {\n return version;\n }\n }\n return 0;\n}\n\nstd::tuple KNearestNeighborIdxCuda(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const int norm,\n const int K,\n int version) {\n // Check inputs are on the same device\n at::TensorArg p1_t{p1, \"p1\", 1}, p2_t{p2, \"p2\", 2},\n lengths1_t{lengths1, \"lengths1\", 3}, lengths2_t{lengths2, \"lengths2\", 4};\n at::CheckedFrom c = \"KNearestNeighborIdxCuda\";\n at::checkAllSameGPU(c, {p1_t, p2_t, lengths1_t, lengths2_t});\n at::checkAllSameType(c, {p1_t, p2_t});\n\n // Set the device for the kernel launch based on the device of the input\n at::cuda::CUDAGuard device_guard(p1.device());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const auto N = p1.size(0);\n const auto P1 = p1.size(1);\n const auto P2 = p2.size(1);\n const auto D = p2.size(2);\n const int64_t K_64 = K;\n\n TORCH_CHECK((norm == 1) || (norm == 2), \"Norm must be 1 or 2.\");\n\n TORCH_CHECK(p1.size(2) == D, \"Point sets must have the same last dimension\");\n auto long_dtype = lengths1.options().dtype(at::kLong);\n auto idxs = at::zeros({N, P1, K}, long_dtype);\n auto dists = at::zeros({N, P1, K}, p1.options());\n\n if (idxs.numel() == 0) {\n AT_CUDA_CHECK(cudaGetLastError());\n return std::make_tuple(idxs, dists);\n }\n\n if (version < 0) {\n version = ChooseVersion(D, K);\n } else if (!KnnCheckVersion(version, D, K)) {\n int new_version = ChooseVersion(D, K);\n std::cout << \"WARNING: Requested KNN version \" << version\n << \" is not compatible with D = \" << D << \"; K = \" << K\n << \". Falling back to version = \" << new_version << std::endl;\n version = new_version;\n }\n\n // At this point we should have a valid version no matter what data the user\n // gave us. But we can check once more to be sure; however this time\n // assert fail since failing at this point means we have a bug in our version\n // selection or checking code.\n AT_ASSERTM(KnnCheckVersion(version, D, K), \"Invalid version\");\n\n const size_t threads = 256;\n const size_t blocks = 256;\n if (version == 0) {\n AT_DISPATCH_FLOATING_TYPES(\n p1.scalar_type(), \"knn_kernel_cuda\", ([&] {\n KNearestNeighborKernelV0<<>>(\n p1.contiguous().data_ptr(),\n p2.contiguous().data_ptr(),\n lengths1.contiguous().data_ptr(),\n lengths2.contiguous().data_ptr(),\n dists.data_ptr(),\n idxs.data_ptr(),\n N,\n P1,\n P2,\n D,\n K,\n norm);\n }));\n } else if (version == 1) {\n AT_DISPATCH_FLOATING_TYPES(p1.scalar_type(), \"knn_kernel_cuda\", ([&] {\n DispatchKernel1D<\n KNearestNeighborV1Functor,\n scalar_t,\n V1_MIN_D,\n V1_MAX_D>(\n D,\n blocks,\n threads,\n p1.contiguous().data_ptr(),\n p2.contiguous().data_ptr(),\n lengths1.contiguous().data_ptr(),\n lengths2.contiguous().data_ptr(),\n dists.data_ptr(),\n idxs.data_ptr(),\n N,\n P1,\n P2,\n K,\n norm);\n }));\n } else if (version == 2) {\n AT_DISPATCH_FLOATING_TYPES(p1.scalar_type(), \"knn_kernel_cuda\", ([&] {\n DispatchKernel2D<\n KNearestNeighborKernelV2Functor,\n scalar_t,\n V2_MIN_D,\n V2_MAX_D,\n V2_MIN_K,\n V2_MAX_K>(\n D,\n K_64,\n blocks,\n threads,\n p1.contiguous().data_ptr(),\n p2.contiguous().data_ptr(),\n lengths1.contiguous().data_ptr(),\n lengths2.contiguous().data_ptr(),\n dists.data_ptr(),\n idxs.data_ptr(),\n N,\n P1,\n P2,\n norm);\n }));\n } else if (version == 3) {\n AT_DISPATCH_FLOATING_TYPES(p1.scalar_type(), \"knn_kernel_cuda\", ([&] {\n DispatchKernel2D<\n KNearestNeighborKernelV3Functor,\n scalar_t,\n V3_MIN_D,\n V3_MAX_D,\n V3_MIN_K,\n V3_MAX_K>(\n D,\n K_64,\n blocks,\n threads,\n p1.contiguous().data_ptr(),\n p2.contiguous().data_ptr(),\n lengths1.contiguous().data_ptr(),\n lengths2.contiguous().data_ptr(),\n dists.data_ptr(),\n idxs.data_ptr(),\n N,\n P1,\n P2,\n norm);\n }));\n }\n AT_CUDA_CHECK(cudaGetLastError());\n return std::make_tuple(idxs, dists);\n}\n\n// ------------------------------------------------------------- //\n// Backward Operators //\n// ------------------------------------------------------------- //\n\n// TODO(gkioxari) support all data types once AtomicAdd supports doubles.\n// Currently, support is for floats only.\n__global__ void KNearestNeighborBackwardKernel(\n const float* __restrict__ p1, // (N, P1, D)\n const float* __restrict__ p2, // (N, P2, D)\n const int64_t* __restrict__ lengths1, // (N,)\n const int64_t* __restrict__ lengths2, // (N,)\n const int64_t* __restrict__ idxs, // (N, P1, K)\n const float* __restrict__ grad_dists, // (N, P1, K)\n float* __restrict__ grad_p1, // (N, P1, D)\n float* __restrict__ grad_p2, // (N, P2, D)\n const size_t N,\n const size_t P1,\n const size_t P2,\n const size_t K,\n const size_t D,\n const size_t norm) {\n const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n const size_t stride = gridDim.x * blockDim.x;\n\n for (size_t i = tid; i < N * P1 * K * D; i += stride) {\n const size_t n = i / (P1 * K * D); // batch index\n size_t rem = i % (P1 * K * D);\n const size_t p1_idx = rem / (K * D); // index of point in p1\n rem = rem % (K * D);\n const size_t k = rem / D; // k-th nearest neighbor\n const size_t d = rem % D; // d-th dimension in the feature vector\n\n const size_t num1 = lengths1[n]; // number of valid points in p1 in batch\n const size_t num2 = lengths2[n]; // number of valid points in p2 in batch\n if ((p1_idx < num1) && (k < num2)) {\n const float grad_dist = grad_dists[n * P1 * K + p1_idx * K + k];\n // index of point in p2 corresponding to the k-th nearest neighbor\n const int64_t p2_idx = idxs[n * P1 * K + p1_idx * K + k];\n // If the index is the pad value of -1 then ignore it\n if (p2_idx == -1) {\n continue;\n }\n float diff = 0.0;\n if (norm == 1) {\n float sign =\n (p1[n * P1 * D + p1_idx * D + d] > p2[n * P2 * D + p2_idx * D + d])\n ? 1.0\n : -1.0;\n diff = grad_dist * sign;\n } else { // norm is 2\n diff = 2.0 * grad_dist *\n (p1[n * P1 * D + p1_idx * D + d] - p2[n * P2 * D + p2_idx * D + d]);\n }\n atomicAdd(grad_p1 + n * P1 * D + p1_idx * D + d, diff);\n atomicAdd(grad_p2 + n * P2 * D + p2_idx * D + d, -1.0f * diff);\n }\n }\n}\n\nstd::tuple KNearestNeighborBackwardCuda(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const at::Tensor& idxs,\n int norm,\n const at::Tensor& grad_dists) {\n // Check inputs are on the same device\n at::TensorArg p1_t{p1, \"p1\", 1}, p2_t{p2, \"p2\", 2},\n lengths1_t{lengths1, \"lengths1\", 3}, lengths2_t{lengths2, \"lengths2\", 4},\n idxs_t{idxs, \"idxs\", 5}, grad_dists_t{grad_dists, \"grad_dists\", 6};\n at::CheckedFrom c = \"KNearestNeighborBackwardCuda\";\n at::checkAllSameGPU(\n c, {p1_t, p2_t, lengths1_t, lengths2_t, idxs_t, grad_dists_t});\n at::checkAllSameType(c, {p1_t, p2_t, grad_dists_t});\n\n // This is nondeterministic because atomicAdd\n at::globalContext().alertNotDeterministic(\"KNearestNeighborBackwardCuda\");\n\n // Set the device for the kernel launch based on the device of the input\n at::cuda::CUDAGuard device_guard(p1.device());\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n const auto N = p1.size(0);\n const auto P1 = p1.size(1);\n const auto P2 = p2.size(1);\n const auto D = p2.size(2);\n const auto K = idxs.size(2);\n\n TORCH_CHECK(p1.size(2) == D, \"Point sets must have the same last dimension\");\n TORCH_CHECK(idxs.size(0) == N, \"KNN idxs must have the same batch dimension\");\n TORCH_CHECK(\n idxs.size(1) == P1, \"KNN idxs must have the same point dimension as p1\");\n TORCH_CHECK(grad_dists.size(0) == N);\n TORCH_CHECK(grad_dists.size(1) == P1);\n TORCH_CHECK(grad_dists.size(2) == K);\n\n auto grad_p1 = at::zeros({N, P1, D}, p1.options());\n auto grad_p2 = at::zeros({N, P2, D}, p2.options());\n\n if (grad_p1.numel() == 0 || grad_p2.numel() == 0) {\n AT_CUDA_CHECK(cudaGetLastError());\n return std::make_tuple(grad_p1, grad_p2);\n }\n\n const int blocks = 64;\n const int threads = 512;\n\n KNearestNeighborBackwardKernel<<>>(\n p1.contiguous().data_ptr(),\n p2.contiguous().data_ptr(),\n lengths1.contiguous().data_ptr(),\n lengths2.contiguous().data_ptr(),\n idxs.contiguous().data_ptr(),\n grad_dists.contiguous().data_ptr(),\n grad_p1.data_ptr(),\n grad_p2.data_ptr(),\n N,\n P1,\n P2,\n K,\n D,\n norm);\n\n AT_CUDA_CHECK(cudaGetLastError());\n return std::make_tuple(grad_p1, grad_p2);\n}" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/knn.h", "content": "/*\n * Copyright (c) Meta Platforms, Inc. and affiliates.\n * All rights reserved.\n *\n * This source code is licensed under the BSD-style license found in the\n * LICENSE file in the root directory of this source tree.\n */\n\n#pragma once\n#include \n#include \n#include \"utils/pytorch3d_cutils.h\"\n\n// Compute indices of K nearest neighbors in pointcloud p2 to points\n// in pointcloud p1.\n//\n// Args:\n// p1: FloatTensor of shape (N, P1, D) giving a batch of pointclouds each\n// containing P1 points of dimension D.\n// p2: FloatTensor of shape (N, P2, D) giving a batch of pointclouds each\n// containing P2 points of dimension D.\n// lengths1: LongTensor, shape (N,), giving actual length of each P1 cloud.\n// lengths2: LongTensor, shape (N,), giving actual length of each P2 cloud.\n// norm: int specifying the norm for the distance (1 for L1, 2 for L2)\n// K: int giving the number of nearest points to return.\n// version: Integer telling which implementation to use.\n//\n// Returns:\n// p1_neighbor_idx: LongTensor of shape (N, P1, K), where\n// p1_neighbor_idx[n, i, k] = j means that the kth nearest\n// neighbor to p1[n, i] in the cloud p2[n] is p2[n, j].\n// It is padded with zeros so that it can be used easily in a later\n// gather() operation.\n//\n// p1_neighbor_dists: FloatTensor of shape (N, P1, K) containing the squared\n// distance from each point p1[n, p, :] to its K neighbors\n// p2[n, p1_neighbor_idx[n, p, k], :].\n\n// CPU implementation.\nstd::tuple KNearestNeighborIdxCpu(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const int norm,\n const int K);\n\n// CUDA implementation\nstd::tuple KNearestNeighborIdxCuda(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const int norm,\n const int K,\n const int version);\n\n// Implementation which is exposed.\nstd::tuple KNearestNeighborIdx(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const int norm,\n const int K,\n const int version) {\n if (p1.is_cuda() || p2.is_cuda()) {\n#ifdef WITH_CUDA\n CHECK_CUDA(p1);\n CHECK_CUDA(p2);\n return KNearestNeighborIdxCuda(\n p1, p2, lengths1, lengths2, norm, K, version);\n#else\n AT_ERROR(\"Not compiled with GPU support.\");\n#endif\n }\n return KNearestNeighborIdxCpu(p1, p2, lengths1, lengths2, norm, K);\n}\n\n// Compute gradients with respect to p1 and p2\n//\n// Args:\n// p1: FloatTensor of shape (N, P1, D) giving a batch of pointclouds each\n// containing P1 points of dimension D.\n// p2: FloatTensor of shape (N, P2, D) giving a batch of pointclouds each\n// containing P2 points of dimension D.\n// lengths1: LongTensor, shape (N,), giving actual length of each P1 cloud.\n// lengths2: LongTensor, shape (N,), giving actual length of each P2 cloud.\n// p1_neighbor_idx: LongTensor of shape (N, P1, K), where\n// p1_neighbor_idx[n, i, k] = j means that the kth nearest\n// neighbor to p1[n, i] in the cloud p2[n] is p2[n, j].\n// It is padded with zeros so that it can be used easily in a later\n// gather() operation. This is computed from the forward pass.\n// norm: int specifying the norm for the distance (1 for L1, 2 for L2)\n// grad_dists: FLoatTensor of shape (N, P1, K) which contains the input\n// gradients.\n//\n// Returns:\n// grad_p1: FloatTensor of shape (N, P1, D) containing the output gradients\n// wrt p1.\n// grad_p2: FloatTensor of shape (N, P2, D) containing the output gradients\n// wrt p2.\n\n// CPU implementation.\nstd::tuple KNearestNeighborBackwardCpu(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const at::Tensor& idxs,\n const int norm,\n const at::Tensor& grad_dists);\n\n// CUDA implementation\nstd::tuple KNearestNeighborBackwardCuda(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const at::Tensor& idxs,\n const int norm,\n const at::Tensor& grad_dists);\n\n// Implementation which is exposed.\nstd::tuple KNearestNeighborBackward(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const at::Tensor& idxs,\n const int norm,\n const at::Tensor& grad_dists) {\n if (p1.is_cuda() || p2.is_cuda()) {\n#ifdef WITH_CUDA\n CHECK_CUDA(p1);\n CHECK_CUDA(p2);\n return KNearestNeighborBackwardCuda(\n p1, p2, lengths1, lengths2, idxs, norm, grad_dists);\n#else\n AT_ERROR(\"Not compiled with GPU support.\");\n#endif\n }\n return KNearestNeighborBackwardCpu(\n p1, p2, lengths1, lengths2, idxs, norm, grad_dists);\n}\n\n// Utility to check whether a KNN version can be used.\n//\n// Args:\n// version: Integer in the range 0 <= version <= 3 indicating one of our\n// KNN implementations.\n// D: Number of dimensions for the input and query point clouds\n// K: Number of neighbors to be found\n//\n// Returns:\n// Whether the indicated KNN version can be used.\nbool KnnCheckVersion(int version, const int64_t D, const int64_t K);" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/knn_cpu.cpp", "content": "/*\n * Copyright (c) Meta Platforms, Inc. and affiliates.\n * All rights reserved.\n *\n * This source code is licensed under the BSD-style license found in the\n * LICENSE file in the root directory of this source tree.\n */\n\n#include \n#include \n#include \n\nstd::tuple KNearestNeighborIdxCpu(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const int norm,\n const int K) {\n const int N = p1.size(0);\n const int P1 = p1.size(1);\n const int D = p1.size(2);\n\n auto long_opts = lengths1.options().dtype(torch::kInt64);\n torch::Tensor idxs = torch::full({N, P1, K}, 0, long_opts);\n torch::Tensor dists = torch::full({N, P1, K}, 0, p1.options());\n\n auto p1_a = p1.accessor();\n auto p2_a = p2.accessor();\n auto lengths1_a = lengths1.accessor();\n auto lengths2_a = lengths2.accessor();\n auto idxs_a = idxs.accessor();\n auto dists_a = dists.accessor();\n\n for (int n = 0; n < N; ++n) {\n const int64_t length1 = lengths1_a[n];\n const int64_t length2 = lengths2_a[n];\n for (int64_t i1 = 0; i1 < length1; ++i1) {\n // Use a priority queue to store (distance, index) tuples.\n std::priority_queue> q;\n for (int64_t i2 = 0; i2 < length2; ++i2) {\n float dist = 0;\n for (int d = 0; d < D; ++d) {\n float diff = p1_a[n][i1][d] - p2_a[n][i2][d];\n if (norm == 1) {\n dist += abs(diff);\n } else { // norm is 2 (default)\n dist += diff * diff;\n }\n }\n int size = static_cast(q.size());\n if (size < K || dist < std::get<0>(q.top())) {\n q.emplace(dist, i2);\n if (size >= K) {\n q.pop();\n }\n }\n }\n while (!q.empty()) {\n auto t = q.top();\n q.pop();\n const int k = q.size();\n dists_a[n][i1][k] = std::get<0>(t);\n idxs_a[n][i1][k] = std::get<1>(t);\n }\n }\n }\n return std::make_tuple(idxs, dists);\n}\n\n// ------------------------------------------------------------- //\n// Backward Operators //\n// ------------------------------------------------------------- //\n\nstd::tuple KNearestNeighborBackwardCpu(\n const at::Tensor& p1,\n const at::Tensor& p2,\n const at::Tensor& lengths1,\n const at::Tensor& lengths2,\n const at::Tensor& idxs,\n const int norm,\n const at::Tensor& grad_dists) {\n const int N = p1.size(0);\n const int P1 = p1.size(1);\n const int D = p1.size(2);\n const int P2 = p2.size(1);\n const int K = idxs.size(2);\n\n torch::Tensor grad_p1 = torch::full({N, P1, D}, 0, p1.options());\n torch::Tensor grad_p2 = torch::full({N, P2, D}, 0, p2.options());\n\n auto p1_a = p1.accessor();\n auto p2_a = p2.accessor();\n auto lengths1_a = lengths1.accessor();\n auto lengths2_a = lengths2.accessor();\n auto idxs_a = idxs.accessor();\n auto grad_dists_a = grad_dists.accessor();\n auto grad_p1_a = grad_p1.accessor();\n auto grad_p2_a = grad_p2.accessor();\n\n for (int n = 0; n < N; ++n) {\n const int64_t length1 = lengths1_a[n];\n int64_t length2 = lengths2_a[n];\n length2 = (length2 < K) ? length2 : K;\n for (int64_t i1 = 0; i1 < length1; ++i1) {\n for (int64_t k = 0; k < length2; ++k) {\n const int64_t i2 = idxs_a[n][i1][k];\n // If the index is the pad value of -1 then ignore it\n if (i2 == -1) {\n continue;\n }\n for (int64_t d = 0; d < D; ++d) {\n float diff = 0.0;\n if (norm == 1) {\n float sign = (p1_a[n][i1][d] > p2_a[n][i2][d]) ? 1.0 : -1.0;\n diff = grad_dists_a[n][i1][k] * sign;\n } else { // norm is 2 (default)\n diff = 2.0f * grad_dists_a[n][i1][k] *\n (p1_a[n][i1][d] - p2_a[n][i2][d]);\n }\n grad_p1_a[n][i1][d] += diff;\n grad_p2_a[n][i2][d] += -1.0f * diff;\n }\n }\n }\n }\n return std::make_tuple(grad_p1, grad_p2);\n}" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/knn_ext.cpp", "content": "#include \n#include \"knn.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n#ifdef WITH_CUDA\n m.def(\"knn_check_version\", &KnnCheckVersion);\n#endif\n m.def(\"knn_points_idx\", &KNearestNeighborIdx);\n m.def(\"knn_points_backward\", &KNearestNeighborBackward);\n}" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/utils/dispatch.cuh", "content": "/*\n * Copyright (c) Meta Platforms, Inc. and affiliates.\n * All rights reserved.\n *\n * This source code is licensed under the BSD-style license found in the\n * LICENSE file in the root directory of this source tree.\n */\n\n// This file provides utilities for dispatching to specialized versions of\n// functions. This is especially useful for CUDA kernels, since specializing\n// them to particular input sizes can often allow the compiler to unroll loops\n// and place arrays into registers, which can give huge performance speedups.\n//\n// As an example, suppose we have the following function which is specialized\n// based on a compile-time int64_t value:\n//\n// template\n// struct SquareOffset {\n// static void run(T y) {\n// T val = x * x + y;\n// std::cout << val << std::endl;\n// }\n// }\n//\n// This function takes one compile-time argument x, and one run-time argument y.\n// We might want to compile specialized versions of this for x=0, x=1, etc and\n// then dispatch to the correct one based on the runtime value of x.\n// One simple way to achieve this is with a lookup table:\n//\n// template\n// void DispatchSquareOffset(const int64_t x, T y) {\n// if (x == 0) {\n// SquareOffset::run(y);\n// } else if (x == 1) {\n// SquareOffset::run(y);\n// } else if (x == 2) {\n// SquareOffset::run(y);\n// }\n// }\n//\n// This function takes both x and y as run-time arguments, and dispatches to\n// different specialized versions of SquareOffset based on the run-time value\n// of x. This works, but it's tedious and error-prone. If we want to change the\n// set of x values for which we provide compile-time specializations, then we\n// will need to do a lot of tedius editing of the dispatch function. Also, if we\n// want to provide compile-time specializations for another function other than\n// SquareOffset, we will need to duplicate the entire lookup table.\n//\n// To solve these problems, we can use the DispatchKernel1D function provided by\n// this file instead:\n//\n// template\n// void DispatchSquareOffset(const int64_t x, T y) {\n// constexpr int64_t xmin = 0;\n// constexpr int64_t xmax = 2;\n// DispatchKernel1D(x, y);\n// }\n//\n// DispatchKernel1D uses template metaprogramming to compile specialized\n// versions of SquareOffset for all values of x with xmin <= x <= xmax, and\n// then dispatches to the correct one based on the run-time value of x. If we\n// want to change the range of x values for which SquareOffset is specialized\n// at compile-time, then all we have to do is change the values of the\n// compile-time constants xmin and xmax.\n//\n// This file also allows us to similarly dispatch functions that depend on two\n// compile-time int64_t values, using the DispatchKernel2D function like this:\n//\n// template\n// struct Sum {\n// static void run(T z, T w) {\n// T val = x + y + z + w;\n// std::cout << val << std::endl;\n// }\n// }\n//\n// template\n// void DispatchSum(const int64_t x, const int64_t y, int z, int w) {\n// constexpr int64_t xmin = 1;\n// constexpr int64_t xmax = 3;\n// constexpr int64_t ymin = 2;\n// constexpr int64_t ymax = 5;\n// DispatchKernel2D(x, y, z, w);\n// }\n//\n// Like its 1D counterpart, DispatchKernel2D uses template metaprogramming to\n// compile specialized versions of sum for all values of (x, y) with\n// xmin <= x <= xmax and ymin <= y <= ymax, then dispatches to the correct\n// specialized version based on the runtime values of x and y.\n\n// Define some helper structs in an anonymous namespace.\nnamespace {\n\n// 1D dispatch: general case.\n// Kernel is the function we want to dispatch to; it should take a typename and\n// an int64_t as template args, and it should define a static void function\n// run which takes any number of arguments of any type.\n// In order to dispatch, we will take an additional template argument curN,\n// and increment it via template recursion until it is equal to the run-time\n// argument N.\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n int64_t curN,\n typename... Args>\nstruct DispatchKernelHelper1D {\n static void run(const int64_t N, Args... args) {\n if (curN == N) {\n // The compile-time value curN is equal to the run-time value N, so we\n // can dispatch to the run method of the Kernel.\n Kernel::run(args...);\n } else if (curN < N) {\n // Increment curN via template recursion\n DispatchKernelHelper1D::run(\n N, args...);\n }\n // We shouldn't get here -- throw an error?\n }\n};\n\n// 1D dispatch: Specialization when curN == maxN\n// We need this base case to avoid infinite template recursion.\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n typename... Args>\nstruct DispatchKernelHelper1D {\n static void run(const int64_t N, Args... args) {\n if (N == maxN) {\n Kernel::run(args...);\n }\n // We shouldn't get here -- throw an error?\n }\n};\n\n// 2D dispatch, general case.\n// This is similar to the 1D case: we take additional template args curN and\n// curM, and increment them via template recursion until they are equal to\n// the run-time values of N and M, at which point we dispatch to the run\n// method of the kernel.\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n int64_t curN,\n int64_t minM,\n int64_t maxM,\n int64_t curM,\n typename... Args>\nstruct DispatchKernelHelper2D {\n static void run(const int64_t N, const int64_t M, Args... args) {\n if (curN == N && curM == M) {\n Kernel::run(args...);\n } else if (curN < N && curM < M) {\n // Increment both curN and curM. This isn't strictly necessary; we could\n // just increment one or the other at each step. But this helps to cut\n // on the number of recursive calls we make.\n DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n curN + 1,\n minM,\n maxM,\n curM + 1,\n Args...>::run(N, M, args...);\n } else if (curN < N) {\n // Increment curN only\n DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n curN + 1,\n minM,\n maxM,\n curM,\n Args...>::run(N, M, args...);\n } else if (curM < M) {\n // Increment curM only\n DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n curN,\n minM,\n maxM,\n curM + 1,\n Args...>::run(N, M, args...);\n }\n }\n};\n\n// 2D dispatch, specialization for curN == maxN\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n int64_t minM,\n int64_t maxM,\n int64_t curM,\n typename... Args>\nstruct DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n maxN,\n minM,\n maxM,\n curM,\n Args...> {\n static void run(const int64_t N, const int64_t M, Args... args) {\n if (maxN == N && curM == M) {\n Kernel::run(args...);\n } else if (curM < maxM) {\n DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n maxN,\n minM,\n maxM,\n curM + 1,\n Args...>::run(N, M, args...);\n }\n // We should not get here -- throw an error?\n }\n};\n\n// 2D dispatch, specialization for curM == maxM\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n int64_t curN,\n int64_t minM,\n int64_t maxM,\n typename... Args>\nstruct DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n curN,\n minM,\n maxM,\n maxM,\n Args...> {\n static void run(const int64_t N, const int64_t M, Args... args) {\n if (curN == N && maxM == M) {\n Kernel::run(args...);\n } else if (curN < maxN) {\n DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n curN + 1,\n minM,\n maxM,\n maxM,\n Args...>::run(N, M, args...);\n }\n // We should not get here -- throw an error?\n }\n};\n\n// 2D dispatch, specialization for curN == maxN, curM == maxM\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n int64_t minM,\n int64_t maxM,\n typename... Args>\nstruct DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n maxN,\n minM,\n maxM,\n maxM,\n Args...> {\n static void run(const int64_t N, const int64_t M, Args... args) {\n if (maxN == N && maxM == M) {\n Kernel::run(args...);\n }\n // We should not get here -- throw an error?\n }\n};\n\n} // namespace\n\n// This is the function we expect users to call to dispatch to 1D functions\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n typename... Args>\nvoid DispatchKernel1D(const int64_t N, Args... args) {\n if (minN <= N && N <= maxN) {\n // Kick off the template recursion by calling the Helper with curN = minN\n DispatchKernelHelper1D::run(\n N, args...);\n }\n // Maybe throw an error if we tried to dispatch outside the allowed range?\n}\n\n// This is the function we expect users to call to dispatch to 2D functions\ntemplate <\n template \n class Kernel,\n typename T,\n int64_t minN,\n int64_t maxN,\n int64_t minM,\n int64_t maxM,\n typename... Args>\nvoid DispatchKernel2D(const int64_t N, const int64_t M, Args... args) {\n if (minN <= N && N <= maxN && minM <= M && M <= maxM) {\n // Kick off the template recursion by calling the Helper with curN = minN\n // and curM = minM\n DispatchKernelHelper2D<\n Kernel,\n T,\n minN,\n maxN,\n minN,\n minM,\n maxM,\n minM,\n Args...>::run(N, M, args...);\n }\n // Maybe throw an error if we tried to dispatch outside the specified range?\n}" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/utils/index_utils.cuh", "content": "/*\n * Copyright (c) Meta Platforms, Inc. and affiliates.\n * All rights reserved.\n *\n * This source code is licensed under the BSD-style license found in the\n * LICENSE file in the root directory of this source tree.\n */\n\n// This converts dynamic array lookups into static array lookups, for small\n// arrays up to size 32.\n//\n// Suppose we have a small thread-local array:\n//\n// float vals[10];\n//\n// Ideally we should only index this array using static indices:\n//\n// for (int i = 0; i < 10; ++i) vals[i] = i * i;\n//\n// If we do so, then the CUDA compiler may be able to place the array into\n// registers, which can have a big performance improvement. However if we\n// access the array dynamically, the the compiler may force the array into\n// local memory, which has the same latency as global memory.\n//\n// These functions convert dynamic array access into static array access\n// using a brute-force lookup table. It can be used like this:\n//\n// float vals[10];\n// int idx = 3;\n// float val = 3.14f;\n// RegisterIndexUtils::set(vals, idx, val);\n// float val2 = RegisterIndexUtils::get(vals, idx);\n//\n// The implementation is based on fbcuda/RegisterUtils.cuh:\n// https://github.com/facebook/fbcuda/blob/master/RegisterUtils.cuh\n// To avoid depending on the entire library, we just reimplement these two\n// functions. The fbcuda implementation is a bit more sophisticated, and uses\n// the preprocessor to generate switch statements that go up to N for each\n// value of N. We are lazy and just have a giant explicit switch statement.\n//\n// We might be able to use a template metaprogramming approach similar to\n// DispatchKernel1D for this. However DispatchKernel1D is intended to be used\n// for dispatching to the correct CUDA kernel on the host, while this is\n// is intended to run on the device. I was concerned that a metaprogramming\n// approach for this might lead to extra function calls at runtime if the\n// compiler fails to optimize them away, which could be very slow on device.\n// However I didn't actually benchmark or test this.\ntemplate \nstruct RegisterIndexUtils {\n __device__ __forceinline__ static T get(const T arr[N], int idx) {\n if (idx < 0 || idx >= N)\n return T();\n switch (idx) {\n case 0:\n return arr[0];\n case 1:\n return arr[1];\n case 2:\n return arr[2];\n case 3:\n return arr[3];\n case 4:\n return arr[4];\n case 5:\n return arr[5];\n case 6:\n return arr[6];\n case 7:\n return arr[7];\n case 8:\n return arr[8];\n case 9:\n return arr[9];\n case 10:\n return arr[10];\n case 11:\n return arr[11];\n case 12:\n return arr[12];\n case 13:\n return arr[13];\n case 14:\n return arr[14];\n case 15:\n return arr[15];\n case 16:\n return arr[16];\n case 17:\n return arr[17];\n case 18:\n return arr[18];\n case 19:\n return arr[19];\n case 20:\n return arr[20];\n case 21:\n return arr[21];\n case 22:\n return arr[22];\n case 23:\n return arr[23];\n case 24:\n return arr[24];\n case 25:\n return arr[25];\n case 26:\n return arr[26];\n case 27:\n return arr[27];\n case 28:\n return arr[28];\n case 29:\n return arr[29];\n case 30:\n return arr[30];\n case 31:\n return arr[31];\n };\n return T();\n }\n\n __device__ __forceinline__ static void set(T arr[N], int idx, T val) {\n if (idx < 0 || idx >= N)\n return;\n switch (idx) {\n case 0:\n arr[0] = val;\n break;\n case 1:\n arr[1] = val;\n break;\n case 2:\n arr[2] = val;\n break;\n case 3:\n arr[3] = val;\n break;\n case 4:\n arr[4] = val;\n break;\n case 5:\n arr[5] = val;\n break;\n case 6:\n arr[6] = val;\n break;\n case 7:\n arr[7] = val;\n break;\n case 8:\n arr[8] = val;\n break;\n case 9:\n arr[9] = val;\n break;\n case 10:\n arr[10] = val;\n break;\n case 11:\n arr[11] = val;\n break;\n case 12:\n arr[12] = val;\n break;\n case 13:\n arr[13] = val;\n break;\n case 14:\n arr[14] = val;\n break;\n case 15:\n arr[15] = val;\n break;\n case 16:\n arr[16] = val;\n break;\n case 17:\n arr[17] = val;\n break;\n case 18:\n arr[18] = val;\n break;\n case 19:\n arr[19] = val;\n break;\n case 20:\n arr[20] = val;\n break;\n case 21:\n arr[21] = val;\n break;\n case 22:\n arr[22] = val;\n break;\n case 23:\n arr[23] = val;\n break;\n case 24:\n arr[24] = val;\n break;\n case 25:\n arr[25] = val;\n break;\n case 26:\n arr[26] = val;\n break;\n case 27:\n arr[27] = val;\n break;\n case 28:\n arr[28] = val;\n break;\n case 29:\n arr[29] = val;\n break;\n case 30:\n arr[30] = val;\n break;\n case 31:\n arr[31] = val;\n break;\n }\n }\n};" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/utils/mink.cuh", "content": "/*\n * Copyright (c) Meta Platforms, Inc. and affiliates.\n * All rights reserved.\n *\n * This source code is licensed under the BSD-style license found in the\n * LICENSE file in the root directory of this source tree.\n */\n\n#pragma once\n#define MINK_H\n\n#include \"index_utils.cuh\"\n\n// A data structure to keep track of the smallest K keys seen so far as well\n// as their associated values, intended to be used in device code.\n// This data structure doesn't allocate any memory; keys and values are stored\n// in arrays passed to the constructor.\n//\n// The implementation is generic; it can be used for any key type that supports\n// the < operator, and can be used with any value type.\n//\n// Example usage:\n//\n// float keys[K];\n// int values[K];\n// MinK mink(keys, values, K);\n// for (...) {\n// // Produce some key and value from somewhere\n// mink.add(key, value);\n// }\n// mink.sort();\n//\n// Now keys and values store the smallest K keys seen so far and the values\n// associated to these keys:\n//\n// for (int k = 0; k < K; ++k) {\n// float key_k = keys[k];\n// int value_k = values[k];\n// }\ntemplate \nclass MinK {\n public:\n // Constructor.\n //\n // Arguments:\n // keys: Array in which to store keys\n // values: Array in which to store values\n // K: How many values to keep track of\n __device__ MinK(key_t* keys, value_t* vals, int K)\n : keys(keys), vals(vals), K(K), _size(0) {}\n\n // Try to add a new key and associated value to the data structure. If the key\n // is one of the smallest K seen so far then it will be kept; otherwise it\n // it will not be kept.\n //\n // This takes O(1) operations if the new key is not kept, or if the structure\n // currently contains fewer than K elements. Otherwise this takes O(K) time.\n //\n // Arguments:\n // key: The key to add\n // val: The value associated to the key\n __device__ __forceinline__ void add(const key_t& key, const value_t& val) {\n if (_size < K) {\n keys[_size] = key;\n vals[_size] = val;\n if (_size == 0 || key > max_key) {\n max_key = key;\n max_idx = _size;\n }\n _size++;\n } else if (key < max_key) {\n keys[max_idx] = key;\n vals[max_idx] = val;\n max_key = key;\n for (int k = 0; k < K; ++k) {\n key_t cur_key = keys[k];\n if (cur_key > max_key) {\n max_key = cur_key;\n max_idx = k;\n }\n }\n }\n }\n\n // Get the number of items currently stored in the structure.\n // This takes O(1) time.\n __device__ __forceinline__ int size() {\n return _size;\n }\n\n // Sort the items stored in the structure using bubble sort.\n // This takes O(K^2) time.\n __device__ __forceinline__ void sort() {\n for (int i = 0; i < _size - 1; ++i) {\n for (int j = 0; j < _size - i - 1; ++j) {\n if (keys[j + 1] < keys[j]) {\n key_t key = keys[j];\n value_t val = vals[j];\n keys[j] = keys[j + 1];\n vals[j] = vals[j + 1];\n keys[j + 1] = key;\n vals[j + 1] = val;\n }\n }\n }\n }\n\n private:\n key_t* keys;\n value_t* vals;\n int K;\n int _size;\n key_t max_key;\n int max_idx;\n};\n\n// This is a version of MinK that only touches the arrays using static indexing\n// via RegisterIndexUtils. If the keys and values are stored in thread-local\n// arrays, then this may allow the compiler to place them in registers for\n// fast access.\n//\n// This has the same API as RegisterMinK, but doesn't support sorting.\n// We found that sorting via RegisterIndexUtils gave very poor performance,\n// and suspect it may have prevented the compiler from placing the arrays\n// into registers.\ntemplate \nclass RegisterMinK {\n public:\n __device__ RegisterMinK(key_t* keys, value_t* vals)\n : keys(keys), vals(vals), _size(0) {}\n\n __device__ __forceinline__ void add(const key_t& key, const value_t& val) {\n if (_size < K) {\n RegisterIndexUtils::set(keys, _size, key);\n RegisterIndexUtils::set(vals, _size, val);\n if (_size == 0 || key > max_key) {\n max_key = key;\n max_idx = _size;\n }\n _size++;\n } else if (key < max_key) {\n RegisterIndexUtils::set(keys, max_idx, key);\n RegisterIndexUtils::set(vals, max_idx, val);\n max_key = key;\n for (int k = 0; k < K; ++k) {\n key_t cur_key = RegisterIndexUtils::get(keys, k);\n if (cur_key > max_key) {\n max_key = cur_key;\n max_idx = k;\n }\n }\n }\n }\n\n __device__ __forceinline__ int size() {\n return _size;\n }\n\n private:\n key_t* keys;\n value_t* vals;\n int _size;\n key_t max_key;\n int max_idx;\n};" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/knn/src/utils/pytorch3d_cutils.h", "content": "/*\n * Copyright (c) Meta Platforms, Inc. and affiliates.\n * All rights reserved.\n *\n * This source code is licensed under the BSD-style license found in the\n * LICENSE file in the root directory of this source tree.\n */\n\n#pragma once\n#include \n\n#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x \" must be a CUDA tensor.\")\n#define CHECK_CONTIGUOUS(x) \\\n TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous.\")\n#define CHECK_CONTIGUOUS_CUDA(x) \\\n CHECK_CUDA(x); \\\n CHECK_CONTIGUOUS(x)" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/__init__.py", "content": "from .arel import ARel\nfrom .confidence import Confidence\nfrom .distill import SelfDistill, TeacherDistill\nfrom .dummy import Dummy\nfrom .local_ssi import EdgeGuidedLocalSSI, LocalSSI\nfrom .regression import Regression\nfrom .silog import SILog\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/arel.py", "content": "import torch\nimport torch.nn as nn\n\nfrom .utils import FNS, masked_mean\n\n\nclass ARel(nn.Module):\n def __init__(\n self,\n weight: float,\n output_fn: str = \"sqrt\",\n input_fn: str = \"linear\",\n eps: float = 1e-5,\n ):\n super().__init__()\n self.name: str = self.__class__.__name__\n self.weight: float = weight\n self.dims = [-2, -1]\n self.output_fn = FNS[output_fn]\n self.input_fn = FNS[input_fn]\n self.eps: float = eps\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor, **kwargs\n ) -> torch.Tensor:\n mask = mask.bool().clone()\n\n input = self.input_fn(input.float())\n target = self.input_fn(target.float())\n\n error = (input - target).norm(dim=1) / target.norm(dim=1).clip(min=0.05)\n mask = mask.squeeze(1)\n\n error_image = masked_mean(data=error, mask=mask, dim=self.dims).squeeze(1, 2)\n error_image = self.output_fn(error_image)\n return error_image\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n output_fn=config[\"output_fn\"],\n input_fn=config[\"input_fn\"],\n )\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/confidence.py", "content": "import torch\nimport torch.nn as nn\n\nfrom .utils import FNS, masked_mean\n\n\nclass Confidence(nn.Module):\n def __init__(\n self,\n weight: float,\n output_fn: str = \"sqrt\",\n input_fn: str = \"linear\",\n rescale: bool = True,\n eps: float = 1e-5,\n ):\n super(Confidence, self).__init__()\n self.name: str = self.__class__.__name__\n self.weight = weight\n self.rescale = rescale\n self.eps = eps\n self.output_fn = FNS[output_fn]\n self.input_fn = FNS[input_fn]\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self,\n input: torch.Tensor,\n target_pred: torch.Tensor,\n target_gt: torch.Tensor,\n mask: torch.Tensor,\n ):\n B, C = target_gt.shape[:2]\n mask = mask.bool()\n target_gt = target_gt.float().reshape(B, C, -1)\n target_pred = target_pred.float().reshape(B, C, -1)\n input = input.float().reshape(B, -1)\n mask = mask.reshape(B, -1)\n\n if self.rescale:\n target_pred = torch.stack(\n [\n p * torch.median(gt[:, m]) / torch.median(p[:, m])\n for p, gt, m in zip(target_pred, target_gt, mask)\n ]\n )\n\n error = torch.abs(\n (self.input_fn(target_pred) - self.input_fn(target_gt)).norm(dim=1) - input\n )\n losses = masked_mean(error, dim=[-1], mask=mask).squeeze(dim=-1)\n losses = self.output_fn(losses)\n\n return losses\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n output_fn=config[\"output_fn\"],\n input_fn=config[\"input_fn\"],\n rescale=config.get(\"rescale\", True),\n )\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/distill.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\n\nfrom .utils import FNS, masked_mean\n\n\nclass SelfDistill(nn.Module):\n def __init__(self, weight: float, output_fn: str = \"sqrt\", eps: float = 1e-5):\n super().__init__()\n self.name: str = self.__class__.__name__\n self.weight: float = weight\n self.dims = (-2, -1)\n self.output_fn = FNS[output_fn]\n self.eps: float = eps\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self,\n input: torch.Tensor,\n intrinsics: torch.Tensor,\n mask: torch.Tensor,\n flips: torch.Tensor,\n downsample_ratio=14,\n ) -> torch.Tensor:\n chunks = input.shape[0] // 2\n mask = F.interpolate(mask.float(), size=input.shape[-2:], mode=\"nearest\")\n\n iters = zip(\n input.chunk(chunks),\n mask.chunk(chunks),\n intrinsics.chunk(chunks),\n flips.chunk(chunks),\n )\n inputs0, inputs1, masks = [], [], []\n for i, (pair_input, pair_mask, pair_cam, pair_flip) in enumerate(iters):\n\n mask0, mask1 = pair_mask\n input0, input1 = pair_input\n cam0, cam1 = pair_cam\n flip0, flip1 = pair_flip\n\n fx_0 = cam0[0, 0] / downsample_ratio\n fx_1 = cam1[0, 0] / downsample_ratio\n cx_0 = cam0[0, 2] / downsample_ratio\n cx_1 = cam1[0, 2] / downsample_ratio\n cy_0 = cam0[1, 2] / downsample_ratio\n cy_1 = cam1[1, 2] / downsample_ratio\n\n # flip image\n if flip0 ^ flip1:\n input0 = torch.flip(input0, dims=(2,))\n mask0 = torch.flip(mask0, dims=(2,))\n cx_0 = input0.shape[-1] - cx_0\n\n # calc zoom\n zoom_x = float(fx_1 / fx_0)\n\n # apply zoom\n input0 = F.interpolate(\n input0.unsqueeze(0), scale_factor=zoom_x, mode=\"bilinear\"\n ).squeeze(0)\n mask0 = F.interpolate(\n mask0.unsqueeze(0), scale_factor=zoom_x, mode=\"nearest\"\n ).squeeze(0)\n\n # calc translation\n change_left = int(cx_1 - (cx_0 - 0.5) * zoom_x - 0.5)\n change_top = int(cy_1 - (cy_0 - 0.5) * zoom_x - 0.5)\n change_right = input1.shape[-1] - change_left - input0.shape[-1]\n change_bottom = input1.shape[-2] - change_top - input0.shape[-2]\n\n # apply translation\n pad_left = max(0, change_left)\n pad_right = max(0, change_right)\n pad_top = max(0, change_top)\n pad_bottom = max(0, change_bottom)\n\n crop_left = max(0, -change_left)\n crop_right = max(0, -change_right)\n crop_top = max(0, -change_top)\n crop_bottom = max(0, -change_bottom)\n\n input0 = F.pad(\n input0,\n (pad_left, pad_right, pad_top, pad_bottom),\n mode=\"constant\",\n value=0,\n )\n mask0 = F.pad(\n mask0,\n (pad_left, pad_right, pad_top, pad_bottom),\n mode=\"constant\",\n value=0,\n )\n input0 = input0[\n :,\n crop_top : input0.shape[-2] - crop_bottom,\n crop_left : input0.shape[-1] - crop_right,\n ]\n mask0 = mask0[\n :,\n crop_top : mask0.shape[-2] - crop_bottom,\n crop_left : mask0.shape[-1] - crop_right,\n ]\n\n mask = torch.logical_and(mask0, mask1)\n\n inputs0.append(input0)\n inputs1.append(input1)\n masks.append(mask)\n\n inputs0 = torch.stack(inputs0, dim=0)\n inputs1 = torch.stack(inputs1, dim=0)\n masks = torch.stack(masks, dim=0)\n loss1 = self.loss(inputs0, inputs1.detach(), masks)\n loss2 = self.loss(inputs1, inputs0.detach(), masks)\n return torch.cat([loss1, loss2], dim=0)\n\n def loss(\n self,\n input: torch.Tensor,\n target: torch.Tensor,\n mask: torch.Tensor,\n ) -> torch.Tensor:\n loss = masked_mean(\n (input - target).square().mean(dim=1), mask=mask, dim=[-2, -1]\n )\n return self.output_fn(loss + self.eps)\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n output_fn=config[\"output_fn\"],\n )\n return obj\n\n\nclass TeacherDistill(nn.Module):\n def __init__(\n self,\n weight: float,\n output_fn: str = \"sqrt\",\n cross: bool = False,\n eps: float = 1e-5,\n ):\n super().__init__()\n assert output_fn in FNS\n self.name: str = self.__class__.__name__\n self.weight: float = weight\n self.dims = (-2, -1)\n self.output_fn = FNS[output_fn]\n self.eps: float = eps\n self.cross = cross\n self.threshold = 0.05\n self.head_dim = 64 # hardcoded for vit\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self,\n student_features: torch.Tensor,\n teacher_features: torch.Tensor,\n student_tokens: torch.Tensor,\n teacher_tokens: torch.Tensor,\n mask: torch.Tensor,\n # metas: List[Dict[str, torch.Tensor]],\n ) -> torch.Tensor:\n B = student_features.shape[0]\n device = student_features.device\n chunks = student_features.shape[0] // 2\n\n mask = (\n F.interpolate(\n mask.float() + 1e-3, size=student_features.shape[-2:], mode=\"nearest\"\n )\n > 0.5\n )\n\n # chunk features as self.head_dim\n student_features = rearrange(\n student_features, \"b (n c) h w -> b c h w n\", c=self.head_dim\n )\n teacher_features = rearrange(\n teacher_features, \"b (n c) h w -> b c h w n\", c=self.head_dim\n )\n student_tokens = rearrange(\n student_tokens, \"b t (n c) -> b t c n\", c=self.head_dim\n )\n teacher_tokens = rearrange(\n teacher_tokens, \"b t (n c) -> b t c n\", c=self.head_dim\n )\n\n distance = (\n (student_features - teacher_features)\n .square()\n .sum(dim=1, keepdim=True)\n .sqrt()\n .mean(dim=-1)\n )\n loss_features = masked_mean(distance, mask=mask, dim=[-2, -1])\n loss_features = self.output_fn(loss_features.clamp(min=self.eps)).squeeze(\n 1, 2, 3\n )\n\n distance = (\n (student_tokens - teacher_tokens).square().sum(dim=-2).sqrt().mean(dim=-1)\n )\n loss_tokens = self.output_fn(distance.clamp(min=self.eps)).squeeze(1)\n\n return loss_features + 0.01 * loss_tokens\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n output_fn=config[\"output_fn\"],\n cross=config[\"cross\"],\n )\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/dummy.py", "content": "import torch\nimport torch.nn as nn\n\n\nclass Dummy(nn.Module):\n def __init__(self, *args, **kwargs):\n super().__init__()\n self.name: str = self.__class__.__name__\n self.weight = 1.0\n\n def forward(self, dummy: torch.Tensor, *args, **kwargs) -> torch.Tensor:\n return torch.tensor([0.0] * dummy.shape[0], device=dummy.device)\n\n @classmethod\n def build(cls, config):\n obj = cls()\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/local_ssi.py", "content": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom unidepth.utils.geometric import erode\n\nfrom .utils import FNS, ind2sub, masked_mean, masked_quantile, ssi\n\n\ndef sample_strong_edges(edges_img, quantile=0.95, reshape=8):\n # flat\n edges_img = F.interpolate(\n edges_img, scale_factor=1 / reshape, mode=\"bilinear\", align_corners=False\n )\n edges_img_flat = edges_img.flatten(1)\n\n # Find strong edges\n edges_mask = edges_img_flat > torch.quantile(\n edges_img_flat, quantile, dim=-1, keepdim=True\n )\n num_samples = edges_mask.sum(dim=-1)\n if (num_samples < 10).any():\n # sample random edges where num_samples < 2\n random = torch.rand_like(edges_img_flat[num_samples < 10, :]) > quantile\n edges_mask[num_samples < 10, :] = torch.logical_or(\n edges_mask[num_samples < 10, :], random\n )\n num_samples = edges_mask.sum(dim=-1)\n\n min_samples = num_samples.min()\n\n # Compute the coordinates of the strong edges as B, N, 2\n edges_coords = torch.stack(\n [torch.nonzero(x, as_tuple=False)[:min_samples].squeeze() for x in edges_mask]\n )\n edges_coords = (\n torch.stack(ind2sub(edges_coords, edges_img.shape[-1]), dim=-1) * reshape\n )\n return edges_coords\n\n\n@torch.jit.script\ndef extract_patches(tensor, sample_coords, patch_size: tuple[int, int] = (32, 32)):\n N, _, H, W = tensor.shape\n device = tensor.device\n dtype = tensor.dtype\n patch_width, patch_height = patch_size\n pad_width = patch_width // 2\n pad_height = patch_height // 2\n\n # Pad the RGB images for both sheep\n tensor_padded = F.pad(\n tensor,\n (pad_width, pad_width, pad_height, pad_height),\n mode=\"constant\",\n value=0.0,\n )\n\n # Adjust edge coordinates to account for padding\n sample_coords_padded = sample_coords + torch.tensor(\n [pad_height, pad_width], dtype=dtype, device=device\n ).reshape(1, 1, 2)\n\n # Calculate the indices for gather operation\n x_centers = sample_coords_padded[:, :, 1].int()\n y_centers = sample_coords_padded[:, :, 0].int()\n\n all_patches = []\n for tensor_i, x_centers_i, y_centers_i in zip(tensor_padded, x_centers, y_centers):\n patches = []\n for x_center, y_center in zip(x_centers_i, y_centers_i):\n y_start, y_end = y_center - pad_height, y_center + pad_height + 1\n x_start, x_end = x_center - pad_width, x_center + pad_width + 1\n patches.append(tensor_i[..., y_start:y_end, x_start:x_end])\n all_patches.append(torch.stack(patches, dim=0))\n\n return torch.stack(all_patches, dim=0).reshape(N, -1, patch_height * patch_width)\n\n\nclass LocalSSI(nn.Module):\n def __init__(\n self,\n weight: float,\n output_fn: str = \"sqrt\",\n patch_size: tuple[int, int] = (32, 32),\n min_samples: int = 4,\n num_levels: int = 4,\n input_fn: str = \"linear\",\n eps: float = 1e-5,\n ):\n super(LocalSSI, self).__init__()\n self.name: str = self.__class__.__name__\n self.weight = weight\n self.output_fn = FNS[output_fn]\n self.input_fn = FNS[input_fn]\n self.min_samples = min_samples\n self.eps = eps\n patch_logrange = np.linspace(\n start=np.log2(min(patch_size)),\n stop=np.log2(max(patch_size)),\n endpoint=True,\n num=num_levels + 1,\n )\n self.patch_logrange = [\n (x, y) for x, y in zip(patch_logrange[:-1], patch_logrange[1:])\n ]\n self.rescale_fn = ssi\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self,\n input: torch.Tensor,\n target: torch.Tensor,\n mask: torch.Tensor,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n mask = mask.bool()\n input = self.input_fn(input.float())\n target = self.input_fn(target.float())\n B, C, H, W = input.shape\n total_errors = []\n\n for ii, patch_logrange in enumerate(self.patch_logrange):\n\n log_kernel = (\n np.random.uniform(*patch_logrange)\n if self.training\n else np.mean(patch_logrange)\n )\n kernel_size = int(\n (2**log_kernel) * min(input.shape[-2:])\n ) # always smaller than min_shape\n kernel_size = (kernel_size, kernel_size)\n stride = (int(kernel_size[0] * 0.9), int(kernel_size[1] * 0.9))\n\n # unfold is always exceeding right/bottom, roll image only negative\n # to have them back in the unfolding window\n max_roll = (\n (W - kernel_size[1]) % stride[1],\n (H - kernel_size[0]) % stride[0],\n )\n roll_x, roll_y = np.random.randint(-max_roll[0], 1), np.random.randint(\n -max_roll[1], 1\n )\n input_fold = torch.roll(input, shifts=(roll_y, roll_x), dims=(2, 3))\n target_fold = torch.roll(target, shifts=(roll_y, roll_x), dims=(2, 3))\n mask_fold = torch.roll(mask.float(), shifts=(roll_y, roll_x), dims=(2, 3))\n\n # unfold in patches\n input_fold = F.unfold(\n input_fold, kernel_size=kernel_size, stride=stride\n ).permute(\n 0, 2, 1\n ) # B N C*H_p*W_p\n target_fold = F.unfold(\n target_fold, kernel_size=kernel_size, stride=stride\n ).permute(0, 2, 1)\n mask_fold = (\n F.unfold(mask_fold, kernel_size=kernel_size, stride=stride)\n .bool()\n .permute(0, 2, 1)\n )\n\n # calculate error patchwise, then mean over patch, then over image based if sample size is significant\n input_fold, target_fold, _ = self.rescale_fn(\n input_fold, target_fold, mask_fold, dim=[-1]\n )\n error = (input_fold - target_fold).abs()\n\n # calculate elements more then 95 percentile and lower than 5percentile of error\n valid_patches = mask_fold.sum(dim=-1) >= self.min_samples\n error_mean_patch = masked_mean(error, mask_fold, dim=[-1]).squeeze(-1)\n error_mean_image = self.output_fn(error_mean_patch.clamp(min=self.eps))\n error_mean_image = masked_mean(\n error_mean_image, mask=valid_patches, dim=[-1]\n )\n total_errors.append(error_mean_image.squeeze(-1))\n\n # global\n input_rescale = input.reshape(B, C, -1)\n target_rescale = target.reshape(B, C, -1)\n mask = mask.reshape(B, 1, -1).clone()\n input, target, mask = self.rescale_fn(\n input_rescale, target_rescale, mask, dim=[-1]\n )\n error = (input - target).abs().squeeze(1)\n\n mask = mask.squeeze(1)\n error_mean_image = masked_mean(error, mask, dim=[-1]).squeeze(-1)\n error_mean_image = self.output_fn(error_mean_image.clamp(min=self.eps))\n\n total_errors.append(error_mean_image)\n\n errors = torch.stack(total_errors).mean(dim=0)\n return errors\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n patch_size=config[\"patch_size\"],\n output_fn=config[\"output_fn\"],\n min_samples=config[\"min_samples\"],\n num_levels=config[\"num_levels\"],\n input_fn=config[\"input_fn\"],\n )\n return obj\n\n\nclass EdgeGuidedLocalSSI(nn.Module):\n def __init__(\n self,\n weight: float,\n output_fn: str = \"sqrt\",\n min_samples: int = 4,\n input_fn: str = \"linear\",\n use_global: bool = True,\n eps: float = 1e-5,\n ):\n super(EdgeGuidedLocalSSI, self).__init__()\n self.name: str = self.__class__.__name__\n self.weight = weight\n self.output_fn = FNS[output_fn]\n self.input_fn = FNS[input_fn]\n self.min_samples = min_samples\n self.eps = eps\n self.use_global = use_global\n self.rescale_fn = ssi\n\n delta_x = torch.tensor(\n [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], requires_grad=False\n )\n delta_y = torch.tensor(\n [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]], requires_grad=False\n )\n self.delta_x = delta_x.reshape(1, 1, 3, 3)\n self.delta_y = delta_y.reshape(1, 1, 3, 3)\n\n try:\n from unidepth.ops.extract_patches import RandomPatchExtractor\n\n self.random_patch_extractor = RandomPatchExtractor()\n except Exception as e:\n self.random_patch_extractor = extract_patches\n print(\n \"EdgeGuidedLocalSSI reverts to a non cuda-optimized operation, \"\n \"you will experince large slowdown, \"\n \"please install it: \",\n \"`cd ./unidepth/ops/extract_patches && bash compile.sh`\",\n )\n\n def get_edge(self, image, mask):\n channels = image.shape[1]\n device = image.device\n delta_x = self.delta_x.to(device).repeat(channels, 1, 1, 1)\n delta_y = self.delta_y.to(device).repeat(channels, 1, 1, 1)\n image_Gx = F.conv2d(image, delta_x, groups=channels, padding=\"same\") / 8\n image_Gy = F.conv2d(image, delta_y, groups=channels, padding=\"same\") / 8\n image_Gx = (\n image_Gx.square().mean(dim=1, keepdim=True).sqrt()\n ) # RMSE over color dim\n image_Gy = image_Gy.square().mean(dim=1, keepdim=True).sqrt()\n edges = torch.sqrt(image_Gx**2 + image_Gy**2)\n edges[:, :, :3, :] = 0\n edges[:, :, -3:, :] = 0\n edges[:, :, :, :3] = 0\n edges[:, :, :, -3:] = 0\n edges[~mask.bool()] = 0\n return edges\n\n def compute_sample_patch_error(\n self, input, target, mask, sampling_coords, kernel_size, image_size\n ):\n B, C, H, W = input.shape\n patch_size = kernel_size[0] * kernel_size[1]\n input = self.random_patch_extractor(\n input, sampling_coords, kernel_size\n ).reshape(B, -1, patch_size)\n target = self.random_patch_extractor(\n target, sampling_coords, kernel_size\n ).reshape(B, -1, patch_size)\n mask = (\n self.random_patch_extractor(mask.float(), sampling_coords, kernel_size)\n .bool()\n .reshape(B, -1, patch_size)\n )\n input, target, mask = self.rescale_fn(input, target, mask, dim=[-1])\n error = (input - target).abs().clamp(min=self.eps)\n\n valid_patches = mask.sum(dim=-1) >= self.min_samples\n error_mean_patch = masked_mean(error, mask, dim=[-1]).squeeze(-1)\n error_mean_image = self.output_fn(error_mean_patch.clamp(min=self.eps))\n error_mean_image = masked_mean(error_mean_image, mask=valid_patches, dim=[-1])\n return error_mean_image\n\n def compute_image_error(self, input, target, mask, image_size):\n H, W = image_size\n input = input.reshape(-1, 1, H * W)\n target = target.reshape(-1, 1, H * W)\n mask = mask.reshape(-1, 1, H * W)\n input, target, mask = self.rescale_fn(input, target, mask, dim=[-1])\n error = (input - target).abs().clamp(min=self.eps)\n\n error_mean_image = masked_mean(error, mask, dim=[-1]).squeeze(-1)\n error_mean_image = self.output_fn(error_mean_image.clamp(min=self.eps))\n return error_mean_image\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self,\n input: torch.Tensor,\n target: torch.Tensor,\n mask: torch.Tensor,\n image: torch.Tensor | None = None,\n validity_mask: torch.Tensor | None = None,\n *args,\n **kwargs,\n ) -> torch.Tensor:\n\n mask = mask.bool()\n input = self.input_fn(input.float())\n target = self.input_fn(target.float())\n B, _, H, W = input.shape\n total_errors = []\n\n # remove border and black border\n if validity_mask is not None:\n validity_mask = erode(validity_mask.float(), kernel_size=3)\n\n edges = self.get_edge(image, validity_mask)\n # quantile was 0.95?\n edges_coords = sample_strong_edges(edges, quantile=0.9, reshape=14)\n log_kernel = np.random.uniform(0.04, 0.08) if self.training else 0.05\n kernel_size = int(\n log_kernel * min(input.shape[-2:])\n ) # always smaller than min_shape\n kernel_size = kernel_size + int(kernel_size % 2 == 0) # odd num\n kernel_size = (kernel_size, kernel_size)\n error_mean_image = self.compute_sample_patch_error(\n input, target, mask, edges_coords, kernel_size, (H, W)\n )\n total_errors.append(error_mean_image.squeeze(-1))\n\n if self.use_global:\n error_mean_image = self.compute_image_error(input, target, mask, (H, W))\n total_errors.append(error_mean_image.squeeze(-1))\n\n errors = torch.stack(total_errors).mean(dim=0)\n return errors\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n output_fn=config[\"output_fn\"],\n input_fn=config[\"input_fn\"],\n use_global=config[\"use_global\"],\n min_samples=config.get(\"min_samples\", 6),\n )\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/regression.py", "content": "import torch\nimport torch.nn as nn\n\nfrom .utils import FNS, REGRESSION_DICT, masked_mean, masked_quantile\n\n\nclass Regression(nn.Module):\n def __init__(\n self,\n weight: float,\n input_fn: str,\n output_fn: str,\n alpha: float,\n gamma: float,\n fn: str,\n dims: list[int] = [-1],\n quantile: float = 0.0,\n **kwargs,\n ):\n super().__init__()\n self.name = self.__class__.__name__\n self.output_fn = FNS[output_fn]\n self.input_fn = FNS[input_fn]\n self.weight = weight\n self.dims = dims\n self.quantile = quantile\n self.alpha = alpha\n self.gamma = gamma\n self.fn = REGRESSION_DICT[fn]\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self,\n input: torch.Tensor,\n target: torch.Tensor,\n mask: torch.Tensor | None = None,\n **kwargs,\n ) -> torch.Tensor:\n if mask is not None: # usually it is just repeated\n mask = mask[:, 0]\n\n input = self.input_fn(input.float())\n target = self.input_fn(target.float())\n error = self.fn(input - target, gamma=self.gamma, alpha=self.alpha).mean(dim=1)\n mean_error = masked_mean(data=error, mask=mask, dim=self.dims).squeeze(\n self.dims\n )\n mean_error = self.output_fn(mean_error)\n return mean_error\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n output_fn=config[\"output_fn\"],\n input_fn=config[\"input_fn\"],\n dims=config.get(\"dims\", (-1,)),\n alpha=config[\"alpha\"],\n gamma=config[\"gamma\"],\n fn=config[\"fn\"],\n )\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/silog.py", "content": "import torch\nimport torch.nn as nn\n\nfrom .utils import (FNS, REGRESSION_DICT, masked_mean, masked_mean_var,\n masked_quantile)\n\n\nclass SILog(nn.Module):\n def __init__(\n self,\n weight: float,\n input_fn: str = \"linear\",\n output_fn: str = \"sqrt\",\n integrated: float = 0.15,\n dims: list[int] = [-3, -2, -1],\n eps: float = 1e-5,\n ):\n super().__init__()\n self.name: str = self.__class__.__name__\n self.weight: float = weight\n\n self.dims = dims\n self.input_fn = FNS[input_fn]\n self.output_fn = FNS[output_fn]\n self.eps: float = eps\n self.integrated = integrated\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def forward(\n self,\n input: torch.Tensor,\n target: torch.Tensor,\n mask: torch.Tensor,\n si: torch.Tensor,\n **kwargs,\n ) -> torch.Tensor:\n mask = mask.bool()\n error = self.input_fn(input.float()) - self.input_fn(target.float())\n mean_error, var_error = masked_mean_var(\n data=error, mask=mask, dim=self.dims, keepdim=False\n )\n if var_error.ndim > 1:\n var_error = var_error.mean(dim=-1)\n\n if self.integrated > 0.0:\n scale_error = mean_error**2\n var_error = var_error + self.integrated * scale_error * (1 - si.int())\n\n out_loss = self.output_fn(var_error)\n return out_loss\n\n @classmethod\n def build(cls, config):\n obj = cls(\n weight=config[\"weight\"],\n dims=config[\"dims\"],\n output_fn=config[\"output_fn\"],\n input_fn=config[\"input_fn\"],\n integrated=config.get(\"integrated\", 0.15),\n )\n return obj\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/losses/utils.py", "content": "from math import prod\nfrom typing import Any, Dict, List, Optional, Tuple\n\nimport torch\n\nFNS = {\n \"sqrt\": lambda x: torch.sqrt(x + 1e-4),\n \"log\": lambda x: torch.log(x + 1e-4),\n \"log1\": lambda x: torch.log(x + 1),\n # if x -> 0 : log(1/x)\n # if x -> inf : log(1+1/x) -> 1/x + hot\n \"log1i\": lambda x: torch.log(1 + 50 / (1e-4 + x)),\n \"linear\": lambda x: x,\n \"square\": torch.square,\n \"disp\": lambda x: 1 / (x + 1e-4),\n \"disp1\": lambda x: 1 / (1 + x),\n}\n\n\nFNS_INV = {\n \"sqrt\": torch.square,\n \"log\": torch.exp,\n \"log1\": lambda x: torch.exp(x) - 1,\n \"linear\": lambda x: x,\n \"square\": torch.sqrt,\n \"disp\": lambda x: 1 / x,\n}\n\n\ndef masked_mean_var(\n data: torch.Tensor, mask: torch.Tensor, dim: List[int], keepdim: bool = True\n):\n if mask is None:\n return data.mean(dim=dim, keepdim=keepdim), data.var(dim=dim, keepdim=keepdim)\n mask = mask.float()\n mask_sum = torch.sum(mask, dim=dim, keepdim=True)\n # data = torch.nan_to_num(data, nan=0.0)\n mask_mean = torch.sum(data * mask, dim=dim, keepdim=True) / torch.clamp(\n mask_sum, min=1.0\n )\n mask_var = torch.sum(\n mask * (data - mask_mean) ** 2, dim=dim, keepdim=True\n ) / torch.clamp(mask_sum, min=1.0)\n if not keepdim:\n mask_mean, mask_var = mask_mean.squeeze(dim), mask_var.squeeze(dim)\n return mask_mean, mask_var\n\n\ndef masked_mean(data: torch.Tensor, mask: torch.Tensor | None, dim: List[int]):\n if mask is None:\n return data.mean(dim=dim, keepdim=True)\n mask = mask.float()\n mask_sum = torch.sum(mask, dim=dim, keepdim=True)\n mask_mean = torch.sum(\n torch.nan_to_num(data, nan=0.0) * mask, dim=dim, keepdim=True\n ) / mask_sum.clamp(min=1.0)\n return mask_mean\n\n\ndef masked_quantile(\n data: torch.Tensor, mask: torch.Tensor | None, dims: List[int], q: float\n):\n \"\"\"\n Compute the quantile of the data only where the mask is 1 along specified dimensions.\n\n Args:\n data (torch.Tensor): The input data tensor.\n mask (torch.Tensor): The mask tensor with the same shape as data, containing 1s where data should be considered.\n dims (list of int): The dimensions to compute the quantile over.\n q (float): The quantile to compute, must be between 0 and 1.\n\n Returns:\n torch.Tensor: The quantile computed over the specified dimensions, ignoring masked values.\n \"\"\"\n masked_data = data * mask if mask is not None else data\n\n # Get a list of all dimensions\n all_dims = list(range(masked_data.dim()))\n\n # Revert negative dimensions\n dims = [d % masked_data.dim() for d in dims]\n\n # Find the dimensions to keep (not included in the `dims` list)\n keep_dims = [d for d in all_dims if d not in dims]\n\n # Permute dimensions to bring `dims` to the front\n permute_order = dims + keep_dims\n permuted_data = masked_data.permute(permute_order)\n\n # Reshape into 2D: (-1, remaining_dims)\n collapsed_shape = (\n -1,\n prod([permuted_data.size(d) for d in range(len(dims), permuted_data.dim())]),\n )\n reshaped_data = permuted_data.reshape(collapsed_shape)\n if mask is None:\n return torch.quantile(reshaped_data, q, dim=0)\n\n permuted_mask = mask.permute(permute_order)\n reshaped_mask = permuted_mask.reshape(collapsed_shape)\n\n # Calculate quantile along the first dimension where mask is true\n quantiles = []\n for i in range(reshaped_data.shape[1]):\n valid_data = reshaped_data[:, i][reshaped_mask[:, i]]\n if valid_data.numel() == 0:\n # print(\"Warning: No valid data found for quantile calculation.\")\n quantiles.append(reshaped_data[:, i].min() * 0.99)\n else:\n quantiles.append(torch.quantile(valid_data, q, dim=0))\n\n # Stack back into a tensor with reduced dimensions\n quantiles = torch.stack(quantiles)\n quantiles = quantiles.reshape(\n [permuted_data.size(d) for d in range(len(dims), permuted_data.dim())]\n )\n\n return quantiles\n\n\ndef masked_median(data: torch.Tensor, mask: torch.Tensor, dim: List[int]):\n ndim = data.ndim\n data = data.flatten(ndim - len(dim))\n mask = mask.flatten(ndim - len(dim))\n mask_median = torch.median(data[..., mask], dim=-1).values\n return mask_median\n\n\ndef masked_median_mad(data: torch.Tensor, mask: torch.Tensor, dim: List[int]):\n ndim = data.ndim\n data = data.flatten(ndim - len(dim))\n mask = mask.flatten(ndim - len(dim))\n mask_median = torch.median(data[mask], dim=-1, keepdim=True).values\n mask_mad = masked_mean((data - mask_median).abs(), mask, dim=[-1])\n return mask_median, mask_mad\n\n\ndef masked_weighted_mean_var(\n data: torch.Tensor, mask: torch.Tensor, weights: torch.Tensor, dim: Tuple[int, ...]\n):\n if mask is None:\n return data.mean(dim=dim, keepdim=True), data.var(dim=dim, keepdim=True)\n mask = mask.float()\n mask_mean = torch.sum(data * mask * weights, dim=dim, keepdim=True) / torch.sum(\n mask * weights, dim=dim, keepdim=True\n ).clamp(min=1.0)\n # V1**2 - V2, V1: sum w_i, V2: sum w_i**2\n denom = torch.sum(weights * mask, dim=dim, keepdim=True).square() - torch.sum(\n (mask * weights).square(), dim=dim, keepdim=True\n )\n # correction is V1 / (V1**2 - V2), if w_i=1 => N/(N**2 - N) => 1/(N-1) (unbiased estimator of variance, cvd)\n correction_factor = torch.sum(mask * weights, dim=dim, keepdim=True) / denom.clamp(\n min=1.0\n )\n mask_var = correction_factor * torch.sum(\n weights * mask * (data - mask_mean) ** 2, dim=dim, keepdim=True\n )\n return mask_mean, mask_var\n\n\ndef ssi(\n input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor, dim: list[int]\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n # recalculate mask with points in 95% confidence interval\n # the statistics are calculated on the stable points and\n # are similar ot median/MAD, but median/MAD gradients\n # are really weird, so this is a workaround\n input_detach = input.detach()\n input_mean, input_var = masked_mean_var(input_detach, mask=mask, dim=dim)\n target_mean, target_var = masked_mean_var(target, mask=mask, dim=dim)\n input_std = (input_var).clip(min=1e-6).sqrt()\n target_std = (target_var).clip(min=1e-6).sqrt()\n stable_points_input = torch.logical_and(\n input_detach > input_mean - 1.96 * input_std,\n input_detach < input_mean + 1.96 * input_std,\n )\n stable_points_target = torch.logical_and(\n target > target_mean - 1.96 * target_std,\n target < target_mean + 1.96 * target_std,\n )\n stable_mask = stable_points_target & stable_points_input & mask\n\n input_mean, input_var = masked_mean_var(input, mask=stable_mask, dim=dim)\n target_mean, target_var = masked_mean_var(target, mask=stable_mask, dim=dim)\n target_normalized = (target - target_mean) / FNS[\"sqrt\"](target_var)\n input_normalized = (input - input_mean) / FNS[\"sqrt\"](input_var)\n return input_normalized, target_normalized, stable_mask\n\n\ndef ind2sub(idx, cols):\n r = idx // cols\n c = idx % cols\n return r, c\n\n\ndef sub2ind(r, c, cols):\n idx = r * cols + c\n return idx\n\n\ndef l2(input_tensor: torch.Tensor, gamma: float = 1.0, *args, **kwargs) -> torch.Tensor:\n return gamma * (input_tensor / gamma) ** 2\n\n\ndef l1(input_tensor: torch.Tensor, gamma: float = 1.0, *args, **kwargs) -> torch.Tensor:\n return torch.abs(input_tensor)\n\n\ndef charbonnier(\n input_tensor: torch.Tensor, gamma: float = 1.0, *args, **kwargs\n) -> torch.Tensor:\n return torch.sqrt(torch.square(input_tensor) + gamma**2) - gamma\n\n\ndef cauchy(\n input_tensor: torch.Tensor, gamma: float = 1.0, *args, **kwargs\n) -> torch.Tensor:\n return gamma * torch.log(torch.square(input_tensor) / gamma + 1)\n\n\ndef geman_mcclure(\n input_tensor: torch.Tensor, gamma: float = 1.0, *args, **kwargs\n) -> torch.Tensor:\n return gamma * torch.square(input_tensor) / (torch.square(input_tensor) + gamma)\n\n\ndef robust_loss(\n input_tensor: torch.Tensor, alpha: float, gamma: float = 1.0, *args, **kwargs\n) -> torch.Tensor:\n coeff = abs(alpha - 2) / alpha\n power = torch.square(input_tensor) / abs(alpha - 2) / (gamma**2) + 1\n return (\n gamma * coeff * (torch.pow(power, alpha / 2) - 1)\n ) # mult gamma to keep grad magnitude invariant wrt gamma\n\n\nREGRESSION_DICT = {\n \"l2\": l2,\n \"l1\": l1,\n \"cauchy\": cauchy,\n \"charbonnier\": charbonnier,\n \"geman_mcclure\": geman_mcclure,\n \"robust_loss\": robust_loss,\n}\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/ops/scheduler.py", "content": "import weakref\n\nimport numpy as np\n\n\nclass PlainCosineScheduler(object):\n def __init__(\n self,\n klass,\n key,\n warmup_iters,\n total_iters,\n overwrite=False,\n init_value=None,\n base_value=None,\n final_value=None,\n step_init=-1,\n ):\n super().__init__()\n self.iter = step_init\n self.overwrite = overwrite\n self.base_value = base_value\n self.init_value = init_value if init_value is not None else base_value\n self.final_value = final_value\n self.total_iters = total_iters\n self.warmup_iters = warmup_iters\n self.key = key\n self.klass = klass\n self.schedulers = [self.get_scheduler()]\n\n def get_scheduler(self):\n init_value = self.init_value\n base_value = self.base_value\n final_value = self.final_value\n warmup_iters = self.warmup_iters\n total_iters = self.total_iters\n\n # normalize in 0,1, then apply function (power) and denormalize\n normalized_schedule = np.linspace(0, 1, warmup_iters, endpoint=True)\n normalized_schedule = np.power(normalized_schedule, 1)\n warmup_schedule = (base_value - init_value) * normalized_schedule + init_value\n\n # main scheduling\n iters = np.arange(total_iters - warmup_iters + 1)\n schedule = final_value + 0.5 * (base_value - final_value) * (\n 1 + np.cos(np.pi * iters / (len(iters) - 1))\n )\n return np.concatenate((warmup_schedule, schedule))\n\n def step(self):\n self.iter = self.iter + 1\n vals = self[self.iter]\n for i, val in enumerate(vals):\n setattr(self.klass, self.key, val)\n\n def __getitem__(self, it):\n it = min(it, self.total_iters)\n return [scheduler[it] for scheduler in self.schedulers]\n\n\nclass CosineScheduler(object):\n def __init__(\n self,\n optimizer,\n warmup_iters,\n total_iters,\n key,\n overwrite=False,\n init_value=None,\n base_value=None,\n final_value=None,\n step_init=-1,\n ):\n super().__init__()\n self.iter = step_init\n self.overwrite = overwrite\n self.optimizer = optimizer\n self.base_value = base_value\n self.init_value = init_value\n self.final_value = final_value\n self.total_iters = total_iters\n self.warmup_iters = warmup_iters\n self.key = key\n self.schedulers = [\n self.get_schedulers(group) for group in optimizer.param_groups\n ]\n\n def get_schedulers(self, group):\n init_value = group.get(self.key + \"_init\", self.init_value)\n base_value = group.get(self.key + \"_base\", self.base_value)\n final_value = group.get(self.key + \"_final\", self.final_value)\n warmup_iters = self.warmup_iters\n total_iters = self.total_iters\n if self.overwrite:\n final_value = self.final_value\n\n # normalize in 0,1, then apply function (power) and denormalize\n normalized_schedule = np.linspace(0, 1, warmup_iters, endpoint=True)\n normalized_schedule = np.power(normalized_schedule, 1)\n warmup_schedule = (base_value - init_value) * normalized_schedule + init_value\n\n # main scheduling\n iters = np.arange(total_iters - warmup_iters + 1)\n schedule = final_value + 0.5 * (base_value - final_value) * (\n 1 + np.cos(np.pi * iters / (len(iters) - 1))\n )\n return np.concatenate((warmup_schedule, schedule))\n\n def step(self):\n self.iter = self.iter + 1\n vals = self[self.iter]\n for group, val in zip(self.optimizer.param_groups, vals):\n if isinstance(group[self.key], (tuple, list)):\n val = (val, *group[self.key][1:])\n group[self.key] = val\n\n def __getitem__(self, it):\n it = min(it, self.total_iters)\n return [scheduler[it] for scheduler in self.schedulers]\n\n def get(self):\n return [group[self.key] for group in self.optimizer.param_groups]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/__init__.py", "content": "from .camera import invert_pinhole\n# from .validation import validate\nfrom .coordinate import coords_grid, normalize_coords\nfrom .distributed import (barrier, get_dist_info, get_rank, is_main_process,\n setup_multi_processes, setup_slurm,\n sync_tensor_across_gpus)\nfrom .evaluation_depth import (DICT_METRICS, DICT_METRICS_3D, eval_3d,\n eval_depth)\nfrom .geometric import spherical_zbuffer_to_euclidean, unproject_points\nfrom .misc import (format_seconds, get_params, identity, recursive_index,\n remove_padding, to_cpu)\nfrom .visualization import colorize, image_grid, log_train_artifacts\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/camera.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom copy import deepcopy\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\n\nfrom .coordinate import coords_grid\nfrom .misc import recursive_to, squeeze_list\n\n\ndef invert_pinhole(K):\n fx = K[..., 0, 0]\n fy = K[..., 1, 1]\n cx = K[..., 0, 2]\n cy = K[..., 1, 2]\n K_inv = torch.zeros_like(K)\n K_inv[..., 0, 0] = 1.0 / fx\n K_inv[..., 1, 1] = 1.0 / fy\n K_inv[..., 0, 2] = -cx / fx\n K_inv[..., 1, 2] = -cy / fy\n K_inv[..., 2, 2] = 1.0\n return K_inv\n\n\nclass Camera:\n \"\"\"\n This is meant to be an abstract parent class, please use the others as actual cameras.\n Pinhole, FIsheye624, MEI, OPENCV, EUCM, Spherical (Equirectangular).\n\n \"\"\"\n\n def __init__(self, params=None, K=None):\n if params.ndim == 1:\n params = params.unsqueeze(0)\n\n if K is None:\n K = (\n torch.eye(3, device=params.device, dtype=params.dtype)\n .unsqueeze(0)\n .repeat(params.shape[0], 1, 1)\n )\n K[..., 0, 0] = params[..., 0]\n K[..., 1, 1] = params[..., 1]\n K[..., 0, 2] = params[..., 2]\n K[..., 1, 2] = params[..., 3]\n\n self.params = params\n self.K = K\n self.overlap_mask = None\n self.projection_mask = None\n\n def project(self, xyz):\n raise NotImplementedError\n\n def unproject(self, uv):\n raise NotImplementedError\n\n def get_projection_mask(self):\n return self.projection_mask\n\n def get_overlap_mask(self):\n return self.overlap_mask\n\n def reconstruct(self, depth):\n id_coords = coords_grid(\n 1, depth.shape[-2], depth.shape[-1], device=depth.device\n )\n rays = self.unproject(id_coords)\n return (\n rays / rays[:, -1:].clamp(min=1e-4) * depth.clamp(min=1e-4)\n ) # assumption z>0!!!\n\n def resize(self, factor):\n self.K[..., :2, :] *= factor\n self.params[..., :4] *= factor\n return self\n\n def to(self, device, non_blocking=False):\n self.params = self.params.to(device, non_blocking=non_blocking)\n self.K = self.K.to(device, non_blocking=non_blocking)\n return self\n\n def get_rays(self, shapes, noisy=False):\n b, h, w = shapes\n uv = coords_grid(1, h, w, device=self.K.device, noisy=noisy)\n rays = self.unproject(uv)\n return rays / torch.norm(rays, dim=1, keepdim=True).clamp(min=1e-4)\n\n def get_pinhole_rays(self, shapes, noisy=False):\n b, h, w = shapes\n uv = coords_grid(b, h, w, device=self.K.device, homogeneous=True, noisy=noisy)\n rays = (invert_pinhole(self.K) @ uv.reshape(b, 3, -1)).reshape(b, 3, h, w)\n return rays / torch.norm(rays, dim=1, keepdim=True).clamp(min=1e-4)\n\n def flip(self, H, W, direction=\"horizontal\"):\n new_cx = (\n W - self.params[:, 2] if direction == \"horizontal\" else self.params[:, 2]\n )\n new_cy = H - self.params[:, 3] if direction == \"vertical\" else self.params[:, 3]\n self.params = torch.stack(\n [self.params[:, 0], self.params[:, 1], new_cx, new_cy], dim=1\n )\n self.K[..., 0, 2] = new_cx\n self.K[..., 1, 2] = new_cy\n return self\n\n def clone(self):\n return deepcopy(self)\n\n def crop(self, left, top, right=None, bottom=None):\n self.K[..., 0, 2] -= left\n self.K[..., 1, 2] -= top\n self.params[..., 2] -= left\n self.params[..., 3] -= top\n return self\n\n # helper function to get how fov changes based on new original size and new size\n def get_new_fov(self, new_shape, original_shape):\n new_hfov = 2 * torch.atan(\n self.params[..., 2] / self.params[..., 0] * new_shape[1] / original_shape[1]\n )\n new_vfov = 2 * torch.atan(\n self.params[..., 3] / self.params[..., 1] * new_shape[0] / original_shape[0]\n )\n return new_hfov, new_vfov\n\n def mask_overlap_projection(self, projected):\n B, _, H, W = projected.shape\n id_coords = coords_grid(B, H, W, device=projected.device)\n\n # check for mask where flow would overlap with other part of the image\n # eleemtns coming from the border are then masked out\n flow = projected - id_coords\n gamma = 0.1\n sample_grid = gamma * flow + id_coords # sample along the flow\n sample_grid[:, 0] = sample_grid[:, 0] / (W - 1) * 2 - 1\n sample_grid[:, 1] = sample_grid[:, 1] / (H - 1) * 2 - 1\n sampled_flow = F.grid_sample(\n flow,\n sample_grid.permute(0, 2, 3, 1),\n mode=\"bilinear\",\n align_corners=False,\n padding_mode=\"border\",\n )\n mask = (\n (1 - gamma) * torch.norm(flow, dim=1, keepdim=True)\n < torch.norm(sampled_flow, dim=1, keepdim=True)\n ) | (torch.norm(flow, dim=1, keepdim=True) < 1)\n return mask\n\n def _pad_params(self):\n # Ensure params are padded to length 16\n if self.params.shape[1] < 16:\n padding = torch.zeros(\n 16 - self.params.shape[1],\n device=self.params.device,\n dtype=self.params.dtype,\n )\n padding = padding.unsqueeze(0).repeat(self.params.shape[0], 1)\n return torch.cat([self.params, padding], dim=1)\n return self.params\n\n @staticmethod\n def flatten_cameras(cameras): # -> list[Camera]:\n # Recursively flatten BatchCamera into primitive cameras\n flattened_cameras = []\n for camera in cameras:\n if isinstance(camera, BatchCamera):\n flattened_cameras.extend(BatchCamera.flatten_cameras(camera.cameras))\n elif isinstance(camera, list):\n flattened_cameras.extend(camera)\n else:\n flattened_cameras.append(camera)\n return flattened_cameras\n\n @staticmethod\n def _stack_or_cat_cameras(cameras, func, **kwargs):\n # Generalized method to handle stacking or concatenation\n flat_cameras = BatchCamera.flatten_cameras(cameras)\n K_matrices = [camera.K for camera in flat_cameras]\n padded_params = [camera._pad_params() for camera in flat_cameras]\n\n stacked_K = func(K_matrices, **kwargs)\n stacked_params = func(padded_params, **kwargs)\n\n # Keep track of the original classes\n original_class = [x.__class__.__name__ for x in flat_cameras]\n return BatchCamera(stacked_params, stacked_K, original_class, flat_cameras)\n\n @classmethod\n def __torch_function__(cls, func, types, args=(), kwargs=None):\n if kwargs is None:\n kwargs = {}\n\n if func is torch.cat:\n return Camera._stack_or_cat_cameras(args[0], func, **kwargs)\n\n if func is torch.stack:\n return Camera._stack_or_cat_cameras(args[0], func, **kwargs)\n\n if func is torch.flatten:\n return Camera._stack_or_cat_cameras(args[0], torch.cat, **kwargs)\n\n return super().__torch_function__(func, types, args, kwargs)\n\n @property\n def device(self):\n return self.K.device\n\n # here we assume that cx,cy are more or less H/2 and W/2\n @property\n def hfov(self):\n return 2 * torch.atan(self.params[..., 2] / self.params[..., 0])\n\n @property\n def vfov(self):\n return 2 * torch.atan(self.params[..., 3] / self.params[..., 1])\n\n @property\n def max_fov(self):\n return 150.0 / 180.0 * np.pi, 150.0 / 180.0 * np.pi\n\n\nclass Pinhole(Camera):\n def __init__(self, params=None, K=None):\n assert params is not None or K is not None\n if params is None:\n params = torch.stack(\n [K[..., 0, 0], K[..., 1, 1], K[..., 0, 2], K[..., 1, 2]], dim=-1\n )\n super().__init__(params=params, K=K)\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def project(self, pcd):\n b, _, h, w = pcd.shape\n pcd_flat = pcd.reshape(b, 3, -1) # [B, 3, H*W]\n cam_coords = self.K @ pcd_flat\n pcd_proj = cam_coords[:, :2] / cam_coords[:, -1:].clamp(min=0.01)\n pcd_proj = pcd_proj.reshape(b, 2, h, w)\n invalid = (\n (pcd_proj[:, 0] >= 0)\n & (pcd_proj[:, 0] < w)\n & (pcd_proj[:, 1] >= 0)\n & (pcd_proj[:, 1] < h)\n )\n self.projection_mask = (~invalid).unsqueeze(1)\n return pcd_proj\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def unproject(self, uv):\n b, _, h, w = uv.shape\n uv_flat = uv.reshape(b, 2, -1) # [B, 2, H*W]\n uv_homogeneous = torch.cat(\n [uv_flat, torch.ones(b, 1, h * w, device=uv.device)], dim=1\n ) # [B, 3, H*W]\n K_inv = torch.inverse(self.K.float())\n xyz = K_inv @ uv_homogeneous\n xyz = xyz / xyz[:, -1:].clip(min=1e-4)\n xyz = xyz.reshape(b, 3, h, w)\n self.unprojection_mask = xyz[:, -1:] > 1e-4\n return xyz\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def reconstruct(self, depth):\n b, _, h, w = depth.shape\n uv = coords_grid(b, h, w, device=depth.device)\n xyz = self.unproject(uv) * depth.clip(min=0.0)\n return xyz\n\n\nclass EUCM(Camera):\n def __init__(self, params):\n super().__init__(params=params, K=None)\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def project(self, xyz):\n H, W = xyz.shape[-2:]\n fx, fy, cx, cy, alpha, beta = self.params[:6].unbind(dim=1)\n x, y, z = xyz.unbind(dim=1)\n d = torch.sqrt(beta * (x**2 + y**2) + z**2)\n\n x = x / (alpha * d + (1 - alpha) * z).clip(min=1e-3)\n y = y / (alpha * d + (1 - alpha) * z).clip(min=1e-3)\n\n Xnorm = fx * x + cx\n Ynorm = fy * y + cy\n\n coords = torch.stack([Xnorm, Ynorm], dim=1)\n\n invalid = (\n (coords[:, 0] < 0)\n | (coords[:, 0] > W)\n | (coords[:, 1] < 0)\n | (coords[:, 1] > H)\n | (z < 0)\n )\n self.projection_mask = (~invalid).unsqueeze(1)\n\n return coords\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def unproject(self, uv):\n u, v = uv.unbind(dim=1)\n fx, fy, cx, cy, alpha, beta = self.params.unbind(dim=1)\n mx = (u - cx) / fx\n my = (v - cy) / fy\n r_square = mx**2 + my**2\n valid_mask = r_square < torch.where(\n alpha < 0.5, 1e6, 1 / (beta * (2 * alpha - 1))\n )\n sqrt_val = 1 - (2 * alpha - 1) * beta * r_square\n mz = (1 - beta * (alpha**2) * r_square) / (\n alpha * torch.sqrt(sqrt_val.clip(min=1e-5)) + (1 - alpha)\n )\n coeff = 1 / torch.sqrt(mx**2 + my**2 + mz**2 + 1e-5)\n\n x = coeff * mx\n y = coeff * my\n z = coeff * mz\n self.unprojection_mask = valid_mask & (z > 1e-3)\n\n xnorm = torch.stack((x, y, z.clamp(1e-3)), dim=1)\n return xnorm\n\n\nclass Spherical(Camera):\n def __init__(self, params):\n # Hfov and Vofv are in radians and halved!\n super().__init__(params=params, K=None)\n\n def resize(self, factor):\n self.K[..., :2, :] *= factor\n self.params[..., :6] *= factor\n return self\n\n def crop(self, left, top, right, bottom):\n self.K[..., 0, 2] -= left\n self.K[..., 1, 2] -= top\n self.params[..., 2] -= left\n self.params[..., 3] -= top\n W, H = self.params[..., 4], self.params[..., 5]\n angle_ratio_W = (W - left - right) / W\n angle_ratio_H = (H - top - bottom) / H\n\n self.params[..., 4] -= left + right\n self.params[..., 5] -= top + bottom\n\n # rescale hfov and vfov\n self.params[..., 6] *= angle_ratio_W\n self.params[..., 7] *= angle_ratio_H\n return self\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def project(self, xyz):\n width, height = self.params[..., 4], self.params[..., 5]\n hfov, vfov = 2 * self.params[..., 6], 2 * self.params[..., 7]\n longitude = torch.atan2(xyz[:, 0], xyz[:, 2])\n latitude = torch.asin(xyz[:, 1] / torch.norm(xyz, dim=1).clamp(min=1e-5))\n\n u = longitude / hfov * (width - 1) + (width - 1) / 2\n v = latitude / vfov * (height - 1) + (height - 1) / 2\n\n return torch.stack([u, v], dim=1)\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def unproject(self, uv):\n u, v = uv.unbind(dim=1)\n\n width, height = self.params[..., 4], self.params[..., 5]\n hfov, vfov = 2 * self.params[..., 6], 2 * self.params[..., 7]\n longitude = (u - (width - 1) / 2) / (width - 1) * hfov\n latitude = (v - (height - 1) / 2) / (height - 1) * vfov\n x = torch.cos(latitude) * torch.sin(longitude)\n z = torch.cos(latitude) * torch.cos(longitude)\n y = torch.sin(latitude)\n unit_sphere = torch.stack([x, y, z], dim=1)\n unit_sphere = unit_sphere / torch.norm(unit_sphere, dim=1, keepdim=True).clip(\n min=1e-5\n )\n\n return unit_sphere\n\n def reconstruct(self, depth):\n id_coords = coords_grid(\n 1, depth.shape[-2], depth.shape[-1], device=depth.device\n )\n return self.unproject(id_coords) * depth\n\n def get_new_fov(self, new_shape, original_shape):\n new_hfov = 2 * self.params[..., 6] * new_shape[1] / original_shape[1]\n new_vfov = 2 * self.params[..., 7] * new_shape[0] / original_shape[0]\n return new_hfov, new_vfov\n\n @property\n def hfov(self):\n return 2 * self.params[..., 6]\n\n @property\n def vfov(self):\n return 2 * self.params[..., 7]\n\n @property\n def max_fov(self):\n return 2 * np.pi, 0.9 * np.pi # avoid strong distortion on tops\n\n\nclass OPENCV(Camera):\n def __init__(self, params):\n super().__init__(params=params, K=None)\n self.use_radial = self.params[..., 4:10].abs().sum() > 1e-6\n assert (\n self.params[..., 7:10].abs().sum() == 0.0\n ), \"Do not support poly division model\"\n self.use_tangential = self.params[..., 10:12].abs().sum() > 1e-6\n self.use_thin_prism = self.params[..., 12:].abs().sum() > 1e-6\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def project(self, xyz):\n eps = 1e-9\n B, _, H, W = xyz.shape\n N = H * W\n xyz = xyz.permute(0, 2, 3, 1).reshape(B, N, 3)\n\n # Radial correction.\n z = xyz[:, :, 2].reshape(B, N, 1)\n z = torch.where(torch.abs(z) < eps, eps * torch.sign(z), z)\n ab = xyz[:, :, :2] / z\n r = torch.norm(ab, dim=-1, p=2, keepdim=True)\n th = r\n\n # Create powers of th (th^3, th^5, ...)\n th_pow = torch.cat([torch.pow(th, 2 + i * 2) for i in range(3)], dim=-1)\n distortion_coeffs_num = self.params[:, 4:7].reshape(B, 1, 3)\n distortion_coeffs_den = self.params[:, 7:10].reshape(B, 1, 3)\n th_num = 1 + torch.sum(th_pow * distortion_coeffs_num, dim=-1, keepdim=True)\n th_den = 1 + torch.sum(th_pow * distortion_coeffs_den, dim=-1, keepdim=True)\n\n xr_yr = ab * th_num / th_den\n uv_dist = xr_yr\n\n # Tangential correction.\n p0 = self.params[..., -6].reshape(B, 1)\n p1 = self.params[..., -5].reshape(B, 1)\n xr = xr_yr[:, :, 0].reshape(B, N)\n yr = xr_yr[:, :, 1].reshape(B, N)\n xr_yr_sq = torch.square(xr_yr)\n xr_sq = xr_yr_sq[:, :, 0].reshape(B, N)\n yr_sq = xr_yr_sq[:, :, 1].reshape(B, N)\n rd_sq = xr_sq + yr_sq\n uv_dist_tu = uv_dist[:, :, 0] + (\n (2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1\n )\n uv_dist_tv = uv_dist[:, :, 1] + (\n (2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0\n )\n uv_dist = torch.stack(\n [uv_dist_tu, uv_dist_tv], dim=-1\n ) # Avoids in-place complaint.\n\n # Thin Prism correction.\n s0 = self.params[..., -4].reshape(B, 1)\n s1 = self.params[..., -3].reshape(B, 1)\n s2 = self.params[..., -2].reshape(B, 1)\n s3 = self.params[..., -1].reshape(B, 1)\n rd_4 = torch.square(rd_sq)\n uv_dist[:, :, 0] = uv_dist[:, :, 0] + (s0 * rd_sq + s1 * rd_4)\n uv_dist[:, :, 1] = uv_dist[:, :, 1] + (s2 * rd_sq + s3 * rd_4)\n\n # Finally, apply standard terms: focal length and camera centers.\n if self.params.shape[-1] == 15:\n fx_fy = self.params[..., 0].reshape(B, 1, 1)\n cx_cy = self.params[..., 1:3].reshape(B, 1, 2)\n else:\n fx_fy = self.params[..., 0:2].reshape(B, 1, 2)\n cx_cy = self.params[..., 2:4].reshape(B, 1, 2)\n result = uv_dist * fx_fy + cx_cy\n\n result = result.reshape(B, H, W, 2).permute(0, 3, 1, 2)\n invalid = (\n (result[:, 0] < 0)\n | (result[:, 0] > W)\n | (result[:, 1] < 0)\n | (result[:, 1] > H)\n )\n self.projection_mask = (~invalid).unsqueeze(1)\n self.overlap_mask = self.mask_overlap_projection(result)\n\n return result\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def unproject(self, uv, max_iters: int = 10):\n eps = 1e-3\n B, _, H, W = uv.shape\n N = H * W\n uv = uv.permute(0, 2, 3, 1).reshape(B, N, 2)\n\n if self.params.shape[-1] == 15:\n fx_fy = self.params[..., 0].reshape(B, 1, 1)\n cx_cy = self.params[..., 1:3].reshape(B, 1, 2)\n else:\n fx_fy = self.params[..., 0:2].reshape(B, 1, 2)\n cx_cy = self.params[..., 2:4].reshape(B, 1, 2)\n\n uv_dist = (uv - cx_cy) / fx_fy\n\n # Compute xr_yr using Newton's method.\n xr_yr = uv_dist.clone() # Initial guess.\n max_iters_tanprism = (\n max_iters if self.use_thin_prism or self.use_tangential else 0\n )\n\n for _ in range(max_iters_tanprism):\n uv_dist_est = xr_yr.clone()\n xr = xr_yr[..., 0].reshape(B, N)\n yr = xr_yr[..., 1].reshape(B, N)\n xr_yr_sq = torch.square(xr_yr)\n xr_sq = xr_yr_sq[..., 0].reshape(B, N)\n yr_sq = xr_yr_sq[..., 1].reshape(B, N)\n rd_sq = xr_sq + yr_sq\n\n if self.use_tangential:\n # Tangential terms.\n p0 = self.params[..., -6].reshape(B, 1)\n p1 = self.params[..., -5].reshape(B, 1)\n uv_dist_est[..., 0] = uv_dist_est[..., 0] + (\n (2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1\n )\n uv_dist_est[..., 1] = uv_dist_est[..., 1] + (\n (2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0\n )\n\n if self.use_thin_prism:\n # Thin Prism terms.\n s0 = self.params[..., -4].reshape(B, 1)\n s1 = self.params[..., -3].reshape(B, 1)\n s2 = self.params[..., -2].reshape(B, 1)\n s3 = self.params[..., -1].reshape(B, 1)\n rd_4 = torch.square(rd_sq)\n uv_dist_est[:, :, 0] = uv_dist_est[:, :, 0] + (s0 * rd_sq + s1 * rd_4)\n uv_dist_est[:, :, 1] = uv_dist_est[:, :, 1] + (s2 * rd_sq + s3 * rd_4)\n\n # Compute the derivative of uv_dist w.r.t. xr_yr.\n duv_dist_dxr_yr = uv.new_ones(B, N, 2, 2)\n\n if self.use_tangential:\n duv_dist_dxr_yr[..., 0, 0] = 1.0 + 6.0 * xr * p0 + 2.0 * yr * p1\n offdiag = 2.0 * (xr * p1 + yr * p0)\n duv_dist_dxr_yr[..., 0, 1] = offdiag\n duv_dist_dxr_yr[..., 1, 0] = offdiag\n duv_dist_dxr_yr[..., 1, 1] = 1.0 + 6.0 * yr * p1 + 2.0 * xr * p0\n\n if self.use_thin_prism:\n xr_yr_sq_norm = xr_sq + yr_sq\n temp1 = 2.0 * (s0 + 2.0 * s1 * xr_yr_sq_norm)\n duv_dist_dxr_yr[..., 0, 0] = duv_dist_dxr_yr[..., 0, 0] + (xr * temp1)\n duv_dist_dxr_yr[..., 0, 1] = duv_dist_dxr_yr[..., 0, 1] + (yr * temp1)\n temp2 = 2.0 * (s2 + 2.0 * s3 * xr_yr_sq_norm)\n duv_dist_dxr_yr[..., 1, 0] = duv_dist_dxr_yr[..., 1, 0] + (xr * temp2)\n duv_dist_dxr_yr[..., 1, 1] = duv_dist_dxr_yr[..., 1, 1] + (yr * temp2)\n\n mat = duv_dist_dxr_yr.reshape(-1, 2, 2)\n a = mat[:, 0, 0].reshape(-1, 1, 1)\n b = mat[:, 0, 1].reshape(-1, 1, 1)\n c = mat[:, 1, 0].reshape(-1, 1, 1)\n d = mat[:, 1, 1].reshape(-1, 1, 1)\n det = 1.0 / ((a * d) - (b * c))\n top = torch.cat([d, -b], dim=-1)\n bot = torch.cat([-c, a], dim=-1)\n inv = det * torch.cat([top, bot], dim=-2)\n inv = inv.reshape(B, N, 2, 2)\n diff = uv_dist - uv_dist_est\n a = inv[..., 0, 0]\n b = inv[..., 0, 1]\n c = inv[..., 1, 0]\n d = inv[..., 1, 1]\n e = diff[..., 0]\n f = diff[..., 1]\n step = torch.stack([a * e + b * f, c * e + d * f], dim=-1)\n # Newton step.\n xr_yr = xr_yr + step\n\n # Compute theta using Newton's method.\n xr_yr_norm = xr_yr.norm(p=2, dim=2).reshape(B, N, 1)\n th = xr_yr_norm.clone()\n max_iters_radial = max_iters if self.use_radial else 0\n c = (\n torch.tensor([2.0 * i + 3 for i in range(3)], device=self.device)\n .reshape(1, 1, 3)\n .repeat(B, 1, 1)\n )\n radial_params_num = self.params[..., 4:7].reshape(B, 1, 3)\n\n # Trust region parameters\n delta = torch.full((B, N, 1), 0.1, device=self.device) # Initial trust radius\n delta_max = torch.tensor(1.0, device=self.device) # Maximum trust radius\n eta = 0.1 # Acceptable reduction threshold\n\n for i in range(max_iters_radial):\n th_sq = th * th # th^2\n # Compute powers of th^2 up to th^(12)\n theta_powers = torch.cat(\n [th_sq ** (i + 1) for i in range(3)], dim=-1\n ) # Shape: (B, N, 6)\n\n # Compute th_radial: radial distortion model applied to th\n th_radial = 1.0 + torch.sum(\n theta_powers * radial_params_num, dim=-1, keepdim=True\n )\n th_radial = th_radial * th # Multiply by th at the end\n\n # Compute derivative dthd_th\n dthd_th = 1.0 + torch.sum(\n c * radial_params_num * theta_powers, dim=-1, keepdim=True\n )\n dthd_th = dthd_th # Already includes derivative terms\n\n # Compute residual\n residual = th_radial - xr_yr_norm # Shape: (B, N, 1)\n residual_norm = torch.norm(residual, dim=2, keepdim=True) # For each pixel\n\n # Check for convergence\n if torch.max(torch.abs(residual)) < eps:\n break\n\n # Avoid division by zero by adding a small epsilon\n safe_dthd_th = dthd_th.clone()\n zero_derivative_mask = dthd_th.abs() < eps\n safe_dthd_th[zero_derivative_mask] = eps\n\n # Compute Newton's step\n step = -residual / safe_dthd_th\n\n # Compute predicted reduction\n predicted_reduction = -(residual * step).sum(dim=2, keepdim=True)\n\n # Adjust step based on trust region\n step_norm = torch.norm(step, dim=2, keepdim=True)\n over_trust_mask = step_norm > delta\n\n # Scale step if it exceeds trust radius\n step_scaled = step.clone()\n step_scaled[over_trust_mask] = step[over_trust_mask] * (\n delta[over_trust_mask] / step_norm[over_trust_mask]\n )\n\n # Update theta\n th_new = th + step_scaled\n\n # Compute new residual\n th_sq_new = th_new * th_new\n theta_powers_new = torch.cat(\n [th_sq_new ** (j + 1) for j in range(3)], dim=-1\n )\n th_radial_new = 1.0 + torch.sum(\n theta_powers_new * radial_params_num, dim=-1, keepdim=True\n )\n th_radial_new = th_radial_new * th_new\n residual_new = th_radial_new - xr_yr_norm\n residual_new_norm = torch.norm(residual_new, dim=2, keepdim=True)\n\n # Compute actual reduction\n actual_reduction = residual_norm - residual_new_norm\n\n # Compute ratio of actual to predicted reduction\n # predicted_reduction[predicted_reduction.abs() < eps] = eps #* torch.sign(predicted_reduction[predicted_reduction.abs() < eps])\n rho = actual_reduction / predicted_reduction\n rho[(actual_reduction == 0) & (predicted_reduction == 0)] = 1.0\n\n # Update trust radius delta\n delta_update_mask = rho > 0.5\n delta[delta_update_mask] = torch.min(\n 2.0 * delta[delta_update_mask], delta_max\n )\n\n delta_decrease_mask = rho < 0.2\n delta[delta_decrease_mask] = 0.25 * delta[delta_decrease_mask]\n\n # Accept or reject the step\n accept_step_mask = rho > eta\n th = torch.where(accept_step_mask, th_new, th)\n\n # Compute the ray direction using theta and xr_yr.\n close_to_zero = torch.logical_and(th.abs() < eps, xr_yr_norm.abs() < eps)\n ray_dir = torch.where(close_to_zero, xr_yr, th / xr_yr_norm * xr_yr)\n\n ray = torch.cat([ray_dir, uv.new_ones(B, N, 1)], dim=2)\n ray = ray.reshape(B, H, W, 3).permute(0, 3, 1, 2)\n\n return ray\n\n\nclass Fisheye624(Camera):\n def __init__(self, params):\n super().__init__(params=params, K=None)\n self.use_radial = self.params[..., 4:10].abs().sum() > 1e-6\n self.use_tangential = self.params[..., 10:12].abs().sum() > 1e-6\n self.use_thin_prism = self.params[..., 12:].abs().sum() > 1e-6\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def project(self, xyz):\n eps = 1e-9\n B, _, H, W = xyz.shape\n N = H * W\n xyz = xyz.permute(0, 2, 3, 1).reshape(B, N, 3)\n\n # Radial correction.\n z = xyz[:, :, 2].reshape(B, N, 1)\n z = torch.where(torch.abs(z) < eps, eps * torch.sign(z), z)\n ab = xyz[:, :, :2] / z\n r = torch.norm(ab, dim=-1, p=2, keepdim=True)\n th = torch.atan(r)\n th_divr = torch.where(r < eps, torch.ones_like(ab), ab / r)\n\n th_pow = torch.cat(\n [torch.pow(th, 3 + i * 2) for i in range(6)], dim=-1\n ) # Create powers of th (th^3, th^5, ...)\n distortion_coeffs = self.params[:, 4:10].reshape(B, 1, 6)\n th_k = th + torch.sum(th_pow * distortion_coeffs, dim=-1, keepdim=True)\n\n xr_yr = th_k * th_divr\n uv_dist = xr_yr\n\n # Tangential correction.\n p0 = self.params[..., -6].reshape(B, 1)\n p1 = self.params[..., -5].reshape(B, 1)\n xr = xr_yr[:, :, 0].reshape(B, N)\n yr = xr_yr[:, :, 1].reshape(B, N)\n xr_yr_sq = torch.square(xr_yr)\n xr_sq = xr_yr_sq[:, :, 0].reshape(B, N)\n yr_sq = xr_yr_sq[:, :, 1].reshape(B, N)\n rd_sq = xr_sq + yr_sq\n uv_dist_tu = uv_dist[:, :, 0] + (\n (2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1\n )\n uv_dist_tv = uv_dist[:, :, 1] + (\n (2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0\n )\n uv_dist = torch.stack(\n [uv_dist_tu, uv_dist_tv], dim=-1\n ) # Avoids in-place complaint.\n\n # Thin Prism correction.\n s0 = self.params[..., -4].reshape(B, 1)\n s1 = self.params[..., -3].reshape(B, 1)\n s2 = self.params[..., -2].reshape(B, 1)\n s3 = self.params[..., -1].reshape(B, 1)\n rd_4 = torch.square(rd_sq)\n uv_dist[:, :, 0] = uv_dist[:, :, 0] + (s0 * rd_sq + s1 * rd_4)\n uv_dist[:, :, 1] = uv_dist[:, :, 1] + (s2 * rd_sq + s3 * rd_4)\n\n # Finally, apply standard terms: focal length and camera centers.\n if self.params.shape[-1] == 15:\n fx_fy = self.params[..., 0].reshape(B, 1, 1)\n cx_cy = self.params[..., 1:3].reshape(B, 1, 2)\n else:\n fx_fy = self.params[..., 0:2].reshape(B, 1, 2)\n cx_cy = self.params[..., 2:4].reshape(B, 1, 2)\n result = uv_dist * fx_fy + cx_cy\n\n result = result.reshape(B, H, W, 2).permute(0, 3, 1, 2)\n invalid = (\n (result[:, 0] < 0)\n | (result[:, 0] > W)\n | (result[:, 1] < 0)\n | (result[:, 1] > H)\n )\n self.projection_mask = (~invalid).unsqueeze(1)\n self.overlap_mask = self.mask_overlap_projection(result)\n\n return result\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def unproject(self, uv, max_iters: int = 10):\n eps = 1e-3\n B, _, H, W = uv.shape\n N = H * W\n uv = uv.permute(0, 2, 3, 1).reshape(B, N, 2)\n\n if self.params.shape[-1] == 15:\n fx_fy = self.params[..., 0].reshape(B, 1, 1)\n cx_cy = self.params[..., 1:3].reshape(B, 1, 2)\n else:\n fx_fy = self.params[..., 0:2].reshape(B, 1, 2)\n cx_cy = self.params[..., 2:4].reshape(B, 1, 2)\n\n uv_dist = (uv - cx_cy) / fx_fy\n\n # Compute xr_yr using Trust-region method.\n xr_yr = uv_dist.clone()\n max_iters_tanprism = (\n max_iters if self.use_thin_prism or self.use_tangential else 0\n )\n\n for _ in range(max_iters_tanprism):\n uv_dist_est = xr_yr.clone()\n xr = xr_yr[..., 0].reshape(B, N)\n yr = xr_yr[..., 1].reshape(B, N)\n xr_yr_sq = torch.square(xr_yr)\n xr_sq = xr_yr_sq[..., 0].reshape(B, N)\n yr_sq = xr_yr_sq[..., 1].reshape(B, N)\n rd_sq = xr_sq + yr_sq\n\n if self.use_tangential:\n # Tangential terms.\n p0 = self.params[..., -6].reshape(B, 1)\n p1 = self.params[..., -5].reshape(B, 1)\n uv_dist_est[..., 0] = uv_dist_est[..., 0] + (\n (2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1\n )\n uv_dist_est[..., 1] = uv_dist_est[..., 1] + (\n (2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0\n )\n\n if self.use_thin_prism:\n # Thin Prism terms.\n s0 = self.params[..., -4].reshape(B, 1)\n s1 = self.params[..., -3].reshape(B, 1)\n s2 = self.params[..., -2].reshape(B, 1)\n s3 = self.params[..., -1].reshape(B, 1)\n rd_4 = torch.square(rd_sq)\n uv_dist_est[:, :, 0] = uv_dist_est[:, :, 0] + (s0 * rd_sq + s1 * rd_4)\n uv_dist_est[:, :, 1] = uv_dist_est[:, :, 1] + (s2 * rd_sq + s3 * rd_4)\n\n # Compute the derivative of uv_dist w.r.t. xr_yr.\n duv_dist_dxr_yr = uv.new_ones(B, N, 2, 2)\n\n if self.use_tangential:\n duv_dist_dxr_yr[..., 0, 0] = 1.0 + 6.0 * xr * p0 + 2.0 * yr * p1\n offdiag = 2.0 * (xr * p1 + yr * p0)\n duv_dist_dxr_yr[..., 0, 1] = offdiag\n duv_dist_dxr_yr[..., 1, 0] = offdiag\n duv_dist_dxr_yr[..., 1, 1] = 1.0 + 6.0 * yr * p1 + 2.0 * xr * p0\n\n if self.use_thin_prism:\n xr_yr_sq_norm = xr_sq + yr_sq\n temp1 = 2.0 * (s0 + 2.0 * s1 * xr_yr_sq_norm)\n duv_dist_dxr_yr[..., 0, 0] = duv_dist_dxr_yr[..., 0, 0] + (xr * temp1)\n duv_dist_dxr_yr[..., 0, 1] = duv_dist_dxr_yr[..., 0, 1] + (yr * temp1)\n temp2 = 2.0 * (s2 + 2.0 * s3 * xr_yr_sq_norm)\n duv_dist_dxr_yr[..., 1, 0] = duv_dist_dxr_yr[..., 1, 0] + (xr * temp2)\n duv_dist_dxr_yr[..., 1, 1] = duv_dist_dxr_yr[..., 1, 1] + (yr * temp2)\n\n mat = duv_dist_dxr_yr.reshape(-1, 2, 2)\n a = mat[:, 0, 0].reshape(-1, 1, 1)\n b = mat[:, 0, 1].reshape(-1, 1, 1)\n c = mat[:, 1, 0].reshape(-1, 1, 1)\n d = mat[:, 1, 1].reshape(-1, 1, 1)\n det = 1.0 / ((a * d) - (b * c))\n top = torch.cat([d, -b], dim=-1)\n bot = torch.cat([-c, a], dim=-1)\n inv = det * torch.cat([top, bot], dim=-2)\n inv = inv.reshape(B, N, 2, 2)\n diff = uv_dist - uv_dist_est\n a = inv[..., 0, 0]\n b = inv[..., 0, 1]\n c = inv[..., 1, 0]\n d = inv[..., 1, 1]\n e = diff[..., 0]\n f = diff[..., 1]\n step = torch.stack([a * e + b * f, c * e + d * f], dim=-1)\n # Newton step.\n xr_yr = xr_yr + step\n\n # Compute theta using Newton's method.\n xr_yr_norm = xr_yr.norm(p=2, dim=2).reshape(B, N, 1)\n th = xr_yr_norm.clone()\n max_iters_radial = max_iters if self.use_radial else 0\n c = (\n torch.tensor([2.0 * i + 3 for i in range(6)], device=self.device)\n .reshape(1, 1, 6)\n .repeat(B, 1, 1)\n )\n radial_params = self.params[..., 4:10].reshape(B, 1, 6)\n\n # Trust region parameters\n delta = torch.full((B, N, 1), 0.1, device=self.device) # Initial trust radius\n delta_max = torch.tensor(1.0, device=self.device) # Maximum trust radius\n eta = 0.1 # Acceptable reduction threshold\n\n for i in range(max_iters_radial):\n th_sq = th * th\n # Compute powers of th^2 up to th^(12)\n theta_powers = torch.cat(\n [th_sq ** (i + 1) for i in range(6)], dim=-1\n ) # Shape: (B, N, 6)\n\n # Compute th_radial: radial distortion model applied to th\n th_radial = 1.0 + torch.sum(\n theta_powers * radial_params, dim=-1, keepdim=True\n )\n th_radial = th_radial * th\n\n # Compute derivative dthd_th\n dthd_th = 1.0 + torch.sum(\n c * radial_params * theta_powers, dim=-1, keepdim=True\n )\n\n # Compute residual\n residual = th_radial - xr_yr_norm # Shape: (B, N, 1)\n residual_norm = torch.norm(residual, dim=2, keepdim=True)\n\n # Check for convergence\n if torch.max(torch.abs(residual)) < eps:\n break\n\n # Avoid division by zero by adding a small epsilon\n safe_dthd_th = dthd_th.clone()\n zero_derivative_mask = dthd_th.abs() < eps\n safe_dthd_th[zero_derivative_mask] = eps\n\n # Compute Newton's step\n step = -residual / safe_dthd_th\n\n # Compute predicted reduction\n predicted_reduction = -(residual * step).sum(dim=2, keepdim=True)\n\n # Adjust step based on trust region\n step_norm = torch.norm(step, dim=2, keepdim=True)\n over_trust_mask = step_norm > delta\n\n # Scale step if it exceeds trust radius\n step_scaled = step.clone()\n step_scaled[over_trust_mask] = step[over_trust_mask] * (\n delta[over_trust_mask] / step_norm[over_trust_mask]\n )\n\n # Update theta\n th_new = th + step_scaled\n\n # Compute new residual\n th_sq_new = th_new * th_new\n theta_powers_new = torch.cat(\n [th_sq_new ** (j + 1) for j in range(6)], dim=-1\n )\n th_radial_new = 1.0 + torch.sum(\n theta_powers_new * radial_params, dim=-1, keepdim=True\n )\n th_radial_new = th_radial_new * th_new\n residual_new = th_radial_new - xr_yr_norm\n residual_new_norm = torch.norm(residual_new, dim=2, keepdim=True)\n\n # Compute actual reduction\n actual_reduction = residual_norm - residual_new_norm\n\n # Compute ratio of actual to predicted reduction\n rho = actual_reduction / predicted_reduction\n rho[(actual_reduction == 0) & (predicted_reduction == 0)] = 1.0\n\n # Update trust radius delta\n delta_update_mask = rho > 0.5\n delta[delta_update_mask] = torch.min(\n 2.0 * delta[delta_update_mask], delta_max\n )\n\n delta_decrease_mask = rho < 0.2\n delta[delta_decrease_mask] = 0.25 * delta[delta_decrease_mask]\n\n # Accept or reject the step\n accept_step_mask = rho > eta\n th = torch.where(accept_step_mask, th_new, th)\n\n # Compute the ray direction using theta and xr_yr.\n close_to_zero = torch.logical_and(th.abs() < eps, xr_yr_norm.abs() < eps)\n ray_dir = torch.where(close_to_zero, xr_yr, torch.tan(th) / xr_yr_norm * xr_yr)\n\n ray = torch.cat([ray_dir, uv.new_ones(B, N, 1)], dim=2)\n ray = ray.reshape(B, H, W, 3).permute(0, 3, 1, 2)\n\n return ray\n\n\nclass MEI(Camera):\n def __init__(self, params):\n super().__init__(params=params, K=None)\n # fx fy cx cy k1 k2 p1 p2 xi\n self.use_radial = self.params[..., 4:6].abs().sum() > 1e-6\n self.use_tangential = self.params[..., 6:8].abs().sum() > 1e-6\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def unproject(self, uv, max_iters: int = 20):\n eps = 1e-6\n B, _, H, W = uv.shape\n N = H * W\n uv = uv.permute(0, 2, 3, 1).reshape(B, N, 2)\n\n k1, k2, p0, p1, xi = self.params[..., 4:9].unbind(dim=1)\n fx_fy = self.params[..., 0:2].reshape(B, 1, 2)\n cx_cy = self.params[..., 2:4].reshape(B, 1, 2)\n\n uv_dist = (uv - cx_cy) / fx_fy\n\n # Compute xr_yr using Newton's method.\n xr_yr = uv_dist.clone() # Initial guess.\n max_iters_tangential = max_iters if self.use_tangential else 0\n for _ in range(max_iters_tangential):\n uv_dist_est = xr_yr.clone()\n\n # Tangential terms.\n xr = xr_yr[..., 0]\n yr = xr_yr[..., 1]\n xr_yr_sq = xr_yr**2\n xr_sq = xr_yr_sq[..., 0]\n yr_sq = xr_yr_sq[..., 1]\n rd_sq = xr_sq + yr_sq\n uv_dist_est[..., 0] = uv_dist_est[..., 0] + (\n (2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1\n )\n uv_dist_est[..., 1] = uv_dist_est[..., 1] + (\n (2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0\n )\n\n # Compute the derivative of uv_dist w.r.t. xr_yr.\n duv_dist_dxr_yr = torch.ones((B, N, 2, 2), device=uv.device)\n duv_dist_dxr_yr[..., 0, 0] = 1.0 + 6.0 * xr * p0 + 2.0 * yr * p1\n offdiag = 2.0 * (xr * p1 + yr * p0)\n duv_dist_dxr_yr[..., 0, 1] = offdiag\n duv_dist_dxr_yr[..., 1, 0] = offdiag\n duv_dist_dxr_yr[..., 1, 1] = 1.0 + 6.0 * yr * p1 + 2.0 * xr * p0\n\n mat = duv_dist_dxr_yr.reshape(-1, 2, 2)\n a = mat[:, 0, 0].reshape(-1, 1, 1)\n b = mat[:, 0, 1].reshape(-1, 1, 1)\n c = mat[:, 1, 0].reshape(-1, 1, 1)\n d = mat[:, 1, 1].reshape(-1, 1, 1)\n det = 1.0 / ((a * d) - (b * c))\n top = torch.cat([d, -b], dim=-1)\n bot = torch.cat([-c, a], dim=-1)\n inv = det * torch.cat([top, bot], dim=-2)\n inv = inv.reshape(B, N, 2, 2)\n\n diff = uv_dist - uv_dist_est\n a = inv[..., 0, 0]\n b = inv[..., 0, 1]\n c = inv[..., 1, 0]\n d = inv[..., 1, 1]\n e = diff[..., 0]\n f = diff[..., 1]\n step = torch.stack([a * e + b * f, c * e + d * f], dim=-1)\n\n # Newton step.\n xr_yr = xr_yr + step\n\n # Compute theta using Newton's method.\n xr_yr_norm = xr_yr.norm(p=2, dim=2).reshape(B, N, 1)\n th = xr_yr_norm.clone()\n max_iters_radial = max_iters if self.use_radial else 0\n for _ in range(max_iters_radial):\n th_radial = 1.0 + k1 * torch.pow(th, 2) + k2 * torch.pow(th, 4)\n dthd_th = 1.0 + 3.0 * k1 * torch.pow(th, 2) + 5.0 * k2 * torch.pow(th, 4)\n th_radial = th_radial * th\n step = (xr_yr_norm - th_radial) / dthd_th\n # handle dthd_th close to 0.\n step = torch.where(\n torch.abs(dthd_th) > eps, step, torch.sign(step) * eps * 10.0\n )\n th = th + step\n\n # Compute the ray direction using theta and xr_yr.\n close_to_zero = (torch.abs(th) < eps) & (torch.abs(xr_yr_norm) < eps)\n ray_dir = torch.where(close_to_zero, xr_yr, th * xr_yr / xr_yr_norm)\n\n # Compute the 3D projective ray\n rho2_u = (\n ray_dir.norm(p=2, dim=2, keepdim=True) ** 2\n ) # B N 1 # x_c * x_c + y_c * y_c\n xi = xi.reshape(B, 1, 1)\n sqrt_term = torch.sqrt(1.0 + (1.0 - xi * xi) * rho2_u)\n P_z = 1.0 - xi * (rho2_u + 1.0) / (xi + sqrt_term)\n\n # Special case when xi is 1.0 (unit sphere projection ??)\n P_z = torch.where(xi == 1.0, (1.0 - rho2_u) / 2.0, P_z)\n\n ray = torch.cat([ray_dir, P_z], dim=-1)\n ray = ray.reshape(B, H, W, 3).permute(0, 3, 1, 2)\n\n return ray\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def project(self, xyz):\n is_flat = xyz.ndim == 3\n B, N = xyz.shape[:2]\n\n if not is_flat:\n B, _, H, W = xyz.shape\n N = H * W\n xyz = xyz.permute(0, 2, 3, 1).reshape(B, N, 3)\n\n k1, k2, p0, p1, xi = self.params[..., 4:].unbind(dim=1)\n fx_fy = self.params[..., 0:2].reshape(B, 1, 2)\n cx_cy = self.params[..., 2:4].reshape(B, 1, 2)\n\n norm = xyz.norm(p=2, dim=-1, keepdim=True)\n ab = xyz[..., :-1] / (xyz[..., -1:] + xi.reshape(B, 1, 1) * norm)\n\n # radial correction\n r = ab.norm(dim=-1, p=2, keepdim=True)\n k1 = self.params[..., 4].reshape(B, 1, 1)\n k2 = self.params[..., 5].reshape(B, 1, 1)\n # ab / r * th * (1 + k1 * (th ** 2) + k2 * (th**4))\n # but here r = th, no spherical distortion\n xr_yr = ab * (1 + k1 * (r**2) + k2 * (r**4))\n\n # Tangential correction.\n uv_dist = xr_yr\n p0 = self.params[:, -3].reshape(B, 1)\n p1 = self.params[:, -2].reshape(B, 1)\n xr = xr_yr[..., 0].reshape(B, N)\n yr = xr_yr[..., 1].reshape(B, N)\n xr_yr_sq = torch.square(xr_yr)\n xr_sq = xr_yr_sq[:, :, 0].reshape(B, N)\n yr_sq = xr_yr_sq[:, :, 1].reshape(B, N)\n rd_sq = xr_sq + yr_sq\n uv_dist_tu = uv_dist[:, :, 0] + (\n (2.0 * xr_sq + rd_sq) * p0 + 2.0 * xr * yr * p1\n )\n uv_dist_tv = uv_dist[:, :, 1] + (\n (2.0 * yr_sq + rd_sq) * p1 + 2.0 * xr * yr * p0\n )\n uv_dist = torch.stack(\n [uv_dist_tu, uv_dist_tv], dim=-1\n ) # Avoids in-place complaint.\n\n result = uv_dist * fx_fy + cx_cy\n\n if not is_flat:\n result = result.reshape(B, H, W, 2).permute(0, 3, 1, 2)\n invalid = (\n (result[:, 0] < 0)\n | (result[:, 0] > W)\n | (result[:, 1] < 0)\n | (result[:, 1] > H)\n )\n self.projection_mask = (~invalid).unsqueeze(1)\n # creates hole in the middle... ??\n # self.overlap_mask = self.mask_overlap_projection(result)\n\n return result\n\n\nclass BatchCamera(Camera):\n \"\"\"\n This is not to be used directly, but to be used as a wrapper around multiple cameras.\n It should expose only the `from_camera` method as it the only way to create a BatchCamera.\n \"\"\"\n\n def __init__(self, params, K, original_class, cameras):\n super().__init__(params, K)\n self.original_class = original_class\n self.cameras = cameras\n\n # Delegate these methods to original camera\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def project(self, points_3d):\n return torch.cat(\n [\n camera.project(points_3d[i : i + 1])\n for i, camera in enumerate(self.cameras)\n ]\n )\n\n @torch.autocast(device_type=\"cuda\", enabled=False, dtype=torch.float32)\n def unproject(self, points_2d):\n val = torch.cat(\n [camera.unproject(points_2d) for i, camera in enumerate(self.cameras)]\n )\n return val\n\n def crop(self, left, top, right=None, bottom=None):\n val = torch.cat(\n [\n camera.crop(left, top, right, bottom)\n for i, camera in enumerate(self.cameras)\n ]\n )\n return val\n\n def resize(self, ratio):\n val = torch.cat([camera.resize(ratio) for i, camera in enumerate(self.cameras)])\n return val\n\n def reconstruct(self, depth):\n val = torch.cat(\n [\n camera.reconstruct(depth[i : i + 1])\n for i, camera in enumerate(self.cameras)\n ]\n )\n return val\n\n def get_projection_mask(self):\n return torch.cat(\n [camera.projection_mask for i, camera in enumerate(self.cameras)]\n )\n\n def to(self, device, non_blocking=False):\n self = super().to(device, non_blocking=non_blocking)\n self.cameras = recursive_to(\n self.cameras, device, non_blocking=non_blocking, cls=Camera\n )\n return self\n\n def reshape(self, *shape):\n # Reshape the intrinsic matrix (K) and params\n # we know that the shape of K is (..., 3, 3) and params is (..., 16)\n reshaped_K = self.K.reshape(*shape, 3, 3)\n reshaped_params = self.params.reshape(*shape, self.params.shape[-1])\n\n self.cameras = np.array(self.cameras, dtype=object).reshape(shape).tolist()\n self.original_class = (\n np.array(self.original_class, dtype=object).reshape(shape).tolist()\n )\n\n # Create a new BatchCamera with reshaped K and params\n return BatchCamera(\n reshaped_params, reshaped_K, self.original_class, self.cameras\n )\n\n def get_new_fov(self, new_shape, original_shape):\n return [\n camera.get_new_fov(new_shape, original_shape)\n for i, camera in enumerate(self.cameras)\n ]\n\n def squeeze(self, dim):\n return BatchCamera(\n self.params.squeeze(dim),\n self.K.squeeze(dim),\n squeeze_list(self.original_class, dim=dim),\n squeeze_list(self.cameras, dim=dim),\n )\n\n def __getitem__(self, idx):\n if isinstance(idx, int):\n return self.cameras[idx]\n\n elif isinstance(idx, slice):\n return BatchCamera(\n self.params[idx],\n self.K[idx],\n self.original_class[idx],\n self.cameras[idx],\n )\n\n raise TypeError(f\"Invalid index type: {type(idx)}\")\n\n def __setitem__(self, idx, value):\n # If it's an integer index, return a single camera\n if isinstance(idx, int):\n self.cameras[idx] = value\n self.params[idx, :] = 0.0\n self.params[idx, : value.params.shape[1]] = value.params[0]\n self.K[idx] = value.K[0]\n\n self.original_class[idx] = getattr(\n value, \"original_class\", value.__class__.__name__\n )\n\n # If it's a slice, return a new BatchCamera with sliced cameras\n elif isinstance(idx, slice):\n # Update each internal attribute using the slice\n self.params[idx] = value.params\n self.K[idx] = value.K\n self.original_class[idx] = value.original_class\n self.cameras[idx] = value.cameras\n\n def __len__(self):\n return len(self.cameras)\n\n @classmethod\n def from_camera(cls, camera):\n return cls(camera.params, camera.K, [camera.__class__.__name__], [camera])\n\n @property\n def is_perspective(self):\n return [isinstance(camera, Pinhole) for camera in self.cameras]\n\n @property\n def is_spherical(self):\n return [isinstance(camera, Spherical) for camera in self.cameras]\n\n @property\n def is_eucm(self):\n return [isinstance(camera, EUCM) for camera in self.cameras]\n\n @property\n def is_fisheye(self):\n return [isinstance(camera, Fisheye624) for camera in self.cameras]\n\n @property\n def is_pinhole(self):\n return [isinstance(camera, Pinhole) for camera in self.cameras]\n\n @property\n def hfov(self):\n return [camera.hfov for camera in self.cameras]\n\n @property\n def vfov(self):\n return [camera.vfov for camera in self.cameras]\n\n @property\n def max_fov(self):\n return [camera.max_fov for camera in self.cameras]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/chamfer_distance.py", "content": "import warnings\nfrom typing import Union\n\nimport torch\n\ntry:\n from unidepth.ops.knn import knn_points\nexcept ImportError as e:\n warnings.warn(\n \"!! To run evaluation you need KNN. Please compile KNN: \"\n \"`cd unidepth/ops/knn with && bash compile.sh`.\"\n )\n knn_points = lambda x : x\n\n\ndef _validate_chamfer_reduction_inputs(\n batch_reduction: Union[str, None], point_reduction: str\n):\n \"\"\"Check the requested reductions are valid.\n\n Args:\n batch_reduction: Reduction operation to apply for the loss across the\n batch, can be one of [\"mean\", \"sum\"] or None.\n point_reduction: Reduction operation to apply for the loss across the\n points, can be one of [\"mean\", \"sum\"].\n \"\"\"\n if batch_reduction is not None and batch_reduction not in [\"mean\", \"sum\"]:\n raise ValueError('batch_reduction must be one of [\"mean\", \"sum\"] or None')\n if point_reduction not in [\"mean\", \"sum\"]:\n raise ValueError('point_reduction must be one of [\"mean\", \"sum\"]')\n\n\ndef _handle_pointcloud_input(\n points: torch.Tensor,\n lengths: Union[torch.Tensor, None],\n normals: Union[torch.Tensor, None],\n):\n \"\"\"\n If points is an instance of Pointclouds, retrieve the padded points tensor\n along with the number of points per batch and the padded normals.\n Otherwise, return the input points (and normals) with the number of points per cloud\n set to the size of the second dimension of `points`.\n \"\"\"\n if points.ndim != 3:\n raise ValueError(\"Expected points to be of shape (N, P, D)\")\n X = points\n if lengths is not None and (lengths.ndim != 1 or lengths.shape[0] != X.shape[0]):\n raise ValueError(\"Expected lengths to be of shape (N,)\")\n if lengths is None:\n lengths = torch.full(\n (X.shape[0],), X.shape[1], dtype=torch.int64, device=points.device\n )\n if normals is not None and normals.ndim != 3:\n raise ValueError(\"Expected normals to be of shape (N, P, 3\")\n\n return X, lengths, normals\n\n\nclass ChamferDistance(torch.nn.Module):\n def forward(\n self,\n x,\n y,\n x_lengths=None,\n y_lengths=None,\n x_normals=None,\n y_normals=None,\n weights=None,\n batch_reduction: Union[str, None] = \"mean\",\n point_reduction: str = \"mean\",\n ):\n \"\"\"\n Chamfer distance between two pointclouds x and y.\n\n Args:\n x: FloatTensor of shape (N, P1, D) or a Pointclouds object representing\n a batch of point clouds with at most P1 points in each batch element,\n batch size N and feature dimension D.\n y: FloatTensor of shape (N, P2, D) or a Pointclouds object representing\n a batch of point clouds with at most P2 points in each batch element,\n batch size N and feature dimension D.\n x_lengths: Optional LongTensor of shape (N,) giving the number of points in each\n cloud in x.\n y_lengths: Optional LongTensor of shape (N,) giving the number of points in each\n cloud in x.\n x_normals: Optional FloatTensor of shape (N, P1, D).\n y_normals: Optional FloatTensor of shape (N, P2, D).\n weights: Optional FloatTensor of shape (N,) giving weights for\n batch elements for reduction operation.\n batch_reduction: Reduction operation to apply for the loss across the\n batch, can be one of [\"mean\", \"sum\"] or None.\n point_reduction: Reduction operation to apply for the loss across the\n points, can be one of [\"mean\", \"sum\"].\n\n Returns:\n 2-element tuple containing\n\n - **loss**: Tensor giving the reduced distance between the pointclouds\n in x and the pointclouds in y.\n - **loss_normals**: Tensor giving the reduced cosine distance of normals\n between pointclouds in x and pointclouds in y. Returns None if\n x_normals and y_normals are None.\n \"\"\"\n _validate_chamfer_reduction_inputs(batch_reduction, point_reduction)\n\n x, x_lengths, x_normals = _handle_pointcloud_input(x, x_lengths, x_normals)\n y, y_lengths, y_normals = _handle_pointcloud_input(y, y_lengths, y_normals)\n\n return_normals = x_normals is not None and y_normals is not None\n\n N, P1, D = x.shape\n P2 = y.shape[1]\n\n # Check if inputs are heterogeneous and create a lengths mask.\n is_x_heterogeneous = (x_lengths != P1).any()\n is_y_heterogeneous = (y_lengths != P2).any()\n x_mask = (\n torch.arange(P1, device=x.device)[None] >= x_lengths[:, None]\n ) # shape [N, P1]\n y_mask = (\n torch.arange(P2, device=y.device)[None] >= y_lengths[:, None]\n ) # shape [N, P2]\n\n if y.shape[0] != N or y.shape[2] != D:\n raise ValueError(\"y does not have the correct shape.\")\n if weights is not None:\n if weights.size(0) != N:\n raise ValueError(\"weights must be of shape (N,).\")\n if not (weights >= 0).all():\n raise ValueError(\"weights cannot be negative.\")\n if weights.sum() == 0.0:\n weights = weights.view(N, 1)\n if batch_reduction in [\"mean\", \"sum\"]:\n return (\n (x.sum((1, 2)) * weights).sum() * 0.0,\n (x.sum((1, 2)) * weights).sum() * 0.0,\n )\n return (\n (x.sum((1, 2)) * weights) * 0.0,\n (x.sum((1, 2)) * weights) * 0.0,\n )\n\n x_nn = knn_points(x, y, lengths1=x_lengths, lengths2=y_lengths, K=1)\n y_nn = knn_points(y, x, lengths1=y_lengths, lengths2=x_lengths, K=1)\n\n cham_x = x_nn.dists[..., 0] # (N, P1)\n cham_y = y_nn.dists[..., 0] # (N, P2)\n\n if is_x_heterogeneous:\n cham_x[x_mask] = 0.0\n if is_y_heterogeneous:\n cham_y[y_mask] = 0.0\n\n if weights is not None:\n cham_x *= weights.view(N, 1)\n cham_y *= weights.view(N, 1)\n\n return cham_x, cham_y, x_nn.idx[..., -1], y_nn.idx[..., -1]\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/constants.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport math\n\nimport torch\n\nOPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)\nOPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)\nIMAGENET_DATASET_MEAN = (0.485, 0.456, 0.406)\nIMAGENET_DATASET_STD = (0.229, 0.224, 0.225)\nDEPTH_BINS = torch.cat(\n (\n torch.logspace(math.log10(0.1), math.log10(180.0), steps=512),\n torch.tensor([260.0]),\n ),\n dim=0,\n)\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/coordinate.py", "content": "import torch\n\n\ndef coords_grid(b, h, w, homogeneous=False, device=None, noisy=False):\n pixel_coords_x = torch.linspace(0.5, w - 0.5, w, device=device)\n pixel_coords_y = torch.linspace(0.5, h - 0.5, h, device=device)\n if noisy: # \\pm 0.5px noise\n pixel_coords_x += torch.rand_like(pixel_coords_x) - 0.5\n pixel_coords_y += torch.rand_like(pixel_coords_y) - 0.5\n\n stacks = [pixel_coords_x.repeat(h, 1), pixel_coords_y.repeat(w, 1).t()]\n if homogeneous:\n ones = torch.ones_like(stacks[0]) # [H, W]\n stacks.append(ones)\n grid = torch.stack(stacks, dim=0).float() # [2, H, W] or [3, H, W]\n grid = grid[None].repeat(b, 1, 1, 1) # [B, 2, H, W] or [B, 3, H, W]\n if device is not None:\n grid = grid.to(device)\n\n return grid\n\n\ndef normalize_coords(coords, h, w):\n c = torch.tensor([(w - 1) / 2.0, (h - 1) / 2.0], device=coords.device).view(\n 1, 2, 1, 1\n )\n return (coords - c) / c\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/distributed.py", "content": "import os\nimport pickle\nimport platform\nimport subprocess\nimport warnings\n\nimport cv2\nimport torch\nimport torch.utils.data.distributed\nfrom torch import distributed as dist\nfrom torch import multiprocessing as mp\n\n_LOCAL_PROCESS_GROUP = None\n\n\ndef is_dist_avail_and_initialized():\n if not dist.is_available():\n return False\n if not dist.is_initialized():\n return False\n return True\n\n\ndef get_rank():\n if not is_dist_avail_and_initialized():\n return 0\n return dist.get_rank()\n\n\ndef get_local_rank() -> int:\n \"\"\"\n Returns:\n The rank of the current process within the local (per-machine) process group.\n \"\"\"\n if not is_dist_avail_and_initialized():\n return 0\n assert _LOCAL_PROCESS_GROUP is not None\n return dist.get_rank(group=_LOCAL_PROCESS_GROUP)\n\n\ndef get_local_size() -> int:\n \"\"\"\n Returns:\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 assert _LOCAL_PROCESS_GROUP is not None\n return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)\n\n\ndef get_world_size():\n if not is_dist_avail_and_initialized():\n return 1\n return dist.get_world_size()\n\n\ndef barrier():\n if not is_dist_avail_and_initialized():\n return\n dist.barrier()\n\n\ndef is_main_process():\n return get_rank() == 0\n\n\ndef is_rank_zero(args):\n return args.rank == 0\n\n\ndef get_dist_info():\n if dist.is_available() and dist.is_initialized():\n rank = dist.get_rank()\n world_size = dist.get_world_size()\n else:\n rank = 0\n world_size = 1\n return rank, world_size\n\n\ndef setup_multi_processes(cfg):\n \"\"\"Setup multi-processing environment variables.\"\"\"\n # set multi-process start method as `fork` to speed up the training\n if platform.system() != \"Windows\":\n mp_start_method = cfg.get(\"mp_start_method\", \"fork\")\n current_method = mp.get_start_method(allow_none=True)\n if current_method is not None and current_method != mp_start_method:\n warnings.warn(\n f\"Multi-processing start method `{mp_start_method}` is \"\n f\"different from the previous setting `{current_method}`.\"\n f\"It will be force set to `{mp_start_method}`. You can change \"\n f\"this behavior by changing `mp_start_method` in your config.\"\n )\n mp.set_start_method(mp_start_method, force=True)\n\n # disable opencv multithreading to avoid system being overloaded\n # opencv_num_threads = cfg.get('opencv_num_threads', 0)\n # cv2.setNumThreads(opencv_num_threads)\n\n # setup OMP threads\n # This code is referred from https://github.com/pytorch/pytorch/blob/master/torch/distributed/run.py # noqa\n # workers_per_gpu = cfg.get('workers_per_gpu', 4)\n\n # if 'OMP_NUM_THREADS' not in os.environ and workers_per_gpu > 1:\n # omp_num_threads = 1\n # warnings.warn(\n # f'Setting OMP_NUM_THREADS environment variable for each process '\n # f'to be {omp_num_threads} in default, to avoid your system being '\n # f'overloaded, please further tune the variable for optimal '\n # f'performance in your application as needed.')\n # os.environ['OMP_NUM_THREADS'] = str(omp_num_threads)\n\n # setup MKL threads\n # if 'MKL_NUM_THREADS' not in os.environ and workers_per_gpu > 1:\n # mkl_num_threads = os.environ.get('OMP_NUM_THREADS', 1)\n # warnings.warn(\n # f'Setting MKL_NUM_THREADS environment variable for each process '\n # f'to be {mkl_num_threads} in default, to avoid your system being '\n # f'overloaded, please further tune the variable for optimal '\n # f'performance in your application as needed.')\n # os.environ['MKL_NUM_THREADS'] = str(mkl_num_threads)\n\n\ndef setup_slurm(backend: str, port: str) -> None:\n proc_id = int(os.environ[\"SLURM_PROCID\"])\n ntasks = int(os.environ[\"SLURM_NTASKS\"])\n node_list = os.environ[\"SLURM_NODELIST\"]\n\n num_gpus = torch.cuda.device_count()\n\n torch.cuda.set_device(proc_id % num_gpus)\n addr = subprocess.getoutput(f\"scontrol show hostname {node_list} | head -n1\")\n os.environ[\"MASTER_PORT\"] = str(port)\n os.environ[\"MASTER_ADDR\"] = addr\n os.environ[\"WORLD_SIZE\"] = str(ntasks)\n os.environ[\"LOCAL_RANK\"] = str(proc_id % num_gpus)\n os.environ[\"RANK\"] = str(proc_id)\n print(\n proc_id,\n ntasks,\n num_gpus,\n proc_id % num_gpus,\n node_list,\n addr,\n os.environ[\"MASTER_PORT\"],\n os.system(\"nvidia-smi -L\"),\n )\n dist.init_process_group(backend, rank=proc_id, world_size=ntasks)\n\n\ndef sync_tensor_across_gpus(t, dim=0, cat=True):\n if t is None or not (dist.is_available() and dist.is_initialized()):\n return t\n t = torch.atleast_1d(t)\n group = dist.group.WORLD\n group_size = torch.distributed.get_world_size(group)\n\n local_size = torch.tensor(t.size(dim), device=t.device)\n all_sizes = [torch.zeros_like(local_size) for _ in range(group_size)]\n dist.all_gather(all_sizes, local_size)\n max_size = max(all_sizes)\n size_diff = max_size.item() - local_size.item()\n if size_diff:\n padding = torch.zeros(size_diff, device=t.device, dtype=t.dtype)\n t = torch.cat((t, padding))\n\n gather_t_tensor = [torch.zeros_like(t) for _ in range(group_size)]\n dist.all_gather(gather_t_tensor, t)\n all_ts = []\n for t, size in zip(gather_t_tensor, all_sizes):\n all_ts.append(t[:size])\n if cat:\n return torch.cat(all_ts, dim=0)\n return all_ts\n\n\ndef sync_string_across_gpus(keys: list[str], device, dim=0):\n keys_serialized = pickle.dumps(keys, protocol=pickle.HIGHEST_PROTOCOL)\n keys_serialized_tensor = (\n torch.frombuffer(keys_serialized, dtype=torch.uint8).clone().to(device)\n )\n keys_serialized_tensor = sync_tensor_across_gpus(\n keys_serialized_tensor, dim=0, cat=False\n )\n keys = [\n key\n for keys in keys_serialized_tensor\n for key in pickle.loads(bytes(keys.cpu().tolist()))\n ]\n return keys\n\n\ndef create_local_process_group() -> None:\n num_workers_per_machine = torch.cuda.device_count()\n global _LOCAL_PROCESS_GROUP\n assert _LOCAL_PROCESS_GROUP is None\n assert get_world_size() % num_workers_per_machine == 0\n num_machines = get_world_size() // num_workers_per_machine\n machine_rank = get_rank() // num_workers_per_machine\n for i in range(num_machines):\n ranks_on_i = list(\n range(i * num_workers_per_machine, (i + 1) * num_workers_per_machine)\n )\n pg = dist.new_group(ranks_on_i)\n if i == machine_rank:\n _LOCAL_PROCESS_GROUP = pg\n\n\ndef _get_global_gloo_group():\n if dist.get_backend() == \"nccl\":\n return dist.new_group(backend=\"gloo\")\n else:\n return dist.group.WORLD\n\n\ndef all_gather(data, group=None):\n if get_world_size() == 1:\n return [data]\n if group is None:\n group = (\n _get_global_gloo_group()\n ) # use CPU group by default, to reduce GPU RAM usage.\n world_size = dist.get_world_size(group)\n if world_size == 1:\n return [data]\n\n output = [None for _ in range(world_size)]\n dist.all_gather_object(output, data, group=group)\n return output\n\n\ndef local_broadcast_process_authkey():\n if get_local_size() == 1:\n return\n local_rank = get_local_rank()\n authkey = bytes(mp.current_process().authkey)\n all_keys = all_gather(authkey)\n local_leader_key = all_keys[get_rank() - local_rank]\n if authkey != local_leader_key:\n # print(\"Process authkey is different from the key of local leader! workers are launched independently ??\")\n # print(\"Overwriting local authkey ...\")\n mp.current_process().authkey = local_leader_key\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/ema_torch.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom __future__ import division, unicode_literals\n\nimport contextlib\nimport copy\nimport weakref\nfrom math import tanh\nfrom typing import Iterable, Optional\n\nimport torch\n\n\nclass DummyExponentialMovingAverage:\n def __init__(self, *args, **kwargs):\n pass\n\n def _get_parameters(self, *args, **kwargs):\n pass\n\n def get_current_decay(self, *args, **kwargs):\n pass\n\n def update(self, *args, **kwargs):\n pass\n\n def copy_to(self, *args, **kwargs):\n pass\n\n def store(self, *args, **kwargs):\n return\n\n def restore(self, *args, **kwargs):\n return\n\n @contextlib.contextmanager\n def average_parameters(self, *args, **kwargs):\n try:\n yield\n finally:\n pass\n\n def to(self, *args, **kwargs):\n pass\n\n def state_dict(self, *args, **kwargs):\n pass\n\n def load_state_dict(self, *args, **kwargs):\n pass\n\n\nclass ExponentialMovingAverage:\n \"\"\"\n Maintains (exponential) moving average of a set of parameters.\n\n Args:\n parameters: Iterable of `torch.nn.Parameter` (typically from\n `model.parameters()`).\n Note that EMA is computed on *all* provided parameters,\n regardless of whether or not they have `requires_grad = True`;\n this allows a single EMA object to be consistantly used even\n if which parameters are trainable changes step to step.\n\n If you want to some parameters in the EMA, do not pass them\n to the object in the first place. For example:\n\n ExponentialMovingAverage(\n parameters=[p for p in model.parameters() if p.requires_grad],\n decay=0.9\n )\n\n will ignore parameters that do not require grad.\n\n decay: The exponential decay.\n\n use_num_updates: Whether to use number of updates when computing\n averages.\n \"\"\"\n\n def __init__(\n self,\n parameters: Iterable[torch.nn.Parameter],\n decay: float,\n use_num_updates: bool = True,\n update_after_step: int = 10000,\n tau: int = 20000,\n switch: bool = False,\n ):\n if decay < 0.0 or decay > 1.0:\n raise ValueError(\"Decay must be between 0 and 1\")\n self.decay = decay\n self.switch = switch # fi keeping EMA params in model after epochs\n self.num_updates = 0 if use_num_updates else None\n parameters = list(parameters)\n self.shadow_params = [p.clone().detach() for p in parameters]\n self.collected_params = None\n # By maintaining only a weakref to each parameter,\n # we maintain the old GC behaviour of ExponentialMovingAverage:\n # if the model goes out of scope but the ExponentialMovingAverage\n # is kept, no references to the model or its parameters will be\n # maintained, and the model will be cleaned up.\n self._params_refs = [weakref.ref(p) for p in parameters]\n self.update_after_step = update_after_step\n self.tau = tau\n\n def _get_parameters(\n self, parameters: Optional[Iterable[torch.nn.Parameter]]\n ) -> Iterable[torch.nn.Parameter]:\n if parameters is None:\n parameters = [p() for p in self._params_refs]\n if any(p is None for p in parameters):\n raise ValueError(\n \"(One of) the parameters with which this ExponentialMovingAverage was initialized no longer exists (was garbage collected);\"\n \" please either provide `parameters` explicitly or keep the model to which they belong from being garbage collected.\"\n )\n return parameters\n else:\n parameters = list(parameters)\n if len(parameters) != len(self.shadow_params):\n raise ValueError(\n \"Number of parameters passed as argument is different \"\n \"from number of shadow parameters maintained by this \"\n \"ExponentialMovingAverage\"\n )\n return parameters\n\n def get_current_decay(self):\n epoch = max(self.num_updates - self.update_after_step - 1, 0.0)\n if epoch <= 0:\n return 0.0\n value = tanh(epoch / self.tau) * self.decay\n return value\n\n def update(self, parameters: Optional[Iterable[torch.nn.Parameter]] = None) -> None:\n \"\"\"\n Update currently maintained parameters.\n\n Call this every time the parameters are updated, such as the result of\n the `optimizer.step()` call.\n\n Args:\n parameters: Iterable of `torch.nn.Parameter`; usually the same set of\n parameters used to initialize this object. If `None`, the\n parameters with which this `ExponentialMovingAverage` was\n initialized will be used.\n \"\"\"\n parameters = self._get_parameters(parameters)\n decay = self.get_current_decay()\n if self.num_updates is not None:\n self.num_updates += 1\n\n one_minus_decay = 1.0 - decay\n with torch.no_grad():\n for s_param, param in zip(self.shadow_params, parameters):\n tmp = s_param - param\n # tmp will be a new tensor so we can do in-place\n tmp.mul_(one_minus_decay)\n s_param.sub_(tmp)\n\n def copy_to(\n self, parameters: Optional[Iterable[torch.nn.Parameter]] = None\n ) -> None:\n \"\"\"\n Copy current averaged parameters into given collection of parameters.\n\n Args:\n parameters: Iterable of `torch.nn.Parameter`; the parameters to be\n updated with the stored moving averages. If `None`, the\n parameters with which this `ExponentialMovingAverage` was\n initialized will be used.\n \"\"\"\n parameters = self._get_parameters(parameters)\n for s_param, param in zip(self.shadow_params, parameters):\n param.data.copy_(s_param.data)\n\n def store(self, parameters: Optional[Iterable[torch.nn.Parameter]] = None) -> None:\n \"\"\"\n Save the current parameters for restoring later.\n\n Args:\n parameters: Iterable of `torch.nn.Parameter`; the parameters to be\n temporarily stored. If `None`, the parameters of with which this\n `ExponentialMovingAverage` was initialized will be used.\n \"\"\"\n parameters = self._get_parameters(parameters)\n self.collected_params = [param.detach().clone() for param in parameters]\n\n def restore(\n self, parameters: Optional[Iterable[torch.nn.Parameter]] = None\n ) -> None:\n \"\"\"\n Restore the parameters stored with the `store` method.\n Useful to validate the model with EMA parameters without affecting the\n original optimization process. Store the parameters before the\n `copy_to` method. After validation (or model saving), use this to\n restore the former parameters.\n\n Args:\n parameters: Iterable of `torch.nn.Parameter`; the parameters to be\n updated with the stored parameters. If `None`, the\n parameters with which this `ExponentialMovingAverage` was\n initialized will be used.\n \"\"\"\n if self.collected_params is None:\n raise RuntimeError(\n \"This ExponentialMovingAverage has no `store()`ed weights \"\n \"to `restore()`\"\n )\n parameters = self._get_parameters(parameters)\n for c_param, param in zip(self.collected_params, parameters):\n param.data.copy_(c_param.data)\n\n @contextlib.contextmanager\n def average_parameters(\n self, parameters: Optional[Iterable[torch.nn.Parameter]] = None\n ):\n r\"\"\"\n Context manager for validation/inference with averaged parameters.\n\n Equivalent to:\n\n ema.store()\n ema.copy_to()\n try:\n ...\n finally:\n ema.restore()\n\n Args:\n parameters: Iterable of `torch.nn.Parameter`; the parameters to be\n updated with the stored parameters. If `None`, the\n parameters with which this `ExponentialMovingAverage` was\n initialized will be used.\n \"\"\"\n parameters = self._get_parameters(parameters)\n self.store(parameters)\n self.copy_to(parameters)\n try:\n yield\n finally:\n if not self.switch:\n self.restore(parameters)\n\n def to(self, device=None, dtype=None) -> None:\n r\"\"\"Move internal buffers of the ExponentialMovingAverage to `device`.\n\n Args:\n device: like `device` argument to `torch.Tensor.to`\n \"\"\"\n # .to() on the tensors handles None correctly\n self.shadow_params = [\n (\n p.to(device=device, dtype=dtype)\n if p.is_floating_point()\n else p.to(device=device)\n )\n for p in self.shadow_params\n ]\n if self.collected_params is not None:\n self.collected_params = [\n (\n p.to(device=device, dtype=dtype)\n if p.is_floating_point()\n else p.to(device=device)\n )\n for p in self.collected_params\n ]\n return\n\n def state_dict(self) -> dict:\n r\"\"\"Returns the state of the ExponentialMovingAverage as a dict.\"\"\"\n # Following PyTorch conventions, references to tensors are returned:\n # \"returns a reference to the state and not its copy!\" -\n # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict\n return {\n \"decay\": self.decay,\n \"num_updates\": self.num_updates,\n \"shadow_params\": self.shadow_params,\n \"collected_params\": self.collected_params,\n }\n\n def load_state_dict(self, state_dict: dict) -> None:\n r\"\"\"Loads the ExponentialMovingAverage state.\n\n Args:\n state_dict (dict): EMA state. Should be an object returned\n from a call to :meth:`state_dict`.\n \"\"\"\n # deepcopy, to be consistent with module API\n state_dict = copy.deepcopy(state_dict)\n self.decay = state_dict[\"decay\"]\n if self.decay < 0.0 or self.decay > 1.0:\n raise ValueError(\"Decay must be between 0 and 1\")\n self.num_updates = state_dict[\"num_updates\"]\n assert self.num_updates is None or isinstance(\n self.num_updates, int\n ), \"Invalid num_updates\"\n\n self.shadow_params = state_dict[\"shadow_params\"]\n assert isinstance(self.shadow_params, list), \"shadow_params must be a list\"\n assert all(\n isinstance(p, torch.Tensor) for p in self.shadow_params\n ), \"shadow_params must all be Tensors\"\n\n self.collected_params = state_dict[\"collected_params\"]\n if self.collected_params is not None:\n assert isinstance(\n self.collected_params, list\n ), \"collected_params must be a list\"\n assert all(\n isinstance(p, torch.Tensor) for p in self.collected_params\n ), \"collected_params must all be Tensors\"\n assert len(self.collected_params) == len(\n self.shadow_params\n ), \"collected_params and shadow_params had different lengths\"\n\n if len(self.shadow_params) == len(self._params_refs):\n # Consistant with torch.optim.Optimizer, cast things to consistant\n # device and dtype with the parameters\n params = [p() for p in self._params_refs]\n # If parameters have been garbage collected, just load the state\n # we were given without change.\n if not any(p is None for p in params):\n # ^ parameter references are still good\n for i, p in enumerate(params):\n self.shadow_params[i] = self.shadow_params[i].to(\n device=p.device, dtype=p.dtype\n )\n if self.collected_params is not None:\n self.collected_params[i] = self.collected_params[i].to(\n device=p.device, dtype=p.dtype\n )\n else:\n raise ValueError(\n \"Tried to `load_state_dict()` with the wrong number of \"\n \"parameters in the saved state.\"\n )\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/evaluation_depth.py", "content": "from collections import defaultdict\nfrom functools import partial\n\nimport torch\nimport torch.nn.functional as F\n\nfrom unidepth.utils.chamfer_distance import ChamferDistance\n\nchamfer_cls = ChamferDistance()\n\n\ndef chamfer_dist(tensor1, tensor2):\n x_lengths = torch.tensor((tensor1.shape[1],), device=tensor1.device)\n y_lengths = torch.tensor((tensor2.shape[1],), device=tensor2.device)\n dist1, dist2, idx1, idx2 = chamfer_cls(\n tensor1, tensor2, x_lengths=x_lengths, y_lengths=y_lengths\n )\n return (torch.sqrt(dist1) + torch.sqrt(dist2)) / 2\n\n\ndef auc(tensor1, tensor2, thresholds):\n x_lengths = torch.tensor((tensor1.shape[1],), device=tensor1.device)\n y_lengths = torch.tensor((tensor2.shape[1],), device=tensor2.device)\n dist1, dist2, idx1, idx2 = chamfer_cls(\n tensor1, tensor2, x_lengths=x_lengths, y_lengths=y_lengths\n )\n # compute precision recall\n precisions = [(dist1 < threshold).sum() / dist1.numel() for threshold in thresholds]\n recalls = [(dist2 < threshold).sum() / dist2.numel() for threshold in thresholds]\n auc_value = torch.trapz(\n torch.tensor(precisions, device=tensor1.device),\n torch.tensor(recalls, device=tensor1.device),\n )\n return auc_value\n\n\ndef delta(tensor1, tensor2, exponent):\n inlier = torch.maximum((tensor1 / tensor2), (tensor2 / tensor1))\n return (inlier < 1.25**exponent).to(torch.float32).mean()\n\n\ndef tau(tensor1, tensor2, perc):\n inlier = torch.maximum((tensor1 / tensor2), (tensor2 / tensor1))\n return (inlier < (1.0 + perc)).to(torch.float32).mean()\n\n\ndef ssi(tensor1, tensor2):\n stability_mat = 1e-9 * torch.eye(2, device=tensor1.device)\n tensor2_one = torch.stack(\n [tensor2.detach(), torch.ones_like(tensor2).detach()], dim=1\n )\n scale_shift = torch.inverse(tensor2_one.T @ tensor2_one + stability_mat) @ (\n tensor2_one.T @ tensor1.unsqueeze(1)\n )\n scale, shift = scale_shift.squeeze().chunk(2, dim=0)\n return tensor2 * scale + shift\n\n\ndef si(tensor1, tensor2):\n return tensor2 * torch.median(tensor1) / torch.median(tensor2)\n\n\ndef arel(tensor1, tensor2):\n tensor2 = tensor2 * torch.median(tensor1) / torch.median(tensor2)\n return (torch.abs(tensor1 - tensor2) / tensor1).mean()\n\n\ndef d_auc(tensor1, tensor2):\n exponents = torch.linspace(0.01, 5.0, steps=100, device=tensor1.device)\n deltas = [delta(tensor1, tensor2, exponent) for exponent in exponents]\n return torch.trapz(torch.tensor(deltas, device=tensor1.device), exponents) / 5.0\n\n\ndef f1_score(tensor1, tensor2, thresholds):\n x_lengths = torch.tensor((tensor1.shape[1],), device=tensor1.device)\n y_lengths = torch.tensor((tensor2.shape[1],), device=tensor2.device)\n dist1, dist2, idx1, idx2 = chamfer_cls(\n tensor1, tensor2, x_lengths=x_lengths, y_lengths=y_lengths\n )\n # compute precision recall\n precisions = [(dist1 < threshold).sum() / dist1.numel() for threshold in thresholds]\n recalls = [(dist2 < threshold).sum() / dist2.numel() for threshold in thresholds]\n precisions = torch.tensor(precisions, device=tensor1.device)\n recalls = torch.tensor(recalls, device=tensor1.device)\n f1_thresholds = 2 * precisions * recalls / (precisions + recalls)\n f1_thresholds = torch.where(\n torch.isnan(f1_thresholds), torch.zeros_like(f1_thresholds), f1_thresholds\n )\n f1_value = torch.trapz(f1_thresholds) / len(thresholds)\n return f1_value\n\n\nDICT_METRICS = {\n \"d1\": partial(delta, exponent=1.0),\n \"d2\": partial(delta, exponent=2.0),\n \"d3\": partial(delta, exponent=3.0),\n \"rmse\": lambda gt, pred: torch.sqrt(((gt - pred) ** 2).mean()),\n \"rmselog\": lambda gt, pred: torch.sqrt(\n ((torch.log(gt) - torch.log(pred)) ** 2).mean()\n ),\n \"arel\": lambda gt, pred: (torch.abs(gt - pred) / gt).mean(),\n \"sqrel\": lambda gt, pred: (((gt - pred) ** 2) / gt).mean(),\n \"log10\": lambda gt, pred: torch.abs(torch.log10(pred) - torch.log10(gt)).mean(),\n \"silog\": lambda gt, pred: 100 * torch.std(torch.log(pred) - torch.log(gt)).mean(),\n \"medianlog\": lambda gt, pred: 100\n * (torch.log(pred) - torch.log(gt)).median().abs(),\n \"d_auc\": d_auc,\n \"tau\": partial(tau, perc=0.03),\n}\n\n\nDICT_METRICS_3D = {\n \"MSE_3d\": lambda gt, pred, thresholds: torch.norm(gt - pred, dim=0, p=2),\n \"chamfer\": lambda gt, pred, thresholds: chamfer_dist(\n gt.unsqueeze(0).permute(0, 2, 1), pred.unsqueeze(0).permute(0, 2, 1)\n ),\n \"F1\": lambda gt, pred, thresholds: f1_score(\n gt.unsqueeze(0).permute(0, 2, 1),\n pred.unsqueeze(0).permute(0, 2, 1),\n thresholds=thresholds,\n ),\n}\n\nDICT_METRICS_D = {\n \"a1\": lambda gt, pred: (torch.maximum((gt / pred), (pred / gt)) > 1.25**1.0).to(\n torch.float32\n ),\n \"abs_rel\": lambda gt, pred: (torch.abs(gt - pred) / gt),\n}\n\n\ndef eval_depth(\n gts: torch.Tensor, preds: torch.Tensor, masks: torch.Tensor, max_depth=None\n):\n summary_metrics = defaultdict(list)\n preds = F.interpolate(preds, gts.shape[-2:], mode=\"bilinear\")\n for i, (gt, pred, mask) in enumerate(zip(gts, preds, masks)):\n if max_depth is not None:\n mask = mask & (gt <= max_depth)\n for name, fn in DICT_METRICS.items():\n if name in [\"tau\", \"d1\", \"arel\"]:\n for rescale_fn in [\"ssi\", \"si\"]:\n summary_metrics[f\"{name}_{rescale_fn}\"].append(\n fn(gt[mask], eval(rescale_fn)(gt[mask], pred[mask]))\n )\n summary_metrics[name].append(fn(gt[mask], pred[mask]).mean())\n return {name: torch.stack(vals, dim=0) for name, vals in summary_metrics.items()}\n\n\ndef eval_3d(\n gts: torch.Tensor, preds: torch.Tensor, masks: torch.Tensor, thresholds=None\n):\n summary_metrics = defaultdict(list)\n ratio = min(\n 1.0, (240 * 320 / masks.sum()) ** 0.5\n ) # rescale to avoid OOM during eval, FIXME\n h_max, w_max = int(gts.shape[-2] * ratio), int(gts.shape[-1] * ratio)\n gts = F.interpolate(gts, size=(h_max, w_max), mode=\"nearest-exact\")\n preds = F.interpolate(preds, size=(h_max, w_max), mode=\"nearest-exact\")\n masks = F.interpolate(\n masks.float(), size=(h_max, w_max), mode=\"nearest-exact\"\n ).bool()\n for i, (gt, pred, mask) in enumerate(zip(gts, preds, masks)):\n if not torch.any(mask):\n continue\n for name, fn in DICT_METRICS_3D.items():\n summary_metrics[name].append(\n fn(gt[:, mask.squeeze()], pred[:, mask.squeeze()], thresholds).mean()\n )\n return {name: torch.stack(vals, dim=0) for name, vals in summary_metrics.items()}\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/geometric.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom typing import Tuple\n\nimport torch\nfrom torch.nn import functional as F\n\n\n@torch.jit.script\ndef generate_rays(\n camera_intrinsics: torch.Tensor, image_shape: Tuple[int, int], noisy: bool = False\n):\n batch_size, device, dtype = (\n camera_intrinsics.shape[0],\n camera_intrinsics.device,\n camera_intrinsics.dtype,\n )\n height, width = image_shape\n # Generate grid of pixel coordinates\n pixel_coords_x = torch.linspace(0, width - 1, width, device=device, dtype=dtype)\n pixel_coords_y = torch.linspace(0, height - 1, height, device=device, dtype=dtype)\n if noisy:\n pixel_coords_x += torch.rand_like(pixel_coords_x) - 0.5\n pixel_coords_y += torch.rand_like(pixel_coords_y) - 0.5\n pixel_coords = torch.stack(\n [pixel_coords_x.repeat(height, 1), pixel_coords_y.repeat(width, 1).t()], dim=2\n ) # (H, W, 2)\n pixel_coords = pixel_coords + 0.5\n\n # Calculate ray directions\n intrinsics_inv = torch.eye(3, device=device).unsqueeze(0).repeat(batch_size, 1, 1)\n intrinsics_inv[:, 0, 0] = 1.0 / camera_intrinsics[:, 0, 0]\n intrinsics_inv[:, 1, 1] = 1.0 / camera_intrinsics[:, 1, 1]\n intrinsics_inv[:, 0, 2] = -camera_intrinsics[:, 0, 2] / camera_intrinsics[:, 0, 0]\n intrinsics_inv[:, 1, 2] = -camera_intrinsics[:, 1, 2] / camera_intrinsics[:, 1, 1]\n homogeneous_coords = torch.cat(\n [pixel_coords, torch.ones_like(pixel_coords[:, :, :1])], dim=2\n ) # (H, W, 3)\n ray_directions = torch.matmul(\n intrinsics_inv, homogeneous_coords.permute(2, 0, 1).flatten(1)\n ) # (3, H*W)\n ray_directions = F.normalize(ray_directions, dim=1) # (B, 3, H*W)\n ray_directions = ray_directions.permute(0, 2, 1) # (B, H*W, 3)\n\n theta = torch.atan2(ray_directions[..., 0], ray_directions[..., -1])\n phi = torch.acos(ray_directions[..., 1])\n # pitch = torch.asin(ray_directions[..., 1])\n # roll = torch.atan2(ray_directions[..., 0], - ray_directions[..., 1])\n angles = torch.stack([theta, phi], dim=-1)\n return ray_directions, angles\n\n\n@torch.jit.script\ndef spherical_zbuffer_to_euclidean(spherical_tensor: torch.Tensor) -> torch.Tensor:\n theta = spherical_tensor[..., 0] # Extract polar angle\n phi = spherical_tensor[..., 1] # Extract azimuthal angle\n z = spherical_tensor[..., 2] # Extract zbuffer depth\n\n # y = r * cos(phi)\n # x = r * sin(phi) * sin(theta)\n # z = r * sin(phi) * cos(theta)\n # =>\n # r = z / sin(phi) / cos(theta)\n # y = z / (sin(phi) / cos(phi)) / cos(theta)\n # x = z * sin(theta) / cos(theta)\n x = z * torch.tan(theta)\n y = z / torch.tan(phi) / torch.cos(theta)\n\n euclidean_tensor = torch.stack((x, y, z), dim=-1)\n return euclidean_tensor\n\n\n@torch.jit.script\ndef spherical_to_euclidean(spherical_tensor: torch.Tensor) -> torch.Tensor:\n theta = spherical_tensor[..., 0] # Extract polar angle\n phi = spherical_tensor[..., 1] # Extract azimuthal angle\n r = spherical_tensor[..., 2] # Extract radius\n # y = r * cos(phi)\n # x = r * sin(phi) * sin(theta)\n # z = r * sin(phi) * cos(theta)\n x = r * torch.sin(phi) * torch.sin(theta)\n y = r * torch.cos(phi)\n z = r * torch.cos(theta) * torch.sin(phi)\n\n euclidean_tensor = torch.stack((x, y, z), dim=-1)\n return euclidean_tensor\n\n\n@torch.jit.script\ndef euclidean_to_spherical(spherical_tensor: torch.Tensor) -> torch.Tensor:\n x = spherical_tensor[..., 0] # Extract polar angle\n y = spherical_tensor[..., 1] # Extract azimuthal angle\n z = spherical_tensor[..., 2] # Extract radius\n # y = r * cos(phi)\n # x = r * sin(phi) * sin(theta)\n # z = r * sin(phi) * cos(theta)\n r = torch.sqrt(x**2 + y**2 + z**2)\n theta = torch.atan2(x / r, z / r)\n phi = torch.acos(y / r)\n\n euclidean_tensor = torch.stack((theta, phi, r), dim=-1)\n return euclidean_tensor\n\n\n@torch.jit.script\ndef euclidean_to_spherical_zbuffer(euclidean_tensor: torch.Tensor) -> torch.Tensor:\n pitch = torch.asin(euclidean_tensor[..., 1])\n yaw = torch.atan2(euclidean_tensor[..., 0], euclidean_tensor[..., -1])\n z = euclidean_tensor[..., 2] # Extract zbuffer depth\n euclidean_tensor = torch.stack((pitch, yaw, z), dim=-1)\n return euclidean_tensor\n\n\n@torch.jit.script\ndef unproject_points(\n depth: torch.Tensor, camera_intrinsics: torch.Tensor\n) -> torch.Tensor:\n \"\"\"\n Unprojects a batch of depth maps to 3D point clouds using camera intrinsics.\n\n Args:\n depth (torch.Tensor): Batch of depth maps of shape (B, 1, H, W).\n camera_intrinsics (torch.Tensor): Camera intrinsic matrix of shape (B, 3, 3).\n\n Returns:\n torch.Tensor: Batch of 3D point clouds of shape (B, 3, H, W).\n \"\"\"\n batch_size, _, height, width = depth.shape\n device = depth.device\n\n # Create pixel grid\n y_coords, x_coords = torch.meshgrid(\n torch.arange(height, device=device),\n torch.arange(width, device=device),\n indexing=\"ij\",\n )\n pixel_coords = torch.stack((x_coords, y_coords), dim=-1) # (H, W, 2)\n\n # Get homogeneous coords (u v 1)\n pixel_coords_homogeneous = torch.cat(\n (pixel_coords, torch.ones((height, width, 1), device=device)), dim=-1\n )\n pixel_coords_homogeneous = pixel_coords_homogeneous.permute(2, 0, 1).flatten(\n 1\n ) # (3, H*W)\n # Apply K^-1 @ (u v 1): [B, 3, 3] @ [3, H*W] -> [B, 3, H*W]\n unprojected_points = torch.matmul(\n torch.inverse(camera_intrinsics), pixel_coords_homogeneous\n ) # (B, 3, H*W)\n unprojected_points = unprojected_points.view(\n batch_size, 3, height, width\n ) # (B, 3, H, W)\n unprojected_points = unprojected_points * depth # (B, 3, H, W)\n return unprojected_points\n\n\n@torch.jit.script\ndef project_points(\n points_3d: torch.Tensor,\n intrinsic_matrix: torch.Tensor,\n image_shape: Tuple[int, int],\n) -> torch.Tensor:\n # Project 3D points onto the image plane via intrinsics (u v w) = (x y z) @ K^T\n points_2d = torch.matmul(points_3d, intrinsic_matrix.transpose(1, 2))\n\n # Normalize projected points: (u v w) -> (u / w, v / w, 1)\n points_2d = points_2d[..., :2] / points_2d[..., 2:]\n\n points_2d = points_2d.int()\n\n # points need to be inside the image (can it diverge onto all points out???)\n valid_mask = (\n (points_2d[..., 0] >= 0)\n & (points_2d[..., 0] < image_shape[1])\n & (points_2d[..., 1] >= 0)\n & (points_2d[..., 1] < image_shape[0])\n )\n\n # Calculate the flat indices of the valid pixels\n flat_points_2d = points_2d[..., 0] + points_2d[..., 1] * image_shape[1]\n flat_indices = flat_points_2d.long()\n\n # Create depth maps and counts using scatter_add, (B, H, W)\n depth_maps = torch.zeros(\n [points_3d.shape[0], *image_shape], device=points_3d.device\n )\n counts = torch.zeros([points_3d.shape[0], *image_shape], device=points_3d.device)\n\n # Loop over batches to apply masks and accumulate depth/count values\n for i in range(points_3d.shape[0]):\n valid_indices = flat_indices[i, valid_mask[i]]\n depth_maps[i].view(-1).scatter_add_(\n 0, valid_indices, points_3d[i, valid_mask[i], 2]\n )\n counts[i].view(-1).scatter_add_(\n 0, valid_indices, torch.ones_like(points_3d[i, valid_mask[i], 2])\n )\n\n # Calculate mean depth for each pixel in each batch\n mean_depth_maps = depth_maps / counts.clamp(min=1.0)\n return mean_depth_maps.reshape(-1, 1, *image_shape) # (B, 1, H, W)\n\n\n@torch.jit.script\ndef downsample(data: torch.Tensor, downsample_factor: int = 2):\n N, _, H, W = data.shape\n data = data.view(\n N,\n H // downsample_factor,\n downsample_factor,\n W // downsample_factor,\n downsample_factor,\n 1,\n )\n data = data.permute(0, 1, 3, 5, 2, 4).contiguous()\n data = data.view(-1, downsample_factor * downsample_factor)\n data_tmp = torch.where(data == 0.0, 1e5 * torch.ones_like(data), data)\n data = torch.min(data_tmp, dim=-1).values\n data = data.view(N, 1, H // downsample_factor, W // downsample_factor)\n data = torch.where(data > 1000, torch.zeros_like(data), data)\n return data\n\n\n@torch.jit.script\ndef flat_interpolate(\n flat_tensor: torch.Tensor,\n old: Tuple[int, int],\n new: Tuple[int, int],\n antialias: bool = True,\n mode: str = \"bilinear\",\n) -> torch.Tensor:\n if old[0] == new[0] and old[1] == new[1]:\n return flat_tensor\n tensor = flat_tensor.view(flat_tensor.shape[0], old[0], old[1], -1).permute(\n 0, 3, 1, 2\n ) # b c h w\n tensor_interp = F.interpolate(\n tensor,\n size=(new[0], new[1]),\n mode=mode,\n align_corners=False,\n antialias=antialias,\n )\n flat_tensor_interp = tensor_interp.view(\n flat_tensor.shape[0], -1, new[0] * new[1]\n ).permute(\n 0, 2, 1\n ) # b (h w) c\n return flat_tensor_interp.contiguous()\n\n\n@torch.jit.script\ndef dilate(image, kernel_size: int | tuple[int, int]):\n if isinstance(kernel_size, int):\n kernel_size = (kernel_size, kernel_size)\n device, dtype = image.device, image.dtype\n padding = (kernel_size[0] // 2, kernel_size[1] // 2)\n kernel = torch.ones((1, 1, *kernel_size), dtype=torch.float32, device=image.device)\n dilated_image = F.conv2d(image.float(), kernel, padding=padding, stride=1)\n dilated_image = torch.where(\n dilated_image > 0,\n torch.tensor(1.0, device=device),\n torch.tensor(0.0, device=device),\n )\n return dilated_image.to(dtype)\n\n\n@torch.jit.script\ndef erode(image, kernel_size: int | tuple[int, int]):\n if isinstance(kernel_size, int):\n kernel_size = (kernel_size, kernel_size)\n device, dtype = image.device, image.dtype\n padding = (kernel_size[0] // 2, kernel_size[1] // 2)\n kernel = torch.ones((1, 1, *kernel_size), dtype=torch.float32, device=image.device)\n eroded_image = F.conv2d(image.float(), kernel, padding=padding, stride=1)\n eroded_image = torch.where(\n eroded_image == (kernel_size[0] * kernel_size[1]),\n torch.tensor(1.0, device=device),\n torch.tensor(0.0, device=device),\n )\n return eroded_image.to(dtype)\n\n\n@torch.jit.script\ndef iou(mask1: torch.Tensor, mask2: torch.Tensor) -> torch.Tensor:\n device = mask1.device\n\n # Ensure the masks are binary (0 or 1)\n mask1 = mask1.to(torch.bool)\n mask2 = mask2.to(torch.bool)\n\n # Compute intersection and union\n intersection = torch.sum(mask1 & mask2).to(torch.float32)\n union = torch.sum(mask1 | mask2).to(torch.float32)\n\n # Compute IoU\n iou = intersection / union.clip(min=1.0)\n\n return iou\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/misc.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom functools import wraps\nfrom time import time\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange, reduce, repeat\nfrom scipy import interpolate\n\n\n@torch.jit.script\ndef max_stack(tensors: list[torch.Tensor]) -> torch.Tensor:\n if len(tensors) == 1:\n return tensors[0]\n return torch.stack(tensors, dim=-1).max(dim=-1).values\n\n\ndef last_stack(tensors: list[torch.Tensor]) -> torch.Tensor:\n return tensors[-1]\n\n\ndef first_stack(tensors: list[torch.Tensor]) -> torch.Tensor:\n return tensors[0]\n\n\n@torch.jit.script\ndef softmax_stack(\n tensors: list[torch.Tensor], temperature: float = 1.0\n) -> torch.Tensor:\n if len(tensors) == 1:\n return tensors[0]\n return F.softmax(torch.stack(tensors, dim=-1) / temperature, dim=-1).sum(dim=-1)\n\n\n@torch.jit.script\ndef mean_stack(tensors: list[torch.Tensor]) -> torch.Tensor:\n if len(tensors) == 1:\n return tensors[0]\n return torch.stack(tensors, dim=-1).mean(dim=-1)\n\n\n@torch.jit.script\ndef sum_stack(tensors: list[torch.Tensor]) -> torch.Tensor:\n if len(tensors) == 1:\n return tensors[0]\n return torch.stack(tensors, dim=-1).sum(dim=-1)\n\n\ndef convert_module_to_f16(l):\n \"\"\"\n Convert primitive modules to float16.\n \"\"\"\n if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):\n l.weight.data = l.weight.data.half()\n if l.bias is not None:\n l.bias.data = l.bias.data.half()\n\n\ndef convert_module_to_f32(l):\n \"\"\"\n Convert primitive modules to float32, undoing convert_module_to_f16().\n \"\"\"\n if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):\n l.weight.data = l.weight.data.float()\n if l.bias is not None:\n l.bias.data = l.bias.data.float()\n\n\ndef format_seconds(seconds):\n minutes, seconds = divmod(seconds, 60)\n hours, minutes = divmod(minutes, 60)\n return f\"{hours:d}:{minutes:02d}:{seconds:02d}\"\n\n\ndef get_params(module, lr, wd):\n skip_list = {}\n skip_keywords = {}\n if hasattr(module, \"no_weight_decay\"):\n skip_list = module.no_weight_decay()\n if hasattr(module, \"no_weight_decay_keywords\"):\n skip_keywords = module.no_weight_decay_keywords()\n has_decay = []\n no_decay = []\n for name, param in module.named_parameters():\n if not param.requires_grad:\n continue # frozen weights\n if (\n (name in skip_list)\n or any((kw in name for kw in skip_keywords))\n or len(param.shape) == 1\n or name.endswith(\".gamma\")\n or name.endswith(\".beta\")\n or name.endswith(\".bias\")\n ):\n # if (name in skip_list) or any((kw in name for kw in skip_keywords)) or len(param.shape) == 1:\n no_decay.append(param)\n else:\n has_decay.append(param)\n\n group1 = {\n \"params\": has_decay,\n \"weight_decay\": wd,\n \"lr\": lr,\n \"weight_decay_init\": wd,\n \"weight_decay_base\": wd,\n # \"lr_init\": lr,\n \"lr_base\": lr,\n }\n group2 = {\n \"params\": no_decay,\n \"weight_decay\": 0.0,\n \"lr\": lr,\n \"weight_decay_init\": 0.0,\n \"weight_decay_base\": 0.0,\n \"weight_decay_final\": 0.0,\n # \"lr_init\": lr,\n \"lr_base\": lr,\n }\n return [group1, group2], [lr, lr]\n\n\ndef get_num_layer_for_swin(var_name, num_max_layer, layers_per_stage):\n if var_name in (\"cls_token\", \"mask_token\", \"pos_embed\", \"absolute_pos_embed\"):\n return 0\n elif var_name.startswith(\"patch_embed\"):\n return 0\n elif var_name.startswith(\"layers\"):\n if var_name.split(\".\")[2] == \"blocks\":\n stage_id = int(var_name.split(\".\")[1])\n layer_id = int(var_name.split(\".\")[3]) + sum(layers_per_stage[:stage_id])\n return layer_id + 1\n elif var_name.split(\".\")[2] == \"downsample\":\n stage_id = int(var_name.split(\".\")[1])\n layer_id = sum(layers_per_stage[: stage_id + 1])\n return layer_id\n else:\n return num_max_layer - 1\n\n\ndef get_params_layerdecayswin(module, lr, wd, ld):\n skip_list = {}\n skip_keywords = {}\n if hasattr(module, \"no_weight_decay\"):\n skip_list = module.no_weight_decay()\n if hasattr(module, \"no_weight_decay_keywords\"):\n skip_keywords = module.no_weight_decay_keywords()\n layers_per_stage = module.depths\n num_layers = sum(layers_per_stage) + 1\n lrs = []\n params = []\n for name, param in module.named_parameters():\n if not param.requires_grad:\n print(f\"{name} frozen\")\n continue # frozen weights\n layer_id = get_num_layer_for_swin(name, num_layers, layers_per_stage)\n lr_cur = lr * ld ** (num_layers - layer_id - 1)\n # if (name in skip_list) or any((kw in name for kw in skip_keywords)) or len(param.shape) == 1 or name.endswith(\".bias\"):\n if (name in skip_list) or any((kw in name for kw in skip_keywords)):\n wd_cur = 0.0\n else:\n wd_cur = wd\n params.append({\"params\": param, \"weight_decay\": wd_cur, \"lr\": lr_cur})\n lrs.append(lr_cur)\n return params, lrs\n\n\ndef log(t, eps: float = 1e-5):\n return torch.log(t.clamp(min=eps))\n\n\ndef l2norm(t):\n return F.normalize(t, dim=-1)\n\n\ndef exists(val):\n return val is not None\n\n\ndef identity(t, *args, **kwargs):\n return t\n\n\ndef divisible_by(numer, denom):\n return (numer % denom) == 0\n\n\ndef first(arr, d=None):\n if len(arr) == 0:\n return d\n return arr[0]\n\n\ndef default(val, d):\n if exists(val):\n return val\n return d() if callable(d) else d\n\n\ndef maybe(fn):\n @wraps(fn)\n def inner(x):\n if not exists(x):\n return x\n return fn(x)\n\n return inner\n\n\ndef once(fn):\n called = False\n\n @wraps(fn)\n def inner(x):\n nonlocal called\n if called:\n return\n called = True\n return fn(x)\n\n return inner\n\n\ndef _many(fn):\n @wraps(fn)\n def inner(tensors, pattern, **kwargs):\n return (fn(tensor, pattern, **kwargs) for tensor in tensors)\n\n return inner\n\n\nrearrange_many = _many(rearrange)\nrepeat_many = _many(repeat)\nreduce_many = _many(reduce)\n\n\ndef load_pretrained(state_dict, checkpoint):\n checkpoint_model = checkpoint[\"model\"]\n if any([True if \"encoder.\" in k else False for k in checkpoint_model.keys()]):\n checkpoint_model = {\n k.replace(\"encoder.\", \"\"): v\n for k, v in checkpoint_model.items()\n if k.startswith(\"encoder.\")\n }\n print(\"Detect pre-trained model, remove [encoder.] prefix.\")\n else:\n print(\"Detect non-pre-trained model, pass without doing anything.\")\n print(f\">>>>>>>>>> Remapping pre-trained keys for SWIN ..........\")\n checkpoint = load_checkpoint_swin(state_dict, checkpoint_model)\n\n\ndef load_checkpoint_swin(model, checkpoint_model):\n state_dict = model.state_dict()\n # Geometric interpolation when pre-trained patch size mismatch with fine-tuned patch size\n all_keys = list(checkpoint_model.keys())\n for key in all_keys:\n if \"relative_position_bias_table\" in key:\n relative_position_bias_table_pretrained = checkpoint_model[key]\n relative_position_bias_table_current = state_dict[key]\n L1, nH1 = relative_position_bias_table_pretrained.size()\n L2, nH2 = relative_position_bias_table_current.size()\n if nH1 != nH2:\n print(f\"Error in loading {key}, passing......\")\n else:\n if L1 != L2:\n print(f\"{key}: Interpolate relative_position_bias_table using geo.\")\n src_size = int(L1**0.5)\n dst_size = int(L2**0.5)\n\n def geometric_progression(a, r, n):\n return a * (1.0 - r**n) / (1.0 - r)\n\n left, right = 1.01, 1.5\n while right - left > 1e-6:\n q = (left + right) / 2.0\n gp = geometric_progression(1, q, src_size // 2)\n if gp > dst_size // 2:\n right = q\n else:\n left = q\n\n # if q > 1.090307:\n # q = 1.090307\n\n dis = []\n cur = 1\n for i in range(src_size // 2):\n dis.append(cur)\n cur += q ** (i + 1)\n\n r_ids = [-_ for _ in reversed(dis)]\n\n x = r_ids + [0] + dis\n y = r_ids + [0] + dis\n\n t = dst_size // 2.0\n dx = np.arange(-t, t + 0.1, 1.0)\n dy = np.arange(-t, t + 0.1, 1.0)\n\n print(\"Original positions = %s\" % str(x))\n print(\"Target positions = %s\" % str(dx))\n\n all_rel_pos_bias = []\n\n for i in range(nH1):\n z = (\n relative_position_bias_table_pretrained[:, i]\n .view(src_size, src_size)\n .float()\n .numpy()\n )\n f_cubic = interpolate.interp2d(x, y, z, kind=\"cubic\")\n all_rel_pos_bias.append(\n torch.Tensor(f_cubic(dx, dy))\n .contiguous()\n .view(-1, 1)\n .to(relative_position_bias_table_pretrained.device)\n )\n\n new_rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)\n checkpoint_model[key] = new_rel_pos_bias\n\n # delete relative_position_index since we always re-init it\n relative_position_index_keys = [\n k for k in checkpoint_model.keys() if \"relative_position_index\" in k\n ]\n for k in relative_position_index_keys:\n del checkpoint_model[k]\n\n # delete relative_coords_table since we always re-init it\n relative_coords_table_keys = [\n k for k in checkpoint_model.keys() if \"relative_coords_table\" in k\n ]\n for k in relative_coords_table_keys:\n del checkpoint_model[k]\n\n # # re-map keys due to name change\n rpe_mlp_keys = [k for k in checkpoint_model.keys() if \"cpb_mlp\" in k]\n for k in rpe_mlp_keys:\n checkpoint_model[k.replace(\"cpb_mlp\", \"rpe_mlp\")] = checkpoint_model.pop(k)\n\n # delete attn_mask since we always re-init it\n attn_mask_keys = [k for k in checkpoint_model.keys() if \"attn_mask\" in k]\n for k in attn_mask_keys:\n del checkpoint_model[k]\n\n encoder_keys = [k for k in checkpoint_model.keys() if k.startswith(\"encoder.\")]\n for k in encoder_keys:\n checkpoint_model[k.replace(\"encoder.\", \"\")] = checkpoint_model.pop(k)\n\n return checkpoint_model\n\n\ndef add_padding_metas(out, image_metas):\n device = out.device\n # left, right, top, bottom\n paddings = [img_meta.get(\"paddings\", [0] * 4) for img_meta in image_metas]\n paddings = torch.stack(paddings).to(device)\n outs = [F.pad(o, padding, value=0.0) for padding, o in zip(paddings, out)]\n return torch.stack(outs)\n\n\n# left, right, top, bottom\ndef remove_padding(out, paddings):\n H, W = out.shape[-2:]\n outs = [\n o[..., padding[2] : H - padding[3], padding[0] : W - padding[1]]\n for padding, o in zip(paddings, out)\n ]\n return torch.stack(outs)\n\n\ndef remove_padding_metas(out, image_metas):\n B, C, H, W = out.shape\n device = out.device\n # left, right, top, bottom\n paddings = [\n torch.tensor(img_meta.get(\"paddings\", [0] * 4)) for img_meta in image_metas\n ]\n return remove_padding(out, paddings)\n\n\ndef ssi_helper(tensor1, tensor2):\n stability_mat = 1e-4 * torch.eye(2, device=tensor1.device)\n tensor2_one = torch.stack([tensor2, torch.ones_like(tensor2)], dim=1)\n scale_shift = torch.inverse(tensor2_one.T @ tensor2_one + stability_mat) @ (\n tensor2_one.T @ tensor1.unsqueeze(1)\n )\n scale, shift = scale_shift.squeeze().chunk(2, dim=0)\n return scale, shift\n\n\ndef calculate_mean_values(names, values):\n # Create a defaultdict to store sum and count for each name\n name_values = {name: {} for name in names}\n\n # Iterate through the lists and accumulate values for each name\n for name, value in zip(names, values):\n name_values[name][\"sum\"] = name_values[name].get(\"sum\", 0.0) + value\n name_values[name][\"count\"] = name_values[name].get(\"count\", 0.0) + 1\n\n # Calculate mean values and create the output dictionary\n output_dict = {\n name: name_values[name][\"sum\"] / name_values[name][\"count\"]\n for name in name_values\n }\n\n return output_dict\n\n\ndef remove_leading_dim(infos):\n if isinstance(infos, dict):\n return {k: remove_leading_dim(v) for k, v in infos.items()}\n elif isinstance(infos, torch.Tensor):\n return infos.squeeze(0)\n else:\n return infos\n\n\ndef recursive_index(infos, index):\n if isinstance(infos, dict):\n return {k: recursive_index(v, index) for k, v in infos.items()}\n elif isinstance(infos, torch.Tensor):\n return infos[index]\n else:\n return infos\n\n\ndef to_cpu(infos):\n if isinstance(infos, dict):\n return {k: to_cpu(v) for k, v in infos.items()}\n elif isinstance(infos, torch.Tensor):\n return infos.detach()\n else:\n return infos\n\n\ndef recursive_to(infos, device, non_blocking, cls):\n if isinstance(infos, dict):\n return {k: recursive_to(v, device, non_blocking, cls) for k, v in infos.items()}\n elif isinstance(infos, list):\n return [recursive_to(v, device, non_blocking, cls) for v in infos]\n elif isinstance(infos, cls):\n return infos.to(device, non_blocking=non_blocking)\n else:\n return infos\n\n\ndef masked_mean(\n data: torch.Tensor,\n mask: torch.Tensor | None = None,\n dim: list[int] | None = None,\n keepdim: bool = False,\n) -> torch.Tensor:\n dim = dim if dim is not None else list(range(data.dim()))\n if mask is None:\n return data.mean(dim=dim, keepdim=keepdim)\n mask = mask.float()\n mask_sum = torch.sum(mask, dim=dim, keepdim=True)\n mask_mean = torch.sum(data * mask, dim=dim, keepdim=True) / torch.clamp(\n mask_sum, min=1.0\n )\n return mask_mean.squeeze(dim) if not keepdim else mask_mean\n\n\nclass ProfileMethod:\n def __init__(self, model, func_name, track_statistics=True, verbose=False):\n self.model = model\n self.func_name = func_name\n self.verbose = verbose\n self.track_statistics = track_statistics\n self.timings = []\n\n def __enter__(self):\n # Start timing\n if self.verbose:\n if torch.cuda.is_available():\n torch.cuda.synchronize()\n self.start_time = time()\n return self\n\n def __exit__(self, exc_type, exc_val, exc_tb):\n if self.verbose:\n if torch.cuda.is_available():\n torch.cuda.synchronize()\n\n self.end_time = time()\n\n elapsed_time = self.end_time - self.start_time\n\n self.timings.append(elapsed_time)\n if self.track_statistics and len(self.timings) > 25:\n\n # Compute statistics if tracking\n timings_array = np.array(self.timings)\n mean_time = np.mean(timings_array)\n std_time = np.std(timings_array)\n quantiles = np.percentile(timings_array, [0, 25, 50, 75, 100])\n print(\n f\"{self.model.__class__.__name__}.{self.func_name} took {elapsed_time:.4f} seconds\"\n )\n print(f\"Mean Time: {mean_time:.4f} seconds\")\n print(f\"Std Time: {std_time:.4f} seconds\")\n print(\n f\"Quantiles: Min={quantiles[0]:.4f}, 25%={quantiles[1]:.4f}, Median={quantiles[2]:.4f}, 75%={quantiles[3]:.4f}, Max={quantiles[4]:.4f}\"\n )\n\n else:\n print(\n f\"{self.model.__class__.__name__}.{self.func_name} took {elapsed_time:.4f} seconds\"\n )\n\n\ndef profile_method(track_statistics=True, verbose=False):\n def decorator(func):\n @wraps(func)\n def wrapper(self, *args, **kwargs):\n with ProfileMethod(self, func.__name__, track_statistics, verbose):\n return func(self, *args, **kwargs)\n\n return wrapper\n\n return decorator\n\n\nclass ProfileFunction:\n def __init__(self, func_name, track_statistics=True, verbose=False):\n self.func_name = func_name\n self.verbose = verbose\n self.track_statistics = track_statistics\n self.timings = []\n\n def __enter__(self):\n # Start timing\n if self.verbose:\n if torch.cuda.is_available():\n torch.cuda.synchronize()\n self.start_time = time()\n return self\n\n def __exit__(self, exc_type, exc_val, exc_tb):\n if self.verbose:\n if torch.cuda.is_available():\n torch.cuda.synchronize()\n\n self.end_time = time()\n\n elapsed_time = self.end_time - self.start_time\n\n self.timings.append(elapsed_time)\n if self.track_statistics and len(self.timings) > 25:\n\n # Compute statistics if tracking\n timings_array = np.array(self.timings)\n mean_time = np.mean(timings_array)\n std_time = np.std(timings_array)\n quantiles = np.percentile(timings_array, [0, 25, 50, 75, 100])\n print(f\"{self.func_name} took {elapsed_time:.4f} seconds\")\n print(f\"Mean Time: {mean_time:.4f} seconds\")\n print(f\"Std Time: {std_time:.4f} seconds\")\n print(\n f\"Quantiles: Min={quantiles[0]:.4f}, 25%={quantiles[1]:.4f}, Median={quantiles[2]:.4f}, 75%={quantiles[3]:.4f}, Max={quantiles[4]:.4f}\"\n )\n\n else:\n print(f\"{self.func_name} took {elapsed_time:.4f} seconds\")\n\n\ndef profile_function(track_statistics=True, verbose=False):\n def decorator(func):\n @wraps(func)\n def wrapper(self, *args, **kwargs):\n with ProfileFunction(func.__name__, track_statistics, verbose):\n return func(self, *args, **kwargs)\n\n return wrapper\n\n return decorator\n\n\ndef squeeze_list(nested_list, dim, current_dim=0):\n # If the current dimension is in the list of indices to squeeze\n if isinstance(nested_list, list) and len(nested_list) == 1 and current_dim == dim:\n return squeeze_list(nested_list[0], dim, current_dim + 1)\n elif isinstance(nested_list, list):\n return [squeeze_list(item, dim, current_dim + 1) for item in nested_list]\n else:\n return nested_list\n\n\ndef match_gt(tensor1, tensor2, padding1, padding2, mode: str = \"bilinear\"):\n \"\"\"\n Transform each item in tensor1 batch to match tensor2's dimensions and padding.\n\n Args:\n tensor1 (torch.Tensor): The input tensor to transform, with shape (batch_size, channels, height, width).\n tensor2 (torch.Tensor): The target tensor to match, with shape (batch_size, channels, height, width).\n padding1 (tuple): Padding applied to tensor1 (pad_left, pad_right, pad_top, pad_bottom).\n padding2 (tuple): Desired padding to be applied to match tensor2 (pad_left, pad_right, pad_top, pad_bottom).\n\n Returns:\n torch.Tensor: The batch of transformed tensors matching tensor2's size and padding.\n \"\"\"\n # Get batch size\n batch_size = len(tensor1)\n src_dtype = tensor1[0].dtype\n tgt_dtype = tensor2[0].dtype\n\n # List to store transformed tensors\n transformed_tensors = []\n\n for i in range(batch_size):\n item1 = tensor1[i]\n item2 = tensor2[i]\n\n h1, w1 = item1.shape[1], item1.shape[2]\n pad1_l, pad1_r, pad1_t, pad1_b = (\n padding1[i] if padding1 is not None else (0, 0, 0, 0)\n )\n pad2_l, pad2_r, pad2_t, pad2_b = (\n padding2[i] if padding2 is not None else (0, 0, 0, 0)\n )\n item1_unpadded = item1[:, pad1_t : h1 - pad1_b, pad1_l : w1 - pad1_r]\n\n h2, w2 = (\n item2.shape[1] - pad2_t - pad2_b,\n item2.shape[2] - pad2_l - pad2_r,\n )\n\n item1_resized = F.interpolate(\n item1_unpadded.unsqueeze(0).to(tgt_dtype), size=(h2, w2), mode=mode\n )\n item1_padded = F.pad(item1_resized, (pad2_l, pad2_r, pad2_t, pad2_b))\n transformed_tensors.append(item1_padded)\n\n transformed_batch = torch.cat(transformed_tensors)\n return transformed_batch.to(src_dtype)\n\n\ndef match_intrinsics(K1, tensor1, tensor2, padding1, padding2):\n \"\"\"\n Adjust camera intrinsics K1 to match the size and padding transformation applied to tensor1\n so that it corresponds correctly to tensor2.\n\n Args:\n K1 (torch.Tensor): The camera intrinsics matrix for tensor1, shape (batch_size, 3, 3).\n tensor1 (torch.Tensor): The original image tensor, shape (batch_size, C, H1, W1).\n tensor2 (torch.Tensor): The target image tensor, shape (batch_size, C, H2, W2).\n padding1 (list of tuples): List of padding applied to tensor1 (pad_left, pad_right, pad_top, pad_bottom).\n padding2 (list of tuples): Desired padding to be applied to match tensor2 (pad_left, pad_right, pad_top, pad_bottom).\n\n Returns:\n torch.Tensor: The adjusted intrinsics matrix of shape (batch_size, 3, 3).\n \"\"\"\n batch_size = K1.shape[0]\n K1_new = K1.clone()\n\n for i in range(batch_size):\n h1, w1 = tensor1.shape[2], tensor1.shape[3]\n h2, w2 = tensor2.shape[2], tensor2.shape[3]\n\n # Remove original padding\n pad1_l, pad1_r, pad1_t, pad1_b = (\n padding1[i] if padding1 is not None else (0, 0, 0, 0)\n )\n w1_unpadded, h1_unpadded = w1 - (pad1_l + pad1_r), h1 - (pad1_t + pad1_b)\n\n # Compute new image size after removing original padding\n pad2_l, pad2_r, pad2_t, pad2_b = (\n padding2[i] if padding2 is not None else (0, 0, 0, 0)\n )\n w2_unpadded, h2_unpadded = w2 - (pad2_l + pad2_r), h2 - (pad2_t + pad2_b)\n\n # Compute scaling factors\n scale_x = w2_unpadded / w1_unpadded\n scale_y = h2_unpadded / h1_unpadded\n\n # Update focal length (fx, fy) and principal point (cx, cy)\n K1_new[i, 0, 0] *= scale_x # fx\n K1_new[i, 1, 1] *= scale_y # fy\n\n K1_new[i, 0, 2] = (K1[i, 0, 2] - pad1_l) * scale_x + pad2_l # cx\n K1_new[i, 1, 2] = (K1[i, 1, 2] - pad1_t) * scale_y + pad2_t # cy\n\n return K1_new\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/positional_embedding.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nfrom math import pi\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nfrom einops import rearrange, repeat\n\n\nclass PositionEmbeddingSine(nn.Module):\n def __init__(\n self, num_pos_feats=64, temperature=10000, normalize=False, scale=None\n ):\n super().__init__()\n self.num_pos_feats = num_pos_feats\n self.temperature = temperature\n self.normalize = normalize\n if scale is not None and normalize is False:\n raise ValueError(\"normalize should be True if scale is passed\")\n if scale is None:\n scale = 2 * pi\n self.scale = scale\n\n def forward(\n self, x: torch.Tensor, mask: Optional[torch.Tensor] = None\n ) -> torch.Tensor:\n if mask is None:\n mask = torch.zeros(\n (x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool\n )\n not_mask = ~mask\n y_embed = not_mask.cumsum(1, dtype=torch.float32)\n x_embed = not_mask.cumsum(2, dtype=torch.float32)\n if self.normalize:\n eps = 1e-6\n y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale\n x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale\n\n dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)\n dim_t = self.temperature ** (\n 2 * torch.div(dim_t, 2, rounding_mode=\"floor\") / self.num_pos_feats\n )\n\n pos_x = x_embed[:, :, :, None] / dim_t\n pos_y = y_embed[:, :, :, None] / dim_t\n pos_x = torch.stack(\n (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4\n ).flatten(3)\n pos_y = torch.stack(\n (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4\n ).flatten(3)\n pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)\n return pos\n\n def __repr__(self, _repr_indent=4):\n head = \"Positional encoding \" + self.__class__.__name__\n body = [\n \"num_pos_feats: {}\".format(self.num_pos_feats),\n \"temperature: {}\".format(self.temperature),\n \"normalize: {}\".format(self.normalize),\n \"scale: {}\".format(self.scale),\n ]\n # _repr_indent = 4\n lines = [head] + [\" \" * _repr_indent + line for line in body]\n return \"\\n\".join(lines)\n\n\nclass LearnedSinusoidalPosEmb(nn.Module):\n def __init__(self, dim):\n super().__init__()\n assert (dim % 2) == 0\n half_dim = dim // 2\n self.weights = nn.Parameter(torch.randn(half_dim))\n\n def forward(self, x):\n x = rearrange(x, \"b -> b 1\")\n freqs = x * rearrange(self.weights, \"d -> 1 d\") * 2 * pi\n fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)\n fouriered = torch.cat((x, fouriered), dim=-1)\n return fouriered\n\n\ndef broadcat(tensors, dim=-1):\n num_tensors = len(tensors)\n shape_lens = set(list(map(lambda t: len(t.shape), tensors)))\n assert len(shape_lens) == 1, \"tensors must all have the same number of dimensions\"\n shape_len = list(shape_lens)[0]\n dim = (dim + shape_len) if dim < 0 else dim\n dims = list(zip(*map(lambda t: list(t.shape), tensors)))\n expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]\n assert all(\n [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]\n ), \"invalid dimensions for broadcastable concatentation\"\n max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))\n expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))\n expanded_dims.insert(dim, (dim, dims[dim]))\n expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))\n tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))\n return torch.cat(tensors, dim=dim)\n\n\ndef rotate_half(x):\n x = rearrange(x, \"... (d r) -> ... d r\", r=2)\n x1, x2 = x.unbind(dim=-1)\n x = torch.stack((-x2, x1), dim=-1)\n return rearrange(x, \"... d r -> ... (d r)\")\n\n\nclass VisionRotaryEmbedding(nn.Module):\n def __init__(\n self,\n dim,\n pt_seq_len,\n ft_seq_len=None,\n custom_freqs=None,\n freqs_for=\"lang\",\n theta=10000,\n max_freq=10,\n num_freqs=1,\n ):\n super().__init__()\n if custom_freqs:\n freqs = custom_freqs\n elif freqs_for == \"lang\":\n freqs = 1.0 / (\n theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)\n )\n elif freqs_for == \"pixel\":\n freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi\n elif freqs_for == \"constant\":\n freqs = torch.ones(num_freqs).float()\n else:\n raise ValueError(f\"unknown modality {freqs_for}\")\n\n if ft_seq_len is None:\n ft_seq_len = pt_seq_len\n t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len\n\n freqs_h = torch.einsum(\"..., f -> ... f\", t, freqs)\n freqs_h = repeat(freqs_h, \"... n -> ... (n r)\", r=2)\n\n freqs_w = torch.einsum(\"..., f -> ... f\", t, freqs)\n freqs_w = repeat(freqs_w, \"... n -> ... (n r)\", r=2)\n\n freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)\n\n self.register_buffer(\"freqs_cos\", freqs.cos())\n self.register_buffer(\"freqs_sin\", freqs.sin())\n\n print(\"======== shape of rope freq\", self.freqs_cos.shape, \"========\")\n\n def forward(self, t, start_index=0):\n rot_dim = self.freqs_cos.shape[-1]\n end_index = start_index + rot_dim\n assert (\n rot_dim <= t.shape[-1]\n ), f\"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}\"\n t_left, t, t_right = (\n t[..., :start_index],\n t[..., start_index:end_index],\n t[..., end_index:],\n )\n t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)\n return torch.cat((t_left, t, t_right), dim=-1)\n\n\nclass VisionRotaryEmbeddingFast(nn.Module):\n def __init__(\n self,\n dim,\n pt_seq_len,\n ft_seq_len=None,\n custom_freqs=None,\n freqs_for=\"lang\",\n theta=10000,\n max_freq=10,\n num_freqs=1,\n ):\n super().__init__()\n if custom_freqs:\n freqs = custom_freqs\n elif freqs_for == \"lang\":\n freqs = 1.0 / (\n theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)\n )\n elif freqs_for == \"pixel\":\n freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi\n elif freqs_for == \"constant\":\n freqs = torch.ones(num_freqs).float()\n else:\n raise ValueError(f\"unknown modality {freqs_for}\")\n\n if ft_seq_len is None:\n ft_seq_len = pt_seq_len\n t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len\n\n freqs = torch.einsum(\"..., f -> ... f\", t, freqs)\n freqs = repeat(freqs, \"... n -> ... (n r)\", r=2)\n freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)\n\n freqs_cos = freqs.cos().view(-1, freqs.shape[-1])\n freqs_sin = freqs.sin().view(-1, freqs.shape[-1])\n\n self.register_buffer(\"freqs_cos\", freqs_cos)\n self.register_buffer(\"freqs_sin\", freqs_sin)\n\n def forward(self, t):\n return t * self.freqs_cos + rotate_half(t) * self.freqs_sin\n\n\nfrom math import log2\n\n\ndef generate_fourier_features(\n x: torch.Tensor,\n dim: int = 512,\n max_freq: int = 64,\n use_cos: bool = False,\n use_log: bool = False,\n cat_orig: bool = False,\n):\n x_orig = x\n device, dtype, input_dim = x.device, x.dtype, x.shape[-1]\n num_bands = dim // (2 * input_dim) if use_cos else dim // input_dim\n\n if use_log:\n scales = 2.0 ** torch.linspace(\n 0.0, log2(max_freq), steps=num_bands, device=device, dtype=dtype\n )\n else:\n scales = torch.linspace(\n 1.0, max_freq / 2, num_bands, device=device, dtype=dtype\n )\n\n x = x.unsqueeze(-1)\n scales = scales[(*((None,) * (len(x.shape) - 1)), Ellipsis)]\n\n x = x * scales * pi\n x = torch.cat(\n (\n [x.sin(), x.cos()]\n if use_cos\n else [\n x.sin(),\n ]\n ),\n dim=-1,\n )\n x = x.flatten(-2)\n if cat_orig:\n return torch.cat((x, x_orig), dim=-1)\n return x\n\n\n# from PIL import Image\n# from unidepth.utils import image_grid, colorize\n# if __name__ == \"__main__\":\n# H, W = 512, 512\n# resolution = 128\n# mesh = torch.meshgrid(torch.linspace(-1, 1, H), torch.linspace(-1, 1, W))\n# mesh = torch.stack(mesh, dim=0).unsqueeze(0)\n# mesh = mesh.view(1, 2, -1).permute(0, 2, 1)\n\n# features = generate_fourier_features(mesh, dim=32, max_freq=resolution, use_log=True)\n# channels = features.shape[-1]\n# print(features.shape)\n\n# features = features[0].view(H, W, channels).permute(2, 0, 1).numpy()\n# Image.fromarray(image_grid([colorize(1+x, 0.0, 2.0, \"viridis\") for x in features], rows=8, cols=4)).save(f\"tmp_{resolution}.png\")\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/sht.py", "content": "\"\"\"Real spherical harmonics in Cartesian form for PyTorch.\n\nThis is an autogenerated file. See\nhttps://github.com/cheind/torch-spherical-harmonics\nfor more information.\n\"\"\"\n\nimport torch\n\n\ndef rsh_cart_0(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 0.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,1) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n ],\n -1,\n )\n\n\ndef rsh_cart_1(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 1.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,4) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n ],\n -1,\n )\n\n\ndef rsh_cart_2(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 2.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,9) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n x2 = x**2\n y2 = y**2\n z2 = z**2\n xy = x * y\n xz = x * z\n yz = y * z\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n 1.09254843059208 * xy,\n -1.09254843059208 * yz,\n 0.94617469575756 * z2 - 0.31539156525252,\n -1.09254843059208 * xz,\n 0.54627421529604 * x2 - 0.54627421529604 * y2,\n ],\n -1,\n )\n\n\ndef rsh_cart_3(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 3.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,16) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n x2 = x**2\n y2 = y**2\n z2 = z**2\n xy = x * y\n xz = x * z\n yz = y * z\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n 1.09254843059208 * xy,\n -1.09254843059208 * yz,\n 0.94617469575756 * z2 - 0.31539156525252,\n -1.09254843059208 * xz,\n 0.54627421529604 * x2 - 0.54627421529604 * y2,\n -0.590043589926644 * y * (3.0 * x2 - y2),\n 2.89061144264055 * xy * z,\n 0.304697199642977 * y * (1.5 - 7.5 * z2),\n 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z,\n 0.304697199642977 * x * (1.5 - 7.5 * z2),\n 1.44530572132028 * z * (x2 - y2),\n -0.590043589926644 * x * (x2 - 3.0 * y2),\n ],\n -1,\n )\n\n\ndef rsh_cart_4(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 4.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,25) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n x2 = x**2\n y2 = y**2\n z2 = z**2\n xy = x * y\n xz = x * z\n yz = y * z\n x4 = x2**2\n y4 = y2**2\n z4 = z2**2\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n 1.09254843059208 * xy,\n -1.09254843059208 * yz,\n 0.94617469575756 * z2 - 0.31539156525252,\n -1.09254843059208 * xz,\n 0.54627421529604 * x2 - 0.54627421529604 * y2,\n -0.590043589926644 * y * (3.0 * x2 - y2),\n 2.89061144264055 * xy * z,\n 0.304697199642977 * y * (1.5 - 7.5 * z2),\n 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z,\n 0.304697199642977 * x * (1.5 - 7.5 * z2),\n 1.44530572132028 * z * (x2 - y2),\n -0.590043589926644 * x * (x2 - 3.0 * y2),\n 2.5033429417967 * xy * (x2 - y2),\n -1.77013076977993 * yz * (3.0 * x2 - y2),\n 0.126156626101008 * xy * (52.5 * z2 - 7.5),\n 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 1.48099765681286\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 0.952069922236839 * z2\n + 0.317356640745613,\n 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5),\n -1.77013076977993 * xz * (x2 - 3.0 * y2),\n -3.75501441269506 * x2 * y2\n + 0.625835735449176 * x4\n + 0.625835735449176 * y4,\n ],\n -1,\n )\n\n\ndef rsh_cart_5(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 5.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,36) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n x2 = x**2\n y2 = y**2\n z2 = z**2\n xy = x * y\n xz = x * z\n yz = y * z\n x4 = x2**2\n y4 = y2**2\n z4 = z2**2\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n 1.09254843059208 * xy,\n -1.09254843059208 * yz,\n 0.94617469575756 * z2 - 0.31539156525252,\n -1.09254843059208 * xz,\n 0.54627421529604 * x2 - 0.54627421529604 * y2,\n -0.590043589926644 * y * (3.0 * x2 - y2),\n 2.89061144264055 * xy * z,\n 0.304697199642977 * y * (1.5 - 7.5 * z2),\n 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z,\n 0.304697199642977 * x * (1.5 - 7.5 * z2),\n 1.44530572132028 * z * (x2 - y2),\n -0.590043589926644 * x * (x2 - 3.0 * y2),\n 2.5033429417967 * xy * (x2 - y2),\n -1.77013076977993 * yz * (3.0 * x2 - y2),\n 0.126156626101008 * xy * (52.5 * z2 - 7.5),\n 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 1.48099765681286\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 0.952069922236839 * z2\n + 0.317356640745613,\n 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5),\n -1.77013076977993 * xz * (x2 - 3.0 * y2),\n -3.75501441269506 * x2 * y2\n + 0.625835735449176 * x4\n + 0.625835735449176 * y4,\n -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 8.30264925952416 * xy * z * (x2 - y2),\n 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2),\n 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.241571547304372\n * y\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n -1.24747010616985 * z * (1.5 * z2 - 0.5)\n + 1.6840846433293\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.498988042467941 * z,\n 0.241571547304372\n * x\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2),\n 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4),\n -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n ],\n -1,\n )\n\n\ndef rsh_cart_6(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 6.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,49) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n x2 = x**2\n y2 = y**2\n z2 = z**2\n xy = x * y\n xz = x * z\n yz = y * z\n x4 = x2**2\n y4 = y2**2\n z4 = z2**2\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n 1.09254843059208 * xy,\n -1.09254843059208 * yz,\n 0.94617469575756 * z2 - 0.31539156525252,\n -1.09254843059208 * xz,\n 0.54627421529604 * x2 - 0.54627421529604 * y2,\n -0.590043589926644 * y * (3.0 * x2 - y2),\n 2.89061144264055 * xy * z,\n 0.304697199642977 * y * (1.5 - 7.5 * z2),\n 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z,\n 0.304697199642977 * x * (1.5 - 7.5 * z2),\n 1.44530572132028 * z * (x2 - y2),\n -0.590043589926644 * x * (x2 - 3.0 * y2),\n 2.5033429417967 * xy * (x2 - y2),\n -1.77013076977993 * yz * (3.0 * x2 - y2),\n 0.126156626101008 * xy * (52.5 * z2 - 7.5),\n 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 1.48099765681286\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 0.952069922236839 * z2\n + 0.317356640745613,\n 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5),\n -1.77013076977993 * xz * (x2 - 3.0 * y2),\n -3.75501441269506 * x2 * y2\n + 0.625835735449176 * x4\n + 0.625835735449176 * y4,\n -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 8.30264925952416 * xy * z * (x2 - y2),\n 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2),\n 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.241571547304372\n * y\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n -1.24747010616985 * z * (1.5 * z2 - 0.5)\n + 1.6840846433293\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.498988042467941 * z,\n 0.241571547304372\n * x\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2),\n 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4),\n -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 4.09910463115149 * x**4 * xy\n - 13.6636821038383 * xy**3\n + 4.09910463115149 * xy * y**4,\n -2.36661916223175 * yz * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 0.00427144889505798 * xy * (x2 - y2) * (5197.5 * z2 - 472.5),\n 0.00584892228263444\n * y\n * (3.0 * x2 - y2)\n * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z),\n 0.0701870673916132\n * xy\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n ),\n 0.221950995245231\n * y\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n ),\n -1.48328138624466\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n + 1.86469659985043\n * z\n * (\n -1.33333333333333 * z * (1.5 * z2 - 0.5)\n + 1.8\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.533333333333333 * z\n )\n + 0.953538034014426 * z2\n - 0.317846011338142,\n 0.221950995245231\n * x\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n ),\n 0.0350935336958066\n * (x2 - y2)\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n ),\n 0.00584892228263444\n * x\n * (x2 - 3.0 * y2)\n * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z),\n 0.0010678622237645 * (5197.5 * z2 - 472.5) * (-6.0 * x2 * y2 + x4 + y4),\n -2.36661916223175 * xz * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 0.683184105191914 * x2**3\n + 10.2477615778787 * x2 * y4\n - 10.2477615778787 * x4 * y2\n - 0.683184105191914 * y2**3,\n ],\n -1,\n )\n\n\ndef rsh_cart_7(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 7.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,64) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n x2 = x**2\n y2 = y**2\n z2 = z**2\n xy = x * y\n xz = x * z\n yz = y * z\n x4 = x2**2\n y4 = y2**2\n z4 = z2**2\n\n return torch.stack(\n [\n xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n 1.09254843059208 * xy,\n -1.09254843059208 * yz,\n 0.94617469575756 * z2 - 0.31539156525252,\n -1.09254843059208 * xz,\n 0.54627421529604 * x2 - 0.54627421529604 * y2,\n -0.590043589926644 * y * (3.0 * x2 - y2),\n 2.89061144264055 * xy * z,\n 0.304697199642977 * y * (1.5 - 7.5 * z2),\n 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z,\n 0.304697199642977 * x * (1.5 - 7.5 * z2),\n 1.44530572132028 * z * (x2 - y2),\n -0.590043589926644 * x * (x2 - 3.0 * y2),\n 2.5033429417967 * xy * (x2 - y2),\n -1.77013076977993 * yz * (3.0 * x2 - y2),\n 0.126156626101008 * xy * (52.5 * z2 - 7.5),\n 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 1.48099765681286\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 0.952069922236839 * z2\n + 0.317356640745613,\n 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5),\n -1.77013076977993 * xz * (x2 - 3.0 * y2),\n -3.75501441269506 * x2 * y2\n + 0.625835735449176 * x4\n + 0.625835735449176 * y4,\n -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 8.30264925952416 * xy * z * (x2 - y2),\n 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2),\n 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.241571547304372\n * y\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n -1.24747010616985 * z * (1.5 * z2 - 0.5)\n + 1.6840846433293\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.498988042467941 * z,\n 0.241571547304372\n * x\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2),\n 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4),\n -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 4.09910463115149 * x**4 * xy\n - 13.6636821038383 * xy**3\n + 4.09910463115149 * xy * y**4,\n -2.36661916223175 * yz * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 0.00427144889505798 * xy * (x2 - y2) * (5197.5 * z2 - 472.5),\n 0.00584892228263444\n * y\n * (3.0 * x2 - y2)\n * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z),\n 0.0701870673916132\n * xy\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n ),\n 0.221950995245231\n * y\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n ),\n -1.48328138624466\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n + 1.86469659985043\n * z\n * (\n -1.33333333333333 * z * (1.5 * z2 - 0.5)\n + 1.8\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.533333333333333 * z\n )\n + 0.953538034014426 * z2\n - 0.317846011338142,\n 0.221950995245231\n * x\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n ),\n 0.0350935336958066\n * (x2 - y2)\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n ),\n 0.00584892228263444\n * x\n * (x2 - 3.0 * y2)\n * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z),\n 0.0010678622237645 * (5197.5 * z2 - 472.5) * (-6.0 * x2 * y2 + x4 + y4),\n -2.36661916223175 * xz * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 0.683184105191914 * x2**3\n + 10.2477615778787 * x2 * y4\n - 10.2477615778787 * x4 * y2\n - 0.683184105191914 * y2**3,\n -0.707162732524596\n * y\n * (7.0 * x2**3 + 21.0 * x2 * y4 - 35.0 * x4 * y2 - y2**3),\n 2.6459606618019 * z * (6.0 * x**4 * xy - 20.0 * xy**3 + 6.0 * xy * y**4),\n 9.98394571852353e-5\n * y\n * (5197.5 - 67567.5 * z2)\n * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 0.00239614697244565\n * xy\n * (x2 - y2)\n * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z),\n 0.00397356022507413\n * y\n * (3.0 * x2 - y2)\n * (\n 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z)\n + 1063.125 * z2\n - 118.125\n ),\n 0.0561946276120613\n * xy\n * (\n -4.8 * z * (52.5 * z2 - 7.5)\n + 2.6\n * z\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n )\n + 48.0 * z\n ),\n 0.206472245902897\n * y\n * (\n -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 2.16666666666667\n * z\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n )\n - 10.9375 * z2\n + 2.1875\n ),\n 1.24862677781952 * z * (1.5 * z2 - 0.5)\n - 1.68564615005635\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 2.02901851395672\n * z\n * (\n -1.45833333333333\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n + 1.83333333333333\n * z\n * (\n -1.33333333333333 * z * (1.5 * z2 - 0.5)\n + 1.8\n * z\n * (\n 1.75\n * z\n * (\n 1.66666666666667 * z * (1.5 * z2 - 0.5)\n - 0.666666666666667 * z\n )\n - 1.125 * z2\n + 0.375\n )\n + 0.533333333333333 * z\n )\n + 0.9375 * z2\n - 0.3125\n )\n - 0.499450711127808 * z,\n 0.206472245902897\n * x\n * (\n -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 2.16666666666667\n * z\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n )\n - 10.9375 * z2\n + 2.1875\n ),\n 0.0280973138060306\n * (x2 - y2)\n * (\n -4.8 * z * (52.5 * z2 - 7.5)\n + 2.6\n * z\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n )\n + 48.0 * z\n ),\n 0.00397356022507413\n * x\n * (x2 - 3.0 * y2)\n * (\n 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z)\n + 1063.125 * z2\n - 118.125\n ),\n 0.000599036743111412\n * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z)\n * (-6.0 * x2 * y2 + x4 + y4),\n 9.98394571852353e-5\n * x\n * (5197.5 - 67567.5 * z2)\n * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 2.6459606618019 * z * (x2**3 + 15.0 * x2 * y4 - 15.0 * x4 * y2 - y2**3),\n -0.707162732524596\n * x\n * (x2**3 + 35.0 * x2 * y4 - 21.0 * x4 * y2 - 7.0 * y2**3),\n ],\n -1,\n )\n\n\n# @torch.jit.script\ndef rsh_cart_8(xyz: torch.Tensor):\n \"\"\"Computes all real spherical harmonics up to degree 8.\n\n This is an autogenerated method. See\n https://github.com/cheind/torch-spherical-harmonics\n for more information.\n\n Params:\n xyz: (N,...,3) tensor of points on the unit sphere\n\n Returns:\n rsh: (N,...,81) real spherical harmonics\n projections of input. Ynm is found at index\n `n*(n+1) + m`, with `0 <= n <= degree` and\n `-n <= m <= n`.\n \"\"\"\n x = xyz[..., 0]\n y = xyz[..., 1]\n z = xyz[..., 2]\n\n x2 = x**2\n y2 = y**2\n z2 = z**2\n xy = x * y\n xz = x * z\n yz = y * z\n x4 = x2**2\n y4 = y2**2\n # z4 = z2**2\n return torch.stack(\n [\n 0.282094791773878 * torch.ones(1, device=xyz.device).expand(xyz.shape[:-1]),\n -0.48860251190292 * y,\n 0.48860251190292 * z,\n -0.48860251190292 * x,\n 1.09254843059208 * xy,\n -1.09254843059208 * yz,\n 0.94617469575756 * z2 - 0.31539156525252,\n -1.09254843059208 * xz,\n 0.54627421529604 * x2 - 0.54627421529604 * y2,\n -0.590043589926644 * y * (3.0 * x2 - y2),\n 2.89061144264055 * xy * z,\n 0.304697199642977 * y * (1.5 - 7.5 * z2),\n 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z,\n 0.304697199642977 * x * (1.5 - 7.5 * z2),\n 1.44530572132028 * z * (x2 - y2),\n -0.590043589926644 * x * (x2 - 3.0 * y2),\n 2.5033429417967 * xy * (x2 - y2),\n -1.77013076977993 * yz * (3.0 * x2 - y2),\n 0.126156626101008 * xy * (52.5 * z2 - 7.5),\n 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 1.48099765681286\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 0.952069922236839 * z2\n + 0.317356640745613,\n 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z),\n 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5),\n -1.77013076977993 * xz * (x2 - 3.0 * y2),\n -3.75501441269506 * x2 * y2\n + 0.625835735449176 * x4\n + 0.625835735449176 * y4,\n -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 8.30264925952416 * xy * z * (x2 - y2),\n 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2),\n 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.241571547304372\n * y\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n -1.24747010616985 * z * (1.5 * z2 - 0.5)\n + 1.6840846433293\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.498988042467941 * z,\n 0.241571547304372\n * x\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n ),\n 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z),\n 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2),\n 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4),\n -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 4.09910463115149 * x**4 * xy\n - 13.6636821038383 * xy**3\n + 4.09910463115149 * xy * y**4,\n -2.36661916223175 * yz * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 0.00427144889505798 * xy * (x2 - y2) * (5197.5 * z2 - 472.5),\n 0.00584892228263444\n * y\n * (3.0 * x2 - y2)\n * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z),\n 0.0701870673916132\n * xy\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n ),\n 0.221950995245231\n * y\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n ),\n -1.48328138624466\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n + 1.86469659985043\n * z\n * (\n -1.33333333333333 * z * (1.5 * z2 - 0.5)\n + 1.8\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.533333333333333 * z\n )\n + 0.953538034014426 * z2\n - 0.317846011338142,\n 0.221950995245231\n * x\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n ),\n 0.0350935336958066\n * (x2 - y2)\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n ),\n 0.00584892228263444\n * x\n * (x2 - 3.0 * y2)\n * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z),\n 0.0010678622237645 * (5197.5 * z2 - 472.5) * (-6.0 * x2 * y2 + x4 + y4),\n -2.36661916223175 * xz * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 0.683184105191914 * x2**3\n + 10.2477615778787 * x2 * y4\n - 10.2477615778787 * x4 * y2\n - 0.683184105191914 * y2**3,\n -0.707162732524596\n * y\n * (7.0 * x2**3 + 21.0 * x2 * y4 - 35.0 * x4 * y2 - y2**3),\n 2.6459606618019 * z * (6.0 * x**4 * xy - 20.0 * xy**3 + 6.0 * xy * y**4),\n 9.98394571852353e-5\n * y\n * (5197.5 - 67567.5 * z2)\n * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 0.00239614697244565\n * xy\n * (x2 - y2)\n * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z),\n 0.00397356022507413\n * y\n * (3.0 * x2 - y2)\n * (\n 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z)\n + 1063.125 * z2\n - 118.125\n ),\n 0.0561946276120613\n * xy\n * (\n -4.8 * z * (52.5 * z2 - 7.5)\n + 2.6\n * z\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n )\n + 48.0 * z\n ),\n 0.206472245902897\n * y\n * (\n -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 2.16666666666667\n * z\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n )\n - 10.9375 * z2\n + 2.1875\n ),\n 1.24862677781952 * z * (1.5 * z2 - 0.5)\n - 1.68564615005635\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 2.02901851395672\n * z\n * (\n -1.45833333333333\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n + 1.83333333333333\n * z\n * (\n -1.33333333333333 * z * (1.5 * z2 - 0.5)\n + 1.8\n * z\n * (\n 1.75\n * z\n * (\n 1.66666666666667 * z * (1.5 * z2 - 0.5)\n - 0.666666666666667 * z\n )\n - 1.125 * z2\n + 0.375\n )\n + 0.533333333333333 * z\n )\n + 0.9375 * z2\n - 0.3125\n )\n - 0.499450711127808 * z,\n 0.206472245902897\n * x\n * (\n -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 2.16666666666667\n * z\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n )\n - 10.9375 * z2\n + 2.1875\n ),\n 0.0280973138060306\n * (x2 - y2)\n * (\n -4.8 * z * (52.5 * z2 - 7.5)\n + 2.6\n * z\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n )\n + 48.0 * z\n ),\n 0.00397356022507413\n * x\n * (x2 - 3.0 * y2)\n * (\n 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z)\n + 1063.125 * z2\n - 118.125\n ),\n 0.000599036743111412\n * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z)\n * (-6.0 * x2 * y2 + x4 + y4),\n 9.98394571852353e-5\n * x\n * (5197.5 - 67567.5 * z2)\n * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 2.6459606618019 * z * (x2**3 + 15.0 * x2 * y4 - 15.0 * x4 * y2 - y2**3),\n -0.707162732524596\n * x\n * (x2**3 + 35.0 * x2 * y4 - 21.0 * x4 * y2 - 7.0 * y2**3),\n 5.83141328139864 * xy * (x2**3 + 7.0 * x2 * y4 - 7.0 * x4 * y2 - y2**3),\n -2.91570664069932\n * yz\n * (7.0 * x2**3 + 21.0 * x2 * y4 - 35.0 * x4 * y2 - y2**3),\n 7.87853281621404e-6\n * (1013512.5 * z2 - 67567.5)\n * (6.0 * x**4 * xy - 20.0 * xy**3 + 6.0 * xy * y**4),\n 5.10587282657803e-5\n * y\n * (5.0 * z * (5197.5 - 67567.5 * z2) + 41580.0 * z)\n * (-10.0 * x2 * y2 + 5.0 * x4 + y4),\n 0.00147275890257803\n * xy\n * (x2 - y2)\n * (\n 3.75 * z * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z)\n - 14293.125 * z2\n + 1299.375\n ),\n 0.0028519853513317\n * y\n * (3.0 * x2 - y2)\n * (\n -7.33333333333333 * z * (52.5 - 472.5 * z2)\n + 3.0\n * z\n * (\n 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z)\n + 1063.125 * z2\n - 118.125\n )\n - 560.0 * z\n ),\n 0.0463392770473559\n * xy\n * (\n -4.125 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n + 2.5\n * z\n * (\n -4.8 * z * (52.5 * z2 - 7.5)\n + 2.6\n * z\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n )\n + 48.0 * z\n )\n + 137.8125 * z2\n - 19.6875\n ),\n 0.193851103820053\n * y\n * (\n 3.2 * z * (1.5 - 7.5 * z2)\n - 2.51428571428571\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n + 2.14285714285714\n * z\n * (\n -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 2.16666666666667\n * z\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25\n * z\n * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n )\n - 10.9375 * z2\n + 2.1875\n )\n + 5.48571428571429 * z\n ),\n 1.48417251362228\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.86581687426801\n * z\n * (\n -1.33333333333333 * z * (1.5 * z2 - 0.5)\n + 1.8\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 0.533333333333333 * z\n )\n + 2.1808249179756\n * z\n * (\n 1.14285714285714 * z * (1.5 * z2 - 0.5)\n - 1.54285714285714\n * z\n * (\n 1.75\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n - 1.125 * z2\n + 0.375\n )\n + 1.85714285714286\n * z\n * (\n -1.45833333333333\n * z\n * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z)\n + 1.83333333333333\n * z\n * (\n -1.33333333333333 * z * (1.5 * z2 - 0.5)\n + 1.8\n * z\n * (\n 1.75\n * z\n * (\n 1.66666666666667 * z * (1.5 * z2 - 0.5)\n - 0.666666666666667 * z\n )\n - 1.125 * z2\n + 0.375\n )\n + 0.533333333333333 * z\n )\n + 0.9375 * z2\n - 0.3125\n )\n - 0.457142857142857 * z\n )\n - 0.954110901614325 * z2\n + 0.318036967204775,\n 0.193851103820053\n * x\n * (\n 3.2 * z * (1.5 - 7.5 * z2)\n - 2.51428571428571\n * z\n * (\n 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n + 2.14285714285714\n * z\n * (\n -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 2.16666666666667\n * z\n * (\n -2.8 * z * (1.5 - 7.5 * z2)\n + 2.2\n * z\n * (\n 2.25\n * z\n * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z)\n + 9.375 * z2\n - 1.875\n )\n - 4.8 * z\n )\n - 10.9375 * z2\n + 2.1875\n )\n + 5.48571428571429 * z\n ),\n 0.0231696385236779\n * (x2 - y2)\n * (\n -4.125 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n + 2.5\n * z\n * (\n -4.8 * z * (52.5 * z2 - 7.5)\n + 2.6\n * z\n * (\n 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z)\n - 91.875 * z2\n + 13.125\n )\n + 48.0 * z\n )\n + 137.8125 * z2\n - 19.6875\n ),\n 0.0028519853513317\n * x\n * (x2 - 3.0 * y2)\n * (\n -7.33333333333333 * z * (52.5 - 472.5 * z2)\n + 3.0\n * z\n * (\n 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z)\n + 1063.125 * z2\n - 118.125\n )\n - 560.0 * z\n ),\n 0.000368189725644507\n * (-6.0 * x2 * y2 + x4 + y4)\n * (\n 3.75 * z * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z)\n - 14293.125 * z2\n + 1299.375\n ),\n 5.10587282657803e-5\n * x\n * (5.0 * z * (5197.5 - 67567.5 * z2) + 41580.0 * z)\n * (-10.0 * x2 * y2 + x4 + 5.0 * y4),\n 7.87853281621404e-6\n * (1013512.5 * z2 - 67567.5)\n * (x2**3 + 15.0 * x2 * y4 - 15.0 * x4 * y2 - y2**3),\n -2.91570664069932\n * xz\n * (x2**3 + 35.0 * x2 * y4 - 21.0 * x4 * y2 - 7.0 * y2**3),\n -20.4099464848952 * x2**3 * y2\n - 20.4099464848952 * x2 * y2**3\n + 0.72892666017483 * x4**2\n + 51.0248662122381 * x4 * y4\n + 0.72892666017483 * y4**2,\n ],\n -1,\n )\n\n\n__all__ = [\n \"rsh_cart_0\",\n \"rsh_cart_1\",\n \"rsh_cart_2\",\n \"rsh_cart_3\",\n \"rsh_cart_4\",\n \"rsh_cart_5\",\n \"rsh_cart_6\",\n \"rsh_cart_7\",\n \"rsh_cart_8\",\n]\n\n\nfrom typing import Optional\n\nimport torch\n\n\nclass SphHarm(torch.nn.Module):\n def __init__(self, m, n, dtype=torch.float32) -> None:\n super().__init__()\n self.dtype = dtype\n m = torch.tensor(list(range(-m + 1, m)))\n n = torch.tensor(list(range(n)))\n self.is_normalized = False\n vals = torch.cartesian_prod(m, n).T\n vals = vals[:, vals[0] <= vals[1]]\n m, n = vals.unbind(0)\n\n self.register_buffer(\"m\", tensor=m)\n self.register_buffer(\"n\", tensor=n)\n self.register_buffer(\"l_max\", tensor=torch.max(self.n))\n\n f_a, f_b, initial_value, d0_mask_3d, d1_mask_3d = self._init_legendre()\n self.register_buffer(\"f_a\", tensor=f_a)\n self.register_buffer(\"f_b\", tensor=f_b)\n self.register_buffer(\"d0_mask_3d\", tensor=d0_mask_3d)\n self.register_buffer(\"d1_mask_3d\", tensor=d1_mask_3d)\n self.register_buffer(\"initial_value\", tensor=initial_value)\n\n @property\n def device(self):\n return next(self.buffers()).device\n\n def forward(self, points: torch.Tensor) -> torch.Tensor:\n \"\"\"Computes the spherical harmonics.\"\"\"\n # Y_l^m = (-1) ^ m c_l^m P_l^m(cos(theta)) exp(i m phi)\n B, N, D = points.shape\n dtype = points.dtype\n theta, phi = points.view(-1, D).to(self.dtype).unbind(-1)\n cos_colatitude = torch.cos(phi)\n legendre = self._gen_associated_legendre(cos_colatitude)\n vals = torch.stack([self.m.abs(), self.n], dim=0)\n vals = torch.cat(\n [\n vals.repeat(1, theta.shape[0]),\n torch.arange(theta.shape[0], device=theta.device)\n .unsqueeze(0)\n .repeat_interleave(vals.shape[1], dim=1),\n ],\n dim=0,\n )\n legendre_vals = legendre[vals[0], vals[1], vals[2]]\n legendre_vals = legendre_vals.reshape(-1, theta.shape[0])\n angle = torch.outer(self.m.abs(), theta)\n vandermonde = torch.complex(torch.cos(angle), torch.sin(angle))\n harmonics = torch.complex(\n legendre_vals * torch.real(vandermonde),\n legendre_vals * torch.imag(vandermonde),\n )\n\n # Negative order.\n m = self.m.unsqueeze(-1)\n harmonics = torch.where(\n m < 0, (-1.0) ** m.abs() * torch.conj(harmonics), harmonics\n )\n harmonics = harmonics.permute(1, 0).reshape(B, N, -1).to(dtype)\n return harmonics\n\n def _gen_recurrence_mask(self) -> tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Generates mask for recurrence relation on the remaining entries.\n\n The remaining entries are with respect to the diagonal and offdiagonal\n entries.\n\n Args:\n l_max: see `gen_normalized_legendre`.\n Returns:\n torch.Tensors representing the mask used by the recurrence relations.\n \"\"\"\n\n # Computes all coefficients.\n m_mat, l_mat = torch.meshgrid(\n torch.arange(0, self.l_max + 1, device=self.device, dtype=self.dtype),\n torch.arange(0, self.l_max + 1, device=self.device, dtype=self.dtype),\n indexing=\"ij\",\n )\n if self.is_normalized:\n c0 = l_mat * l_mat\n c1 = m_mat * m_mat\n c2 = 2.0 * l_mat\n c3 = (l_mat - 1.0) * (l_mat - 1.0)\n d0 = torch.sqrt((4.0 * c0 - 1.0) / (c0 - c1))\n d1 = torch.sqrt(((c2 + 1.0) * (c3 - c1)) / ((c2 - 3.0) * (c0 - c1)))\n else:\n d0 = (2.0 * l_mat - 1.0) / (l_mat - m_mat)\n d1 = (l_mat + m_mat - 1.0) / (l_mat - m_mat)\n\n d0_mask_indices = torch.triu_indices(self.l_max + 1, 1)\n d1_mask_indices = torch.triu_indices(self.l_max + 1, 2)\n\n d_zeros = torch.zeros(\n (self.l_max + 1, self.l_max + 1), dtype=self.dtype, device=self.device\n )\n d_zeros[d0_mask_indices] = d0[d0_mask_indices]\n d0_mask = d_zeros\n\n d_zeros = torch.zeros(\n (self.l_max + 1, self.l_max + 1), dtype=self.dtype, device=self.device\n )\n d_zeros[d1_mask_indices] = d1[d1_mask_indices]\n d1_mask = d_zeros\n\n # Creates a 3D mask that contains 1s on the diagonal plane and 0s elsewhere.\n i = torch.arange(self.l_max + 1, device=self.device)[:, None, None]\n j = torch.arange(self.l_max + 1, device=self.device)[None, :, None]\n k = torch.arange(self.l_max + 1, device=self.device)[None, None, :]\n mask = (i + j - k == 0).to(self.dtype)\n d0_mask_3d = torch.einsum(\"jk,ijk->ijk\", d0_mask, mask)\n d1_mask_3d = torch.einsum(\"jk,ijk->ijk\", d1_mask, mask)\n return (d0_mask_3d, d1_mask_3d)\n\n def _recursive(self, i: int, p_val: torch.Tensor, x: torch.Tensor) -> torch.Tensor:\n coeff_0 = self.d0_mask_3d[i]\n coeff_1 = self.d1_mask_3d[i]\n h = torch.einsum(\n \"ij,ijk->ijk\",\n coeff_0,\n torch.einsum(\"ijk,k->ijk\", torch.roll(p_val, shifts=1, dims=1), x),\n ) - torch.einsum(\"ij,ijk->ijk\", coeff_1, torch.roll(p_val, shifts=2, dims=1))\n p_val = p_val + h\n return p_val\n\n def _init_legendre(self):\n a_idx = torch.arange(1, self.l_max + 1, dtype=self.dtype, device=self.device)\n b_idx = torch.arange(self.l_max, dtype=self.dtype, device=self.device)\n if self.is_normalized:\n # The initial value p(0,0).\n initial_value: torch.Tensor = torch.tensor(\n 0.5 / (torch.pi**0.5), device=self.device\n )\n f_a = torch.cumprod(-1 * torch.sqrt(1.0 + 0.5 / a_idx), dim=0)\n f_b = torch.sqrt(2.0 * b_idx + 3.0)\n else:\n # The initial value p(0,0).\n initial_value = torch.tensor(1.0, device=self.device)\n f_a = torch.cumprod(1.0 - 2.0 * a_idx, dim=0)\n f_b = 2.0 * b_idx + 1.0\n\n d0_mask_3d, d1_mask_3d = self._gen_recurrence_mask()\n return f_a, f_b, initial_value, d0_mask_3d, d1_mask_3d\n\n def _gen_associated_legendre(self, x: torch.Tensor) -> torch.Tensor:\n r\"\"\"Computes associated Legendre functions (ALFs) of the first kind.\n\n The ALFs of the first kind are used in spherical harmonics. The spherical\n harmonic of degree `l` and order `m` can be written as\n `Y_l^m(θ, φ) = N_l^m * P_l^m(cos(θ)) * exp(i m φ)`, where `N_l^m` is the\n normalization factor and θ and φ are the colatitude and longitude,\n repectively. `N_l^m` is chosen in the way that the spherical harmonics form\n a set of orthonormal basis function of L^2(S^2). For the computational\n efficiency of spherical harmonics transform, the normalization factor is\n used in the computation of the ALFs. In addition, normalizing `P_l^m`\n avoids overflow/underflow and achieves better numerical stability. Three\n recurrence relations are used in the computation.\n\n Args:\n l_max: The maximum degree of the associated Legendre function. Both the\n degrees and orders are `[0, 1, 2, ..., l_max]`.\n x: A vector of type `float32`, `float64` containing the sampled points in\n spherical coordinates, at which the ALFs are computed; `x` is essentially\n `cos(θ)`. For the numerical integration used by the spherical harmonics\n transforms, `x` contains the quadrature points in the interval of\n `[-1, 1]`. There are several approaches to provide the quadrature points:\n Gauss-Legendre method (`scipy.special.roots_legendre`), Gauss-Chebyshev\n method (`scipy.special.roots_chebyu`), and Driscoll & Healy\n method (Driscoll, James R., and Dennis M. Healy. \"Computing Fourier\n transforms and convolutions on the 2-sphere.\" Advances in applied\n mathematics 15, no. 2 (1994): 202-250.). The Gauss-Legendre quadrature\n points are nearly equal-spaced along θ and provide exact discrete\n orthogonality, (P^m)^T W P_m = I, where `T` represents the transpose\n operation, `W` is a diagonal matrix containing the quadrature weights,\n and `I` is the identity matrix. The Gauss-Chebyshev points are equally\n spaced, which only provide approximate discrete orthogonality. The\n Driscoll & Healy qudarture points are equally spaced and provide the\n exact discrete orthogonality. The number of sampling points is required to\n be twice as the number of frequency points (modes) in the Driscoll & Healy\n approach, which enables FFT and achieves a fast spherical harmonics\n transform.\n is_normalized: True if the associated Legendre functions are normalized.\n With normalization, `N_l^m` is applied such that the spherical harmonics\n form a set of orthonormal basis functions of L^2(S^2).\n\n Returns:\n The 3D array of shape `(l_max + 1, l_max + 1, len(x))` containing the values\n of the ALFs at `x`; the dimensions in the sequence of order, degree, and\n evalution points.\n \"\"\"\n p = torch.zeros(\n (self.l_max + 1, self.l_max + 1, x.shape[0]), dtype=x.dtype, device=x.device\n )\n p[0, 0] = self.initial_value\n\n # Compute the diagonal entries p(l,l) with recurrence.\n y = torch.cumprod(\n torch.broadcast_to(torch.sqrt(1.0 - x * x), (self.l_max, x.shape[0])), dim=0\n )\n p_diag = self.initial_value * torch.einsum(\"i,ij->ij\", self.f_a, y)\n # torch.diag_indices(l_max + 1)\n diag_indices = torch.stack(\n [torch.arange(0, self.l_max + 1, device=x.device)] * 2, dim=0\n )\n p[(diag_indices[0][1:], diag_indices[1][1:])] = p_diag\n\n diag_indices = torch.stack(\n [torch.arange(0, self.l_max, device=x.device)] * 2, dim=0\n )\n\n # Compute the off-diagonal entries with recurrence.\n p_offdiag = torch.einsum(\n \"ij,ij->ij\",\n torch.einsum(\"i,j->ij\", self.f_b, x),\n p[(diag_indices[0], diag_indices[1])],\n ) # p[torch.diag_indices(l_max)])\n p[(diag_indices[0][: self.l_max], diag_indices[1][: self.l_max] + 1)] = (\n p_offdiag\n )\n\n # Compute the remaining entries with recurrence.\n if self.l_max > 1:\n for i in range(2, self.l_max + 1):\n p = self._recursive(i, p, x)\n return p\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/validation.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport torch\nimport torch.utils.data.distributed\nimport wandb\nfrom torch.nn import functional as F\n\nfrom unidepth.utils import barrier, is_main_process\nfrom unidepth.utils.misc import remove_padding\n\n\ndef original_image(batch, preds=None):\n paddings = [\n torch.tensor(pads)\n for img_meta in batch[\"img_metas\"]\n for pads in img_meta.get(\"paddings\", [[0] * 4])\n ]\n paddings = torch.stack(paddings).to(batch[\"data\"][\"image\"].device)[\n ..., [0, 2, 1, 3]\n ] # lrtb\n\n T, _, H, W = batch[\"data\"][\"depth\"].shape\n batch[\"data\"][\"image\"] = F.interpolate(\n batch[\"data\"][\"image\"],\n (H + paddings[2] + paddings[3], W + paddings[1] + paddings[2]),\n mode=\"bilinear\",\n align_corners=False,\n antialias=True,\n )\n batch[\"data\"][\"image\"] = remove_padding(\n batch[\"data\"][\"image\"], paddings.repeat(T, 1)\n )\n\n if preds is not None:\n for key in [\"depth\"]:\n if key in preds:\n preds[key] = F.interpolate(\n preds[key],\n (H + paddings[2] + paddings[3], W + paddings[1] + paddings[2]),\n mode=\"bilinear\",\n align_corners=False,\n antialias=True,\n )\n preds[key] = remove_padding(preds[key], paddings.repeat(T, 1))\n\n return batch, preds\n\n\ndef log_metrics(metrics_all, step):\n for name_ds, metrics in metrics_all.items():\n for metrics_name, metrics_value in metrics.items():\n try:\n print(f\"Metrics/{name_ds}/{metrics_name} {round(metrics_value, 4)}\")\n wandb.log(\n {f\"Metrics/{name_ds}/{metrics_name}\": metrics_value}, step=step\n )\n except:\n pass\n\n\ndef validate(model, test_loaders, step, context):\n metrics_all = {}\n for name_ds, test_loader in test_loaders.items():\n for i, batch in enumerate(test_loader):\n with context:\n batch[\"data\"] = {\n k: v.to(model.device) for k, v in batch[\"data\"].items()\n }\n # remove temporal dimension of the dataloder, here is always 1!\n batch[\"data\"] = {k: v.squeeze(1) for k, v in batch[\"data\"].items()}\n batch[\"img_metas\"] = [\n {k: v[0] for k, v in meta.items() if isinstance(v, list)}\n for meta in batch[\"img_metas\"]\n ]\n\n preds = model(batch[\"data\"], batch[\"img_metas\"])\n\n batch, _ = original_image(batch, preds=None)\n test_loader.dataset.accumulate_metrics(\n inputs=batch[\"data\"],\n preds=preds,\n keyframe_idx=batch[\"img_metas\"][0].get(\"keyframe_idx\"),\n )\n\n barrier()\n metrics_all[name_ds] = test_loader.dataset.get_evaluation()\n\n barrier()\n if is_main_process():\n log_metrics(metrics_all=metrics_all, step=step)\n return metrics_all\n" }, { "path": "camera_pose_annotation/depth_estimation/UniDepth/unidepth/utils/visualization.py", "content": "\"\"\"\nAuthor: Luigi Piccinelli\nLicensed under the CC-BY NC 4.0 license (http://creativecommons.org/licenses/by-nc/4.0/)\n\"\"\"\n\nimport os\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nimport wandb\nfrom PIL import Image\n\nfrom unidepth.utils.misc import ssi_helper\n\n\ndef colorize(\n value: np.ndarray, vmin: float = None, vmax: float = None, cmap: str = \"magma_r\"\n):\n # if already RGB, do nothing\n if value.ndim > 2:\n if value.shape[-1] > 1:\n return value\n value = value[..., 0]\n invalid_mask = value < 0.0001\n # normalize\n vmin = value.min() if vmin is None else vmin\n vmax = value.max() if vmax is None else vmax\n value = (value - vmin) / (vmax - vmin) # vmin..vmax\n\n # set color\n cmapper = plt.get_cmap(cmap)\n value = cmapper(value, bytes=True) # (nxmx4)\n value[invalid_mask] = 0\n img = value[..., :3]\n return img\n\n\ndef image_grid(imgs: list[np.ndarray], rows: int, cols: int) -> np.ndarray:\n if not len(imgs):\n return None\n assert len(imgs) == rows * cols\n h, w = imgs[0].shape[:2]\n grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n\n for i, img in enumerate(imgs):\n grid.paste(\n Image.fromarray(img.astype(np.uint8)).resize(\n (w, h), resample=Image.BILINEAR\n ),\n box=(i % cols * w, i // cols * h),\n )\n\n return np.array(grid)\n\n\ndef get_pointcloud_from_rgbd(\n image: np.array,\n depth: np.array,\n mask: np.ndarray,\n intrinsic_matrix: np.array,\n extrinsic_matrix: np.array = None,\n):\n depth = np.array(depth).squeeze()\n mask = np.array(mask).squeeze()\n # Mask the depth array\n masked_depth = np.ma.masked_where(mask == False, depth)\n # masked_depth = np.ma.masked_greater(masked_depth, 8000)\n # Create idx array\n idxs = np.indices(masked_depth.shape)\n u_idxs = idxs[1]\n v_idxs = idxs[0]\n # Get only non-masked depth and idxs\n z = masked_depth[~masked_depth.mask]\n compressed_u_idxs = u_idxs[~masked_depth.mask]\n compressed_v_idxs = v_idxs[~masked_depth.mask]\n image = np.stack(\n [image[..., i][~masked_depth.mask] for i in range(image.shape[-1])], axis=-1\n )\n\n # Calculate local position of each point\n # Apply vectorized math to depth using compressed arrays\n cx = intrinsic_matrix[0, 2]\n fx = intrinsic_matrix[0, 0]\n x = (compressed_u_idxs - cx) * z / fx\n cy = intrinsic_matrix[1, 2]\n fy = intrinsic_matrix[1, 1]\n # Flip y as we want +y pointing up not down\n y = -((compressed_v_idxs - cy) * z / fy)\n\n # # Apply camera_matrix to pointcloud as to get the pointcloud in world coords\n # if extrinsic_matrix is not None:\n # # Calculate camera pose from extrinsic matrix\n # camera_matrix = np.linalg.inv(extrinsic_matrix)\n # # Create homogenous array of vectors by adding 4th entry of 1\n # # At the same time flip z as for eye space the camera is looking down the -z axis\n # w = np.ones(z.shape)\n # x_y_z_eye_hom = np.vstack((x, y, -z, w))\n # # Transform the points from eye space to world space\n # x_y_z_world = np.dot(camera_matrix, x_y_z_eye_hom)[:3]\n # return x_y_z_world.T\n # else:\n x_y_z_local = np.stack((x, y, z), axis=-1)\n return np.concatenate([x_y_z_local, image], axis=-1)\n\n\ndef save_file_ply(xyz, rgb, pc_file):\n if rgb.max() < 1.001:\n rgb = rgb * 255.0\n rgb = rgb.astype(np.uint8)\n # print(rgb)\n with open(pc_file, \"w\") as f:\n # headers\n f.writelines(\n [\n \"ply\\n\" \"format ascii 1.0\\n\",\n \"element vertex {}\\n\".format(xyz.shape[0]),\n \"property float x\\n\",\n \"property float y\\n\",\n \"property float z\\n\",\n \"property uchar red\\n\",\n \"property uchar green\\n\",\n \"property uchar blue\\n\",\n \"end_header\\n\",\n ]\n )\n\n for i in range(xyz.shape[0]):\n str_v = \"{:10.6f} {:10.6f} {:10.6f} {:d} {:d} {:d}\\n\".format(\n xyz[i, 0], xyz[i, 1], xyz[i, 2], rgb[i, 0], rgb[i, 1], rgb[i, 2]\n )\n f.write(str_v)\n\n\n# really awful fct... FIXME\ndef log_train_artifacts(rgbs, gts, preds, ds_name, step, infos={}):\n rgbs = [\n (127.5 * (rgb + 1))\n .clip(0, 255)\n .to(torch.uint8)\n .cpu()\n .detach()\n .permute(1, 2, 0)\n .numpy()\n for rgb in rgbs\n ]\n\n new_gts, new_preds = [], []\n if len(gts) > 0:\n for i, gt in enumerate(gts):\n scale, shift = ssi_helper(\n gts[i][gts[i] > 0].cpu().detach(), preds[i][gts[i] > 0].cpu().detach()\n )\n gt = gts[i].cpu().detach().squeeze().numpy()\n pred = (preds[i].cpu().detach() * scale + shift).squeeze().numpy()\n vmin = gt[gt > 0].min() if (gt > 0).any() else 0.0\n vmax = gt.max() if (gt > 0).any() else 0.1\n new_gts.append(colorize(gt, vmin=vmin, vmax=vmax))\n new_preds.append(colorize(pred, vmin=vmin, vmax=vmax))\n gts, preds = new_gts, new_preds\n else:\n preds = [\n colorize(pred.cpu().detach().squeeze().numpy(), 0.0, 80.0)\n for i, pred in enumerate(preds)\n ]\n\n num_additional, additionals = 0, []\n for name, info in infos.items():\n num_additional += 1\n if info.shape[1] == 3:\n additionals.extend(\n [\n (127.5 * (x + 1))\n .clip(0, 255)\n .to(torch.uint8)\n .cpu()\n .detach()\n .permute(1, 2, 0)\n .numpy()\n for x in info[:4]\n ]\n )\n else:\n additionals.extend(\n [\n colorize(x.cpu().detach().squeeze().numpy())\n for i, x in enumerate(info[:4])\n ]\n )\n\n num_rows = 2 + int(len(gts) > 0) + num_additional\n artifacts_grid = image_grid(\n [*rgbs, *gts, *preds, *additionals], num_rows, len(rgbs)\n )\n try:\n wandb.log({f\"{ds_name}_training\": [wandb.Image(artifacts_grid)]}, step=step)\n except:\n Image.fromarray(artifacts_grid).save(\n os.path.join(os.environ[\"HOME\"], \"Workspace\", f\"art_grid{step}.png\")\n )\n print(\"Logging training images failed\")\n" }, { "path": "camera_pose_annotation/depth_estimation/__init__.py", "content": "" }, { "path": "camera_pose_annotation/dynamic_mask/__init__.py", "content": "" }, { "path": "camera_pose_annotation/dynamic_mask/inference_batch.py", "content": "\"\"\"\nBatch inference script for dynamic mask generation using SAM2.\nProcesses video frames to generate dynamic object masks based on motion probabilities.\n\"\"\"\n\nimport os\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom glob import glob\nimport cv2\nfrom scipy import ndimage\nfrom scipy.sparse import csr_matrix\nimport argparse\nimport pandas as pd\nimport subprocess\nfrom multiprocessing import Manager\nimport queue\nimport concurrent.futures\nfrom tqdm import tqdm\n\nfrom sam2.build_sam import build_sam2\nfrom sam2.sam2_image_predictor import SAM2ImagePredictor\n\n\ndef compress(dyn_masks, save_path=None):\n \"\"\"Compress dynamic masks using sparse matrix representation.\"\"\"\n assert save_path.endswith(\".npz\")\n # Transform to sparse matrices\n sparse_matrices_list = [csr_matrix(dyn_mask) for dyn_mask in dyn_masks]\n sparse_matrices = {}\n for i, dyn_mask in enumerate(sparse_matrices_list):\n sparse_matrices[f\"f_{i}_data\"] = dyn_mask.data\n sparse_matrices[f\"f_{i}_indices\"] = dyn_mask.indices\n sparse_matrices[f\"f_{i}_indptr\"] = dyn_mask.indptr\n if i == 0:\n sparse_matrices[\"shape\"] = dyn_mask.shape\n np.savez_compressed(save_path, **sparse_matrices)\n\n\ndef segment_sky(image):\n \"\"\"Segment sky regions from image using HSV color space and morphological operations.\"\"\"\n # Convert RGB to HSV\n hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)\n\n # Define range for blue color and create mask\n lower_blue = np.array([0, 0, 100])\n upper_blue = np.array([30, 255, 255])\n mask = cv2.inRange(hsv, lower_blue, upper_blue).view(bool)\n\n # Add luminous gray regions (likely sky)\n mask |= (hsv[:, :, 1] < 10) & (hsv[:, :, 2] > 150)\n mask |= (hsv[:, :, 1] < 30) & (hsv[:, :, 2] > 180)\n mask |= (hsv[:, :, 1] < 50) & (hsv[:, :, 2] > 220)\n\n # Morphological operations to clean up mask\n kernel = np.ones((5, 5), np.uint8)\n mask2 = ndimage.binary_opening(mask, structure=kernel)\n\n # Keep only largest connected components\n _, labels, stats, _ = cv2.connectedComponentsWithStats(\n mask2.view(np.uint8), connectivity=8\n )\n cc_sizes = stats[1:, cv2.CC_STAT_AREA]\n order = cc_sizes.argsort()[::-1] # Bigger first\n i = 0\n selection = []\n while i < len(order) and cc_sizes[order[i]] > cc_sizes[order[0]] / 2:\n selection.append(1 + order[i])\n i += 1\n mask3 = np.isin(labels, selection).reshape(labels.shape)\n\n # Return as tensor\n return torch.from_numpy(mask3)\n\n\ndef predict_mask(predictor, row, args, device):\n \"\"\"Generate dynamic masks for a video using SAM2 and motion probabilities.\"\"\"\n dir_path = os.path.join(args.dir_path, str(row[\"id\"]))\n if not os.path.exists(dir_path):\n print(f\"Directory '{dir_path}' not found. Exit.\")\n return\n img_dir = os.path.join(args.dir_path, str(row[\"id\"]), \"img\")\n if not os.path.exists(img_dir):\n print(f\"Image directory '{img_dir}' not found. Exit.\")\n return\n rec_dir = os.path.join(dir_path, \"reconstructions\")\n if not os.path.exists(rec_dir):\n print(f\"Reconstructions directory '{rec_dir}' not found. Exit.\")\n return\n prob_file = os.path.join(rec_dir, \"motion_prob.npy\")\n if not os.path.exists(prob_file):\n print(f\"Motion probability file '{prob_file}' not found. Exit.\")\n return\n\n compress_file = os.path.join(rec_dir, \"dyn_masks.npz\")\n if os.path.exists(compress_file):\n return\n\n # Load motion probabilities\n motion_probs = torch.from_numpy(np.load(prob_file)).to(device)\n\n # Load images\n images_list = list(sorted(glob(os.path.join(img_dir, \"*.jpg\"))))\n images = [cv2.imread(img_path) for img_path in images_list]\n\n if len(images) == 0 or len(images) != len(motion_probs):\n print(\n f\"{row['video_path']},Number of frames ({len(images)}) does not match number of motion probabilities ({len(motion_probs)}). Exit.\"\n )\n return\n\n width, height = images[0].shape[1], images[0].shape[0]\n area = width * height\n\n masks = []\n # Process each frame\n for i in range(len(images)):\n motion_prob = motion_probs[i].to(device)\n\n # Segment sky to avoid false detections\n sky_mask = segment_sky(images[i])\n\n predictor.set_image(images[i])\n\n # Adaptive thresholding based on motion probability distribution\n # We use an adaptive thresholding based on motion probability distribution to create initial masks. Then\n prob_min, prob_max = motion_prob.min(), motion_prob.max()\n threshold = (prob_max - prob_min) * 0.4 + prob_min\n if threshold > prob_max - 0.1:\n masks.append(np.zeros((height, width), dtype=np.uint8))\n continue\n\n # Create initial mask from motion probabilities\n mask = (motion_prob < threshold).float()\n mask = F.interpolate(\n mask.unsqueeze(0).unsqueeze(0),\n size=(height, width),\n mode=\"bilinear\",\n align_corners=False,\n ).squeeze()\n\n # Find contours and use them as SAM2 prompts\n mask_np = mask.cpu().numpy().astype(np.uint8)\n contours, _ = cv2.findContours(\n mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE\n )\n merged_mask = np.zeros_like(mask_np)\n for c in contours:\n points = []\n for point in c:\n points.append(point[0])\n points = np.array(points)\n\n # Sample points from contour as prompts\n interval = max(1, len(points) // 3)\n input_points = points[::interval].astype(np.float32)\n\n # Skip if points are in sky region\n if sky_mask[input_points[:, 1], input_points[:, 0]].any():\n continue\n\n input_labels = np.ones(input_points.shape[0], dtype=np.int64)\n\n # Use SAM2 to refine mask\n mask, score, _ = predictor.predict(\n point_coords=input_points,\n point_labels=input_labels,\n multimask_output=False,\n )\n\n # Skip if mask area is too large (likely background)\n if mask[0].sum() > area * 0.3:\n continue\n\n merged_mask = np.logical_or(merged_mask, mask[0])\n\n masks.append(merged_mask)\n\n # Save compressed masks\n masks = np.stack(masks, axis=0)\n compress(masks, compress_file)\n\n\ndef worker(task_queue, progress_queue, args, id):\n \"\"\"Worker function for parallel dynamic mask generation.\"\"\"\n gpu_id = id % args.gpu_num\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id) # Bind to specific GPU\n device = torch.device(f\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\n sam2_model = None\n predictor = None\n\n while True:\n try:\n index, row = task_queue.get_nowait()\n except queue.Empty:\n break\n\n # Initialize SAM2 model and predictor lazily\n if sam2_model is None:\n sam2_model = build_sam2(\n args.model_cfg, args.checkpoints_path, device=device\n )\n if predictor is None:\n predictor = SAM2ImagePredictor(sam2_model)\n\n predictor.reset_predictor()\n predict_mask(predictor, row, args, device)\n progress_queue.put(index)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for dynamic mask generation.\"\"\"\n parser = argparse.ArgumentParser(description=\"SAM2 Image Predictor\")\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\n \"--dir_path\",\n type=str,\n required=True,\n help=\"Path to the directory containing images and masks\",\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=16, help=\"#workers for concurrent.futures\"\n )\n parser.add_argument(\"--gpu_num\", type=int, default=1, help=\"gpu number\")\n parser.add_argument(\n \"--checkpoints_path\",\n type=str,\n default=\"checkpoints\",\n help=\"Path to the model checkpoint\",\n )\n parser.add_argument(\n \"--model_cfg\",\n type=str,\n default=\"configs/sam2.1/sam2.1_hiera_l.yaml\",\n help=\"Path to the model configuration file\",\n )\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n if not os.path.exists(args.csv_path):\n print(f\"Meta file '{args.csv_path}' not found. Exit.\")\n return\n\n # Set SAM2 checkpoint path\n args.checkpoints_path = os.path.join(\n args.checkpoints_path, \"SAM2/sam2.1_hiera_large.pt\"\n )\n\n df = pd.read_csv(args.csv_path)\n\n # Setup multiprocessing\n manager = Manager()\n task_queue = manager.Queue()\n progress_queue = manager.Queue()\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n # Process tasks with multiple workers\n with concurrent.futures.ProcessPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for id in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, progress_queue, args, id))\n\n processed = 0\n total_tasks = len(df)\n with tqdm(total=total_tasks, desc=\"Processing rows\") as pbar:\n while processed < total_tasks:\n try:\n progress_queue.get(timeout=1)\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and progress_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom hydra import initialize_config_module\nfrom hydra.core.global_hydra import GlobalHydra\n\nif not GlobalHydra.instance().is_initialized():\n initialize_config_module(\"sam2\", version_base=\"1.2\")\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/automatic_mask_generator.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n# Adapted from https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/automatic_mask_generator.py\nfrom typing import Any, Dict, List, Optional, Tuple\n\nimport numpy as np\nimport torch\nfrom torchvision.ops.boxes import batched_nms, box_area # type: ignore\n\nfrom sam2.modeling.sam2_base import SAM2Base\nfrom sam2.sam2_image_predictor import SAM2ImagePredictor\nfrom sam2.utils.amg import (\n area_from_rle,\n batch_iterator,\n batched_mask_to_box,\n box_xyxy_to_xywh,\n build_all_layer_point_grids,\n calculate_stability_score,\n coco_encode_rle,\n generate_crop_boxes,\n is_box_near_crop_edge,\n mask_to_rle_pytorch,\n MaskData,\n remove_small_regions,\n rle_to_mask,\n uncrop_boxes_xyxy,\n uncrop_masks,\n uncrop_points,\n)\n\n\nclass SAM2AutomaticMaskGenerator:\n def __init__(\n self,\n model: SAM2Base,\n points_per_side: Optional[int] = 32,\n points_per_batch: int = 64,\n pred_iou_thresh: float = 0.8,\n stability_score_thresh: float = 0.95,\n stability_score_offset: float = 1.0,\n mask_threshold: float = 0.0,\n box_nms_thresh: float = 0.7,\n crop_n_layers: int = 0,\n crop_nms_thresh: float = 0.7,\n crop_overlap_ratio: float = 512 / 1500,\n crop_n_points_downscale_factor: int = 1,\n point_grids: Optional[List[np.ndarray]] = None,\n min_mask_region_area: int = 0,\n output_mode: str = \"binary_mask\",\n use_m2m: bool = False,\n multimask_output: bool = True,\n **kwargs,\n ) -> None:\n \"\"\"\n Using a SAM 2 model, generates masks for the entire image.\n Generates a grid of point prompts over the image, then filters\n low quality and duplicate masks. The default settings are chosen\n for SAM 2 with a HieraL backbone.\n\n Arguments:\n model (Sam): The SAM 2 model to use for mask prediction.\n points_per_side (int or None): The number of points to be sampled\n along one side of the image. The total number of points is\n points_per_side**2. If None, 'point_grids' must provide explicit\n point sampling.\n points_per_batch (int): Sets the number of points run simultaneously\n by the model. Higher numbers may be faster but use more GPU memory.\n pred_iou_thresh (float): A filtering threshold in [0,1], using the\n model's predicted mask quality.\n stability_score_thresh (float): A filtering threshold in [0,1], using\n the stability of the mask under changes to the cutoff used to binarize\n the model's mask predictions.\n stability_score_offset (float): The amount to shift the cutoff when\n calculated the stability score.\n mask_threshold (float): Threshold for binarizing the mask logits\n box_nms_thresh (float): The box IoU cutoff used by non-maximal\n suppression to filter duplicate masks.\n crop_n_layers (int): If >0, mask prediction will be run again on\n crops of the image. Sets the number of layers to run, where each\n layer has 2**i_layer number of image crops.\n crop_nms_thresh (float): The box IoU cutoff used by non-maximal\n suppression to filter duplicate masks between different crops.\n crop_overlap_ratio (float): Sets the degree to which crops overlap.\n In the first crop layer, crops will overlap by this fraction of\n the image length. Later layers with more crops scale down this overlap.\n crop_n_points_downscale_factor (int): The number of points-per-side\n sampled in layer n is scaled down by crop_n_points_downscale_factor**n.\n point_grids (list(np.ndarray) or None): A list over explicit grids\n of points used for sampling, normalized to [0,1]. The nth grid in the\n list is used in the nth crop layer. Exclusive with points_per_side.\n min_mask_region_area (int): If >0, postprocessing will be applied\n to remove disconnected regions and holes in masks with area smaller\n than min_mask_region_area. Requires opencv.\n output_mode (str): The form masks are returned in. Can be 'binary_mask',\n 'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools.\n For large resolutions, 'binary_mask' may consume large amounts of\n memory.\n use_m2m (bool): Whether to add a one step refinement using previous mask predictions.\n multimask_output (bool): Whether to output multimask at each point of the grid.\n \"\"\"\n\n assert (points_per_side is None) != (\n point_grids is None\n ), \"Exactly one of points_per_side or point_grid must be provided.\"\n if points_per_side is not None:\n self.point_grids = build_all_layer_point_grids(\n points_per_side,\n crop_n_layers,\n crop_n_points_downscale_factor,\n )\n elif point_grids is not None:\n self.point_grids = point_grids\n else:\n raise ValueError(\"Can't have both points_per_side and point_grid be None.\")\n\n assert output_mode in [\n \"binary_mask\",\n \"uncompressed_rle\",\n \"coco_rle\",\n ], f\"Unknown output_mode {output_mode}.\"\n if output_mode == \"coco_rle\":\n try:\n from pycocotools import mask as mask_utils # type: ignore # noqa: F401\n except ImportError as e:\n print(\"Please install pycocotools\")\n raise e\n\n self.predictor = SAM2ImagePredictor(\n model,\n max_hole_area=min_mask_region_area,\n max_sprinkle_area=min_mask_region_area,\n )\n self.points_per_batch = points_per_batch\n self.pred_iou_thresh = pred_iou_thresh\n self.stability_score_thresh = stability_score_thresh\n self.stability_score_offset = stability_score_offset\n self.mask_threshold = mask_threshold\n self.box_nms_thresh = box_nms_thresh\n self.crop_n_layers = crop_n_layers\n self.crop_nms_thresh = crop_nms_thresh\n self.crop_overlap_ratio = crop_overlap_ratio\n self.crop_n_points_downscale_factor = crop_n_points_downscale_factor\n self.min_mask_region_area = min_mask_region_area\n self.output_mode = output_mode\n self.use_m2m = use_m2m\n self.multimask_output = multimask_output\n\n @classmethod\n def from_pretrained(cls, model_id: str, **kwargs) -> \"SAM2AutomaticMaskGenerator\":\n \"\"\"\n Load a pretrained model from the Hugging Face hub.\n\n Arguments:\n model_id (str): The Hugging Face repository ID.\n **kwargs: Additional arguments to pass to the model constructor.\n\n Returns:\n (SAM2AutomaticMaskGenerator): The loaded model.\n \"\"\"\n from sam2.build_sam import build_sam2_hf\n\n sam_model = build_sam2_hf(model_id, **kwargs)\n return cls(sam_model, **kwargs)\n\n @torch.no_grad()\n def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:\n \"\"\"\n Generates masks for the given image.\n\n Arguments:\n image (np.ndarray): The image to generate masks for, in HWC uint8 format.\n\n Returns:\n list(dict(str, any)): A list over records for masks. Each record is\n a dict containing the following keys:\n segmentation (dict(str, any) or np.ndarray): The mask. If\n output_mode='binary_mask', is an array of shape HW. Otherwise,\n is a dictionary containing the RLE.\n bbox (list(float)): The box around the mask, in XYWH format.\n area (int): The area in pixels of the mask.\n predicted_iou (float): The model's own prediction of the mask's\n quality. This is filtered by the pred_iou_thresh parameter.\n point_coords (list(list(float))): The point coordinates input\n to the model to generate this mask.\n stability_score (float): A measure of the mask's quality. This\n is filtered on using the stability_score_thresh parameter.\n crop_box (list(float)): The crop of the image used to generate\n the mask, given in XYWH format.\n \"\"\"\n\n # Generate masks\n mask_data = self._generate_masks(image)\n\n # Encode masks\n if self.output_mode == \"coco_rle\":\n mask_data[\"segmentations\"] = [\n coco_encode_rle(rle) for rle in mask_data[\"rles\"]\n ]\n elif self.output_mode == \"binary_mask\":\n mask_data[\"segmentations\"] = [rle_to_mask(rle) for rle in mask_data[\"rles\"]]\n else:\n mask_data[\"segmentations\"] = mask_data[\"rles\"]\n\n # Write mask records\n curr_anns = []\n for idx in range(len(mask_data[\"segmentations\"])):\n ann = {\n \"segmentation\": mask_data[\"segmentations\"][idx],\n \"area\": area_from_rle(mask_data[\"rles\"][idx]),\n \"bbox\": box_xyxy_to_xywh(mask_data[\"boxes\"][idx]).tolist(),\n \"predicted_iou\": mask_data[\"iou_preds\"][idx].item(),\n \"point_coords\": [mask_data[\"points\"][idx].tolist()],\n \"stability_score\": mask_data[\"stability_score\"][idx].item(),\n \"crop_box\": box_xyxy_to_xywh(mask_data[\"crop_boxes\"][idx]).tolist(),\n }\n curr_anns.append(ann)\n\n return curr_anns\n\n def _generate_masks(self, image: np.ndarray) -> MaskData:\n orig_size = image.shape[:2]\n crop_boxes, layer_idxs = generate_crop_boxes(\n orig_size, self.crop_n_layers, self.crop_overlap_ratio\n )\n\n # Iterate over image crops\n data = MaskData()\n for crop_box, layer_idx in zip(crop_boxes, layer_idxs):\n crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)\n data.cat(crop_data)\n\n # Remove duplicate masks between crops\n if len(crop_boxes) > 1:\n # Prefer masks from smaller crops\n scores = 1 / box_area(data[\"crop_boxes\"])\n scores = scores.to(data[\"boxes\"].device)\n keep_by_nms = batched_nms(\n data[\"boxes\"].float(),\n scores,\n torch.zeros_like(data[\"boxes\"][:, 0]), # categories\n iou_threshold=self.crop_nms_thresh,\n )\n data.filter(keep_by_nms)\n data.to_numpy()\n return data\n\n def _process_crop(\n self,\n image: np.ndarray,\n crop_box: List[int],\n crop_layer_idx: int,\n orig_size: Tuple[int, ...],\n ) -> MaskData:\n # Crop the image and calculate embeddings\n x0, y0, x1, y1 = crop_box\n cropped_im = image[y0:y1, x0:x1, :]\n cropped_im_size = cropped_im.shape[:2]\n self.predictor.set_image(cropped_im)\n\n # Get points for this crop\n points_scale = np.array(cropped_im_size)[None, ::-1]\n points_for_image = self.point_grids[crop_layer_idx] * points_scale\n\n # Generate masks for this crop in batches\n data = MaskData()\n for (points,) in batch_iterator(self.points_per_batch, points_for_image):\n batch_data = self._process_batch(\n points, cropped_im_size, crop_box, orig_size, normalize=True\n )\n data.cat(batch_data)\n del batch_data\n self.predictor.reset_predictor()\n\n # Remove duplicates within this crop.\n keep_by_nms = batched_nms(\n data[\"boxes\"].float(),\n data[\"iou_preds\"],\n torch.zeros_like(data[\"boxes\"][:, 0]), # categories\n iou_threshold=self.box_nms_thresh,\n )\n data.filter(keep_by_nms)\n\n # Return to the original image frame\n data[\"boxes\"] = uncrop_boxes_xyxy(data[\"boxes\"], crop_box)\n data[\"points\"] = uncrop_points(data[\"points\"], crop_box)\n data[\"crop_boxes\"] = torch.tensor([crop_box for _ in range(len(data[\"rles\"]))])\n\n return data\n\n def _process_batch(\n self,\n points: np.ndarray,\n im_size: Tuple[int, ...],\n crop_box: List[int],\n orig_size: Tuple[int, ...],\n normalize=False,\n ) -> MaskData:\n orig_h, orig_w = orig_size\n\n # Run model on this batch\n points = torch.as_tensor(\n points, dtype=torch.float32, device=self.predictor.device\n )\n in_points = self.predictor._transforms.transform_coords(\n points, normalize=normalize, orig_hw=im_size\n )\n in_labels = torch.ones(\n in_points.shape[0], dtype=torch.int, device=in_points.device\n )\n masks, iou_preds, low_res_masks = self.predictor._predict(\n in_points[:, None, :],\n in_labels[:, None],\n multimask_output=self.multimask_output,\n return_logits=True,\n )\n\n # Serialize predictions and store in MaskData\n data = MaskData(\n masks=masks.flatten(0, 1),\n iou_preds=iou_preds.flatten(0, 1),\n points=points.repeat_interleave(masks.shape[1], dim=0),\n low_res_masks=low_res_masks.flatten(0, 1),\n )\n del masks\n\n if not self.use_m2m:\n # Filter by predicted IoU\n if self.pred_iou_thresh > 0.0:\n keep_mask = data[\"iou_preds\"] > self.pred_iou_thresh\n data.filter(keep_mask)\n\n # Calculate and filter by stability score\n data[\"stability_score\"] = calculate_stability_score(\n data[\"masks\"], self.mask_threshold, self.stability_score_offset\n )\n if self.stability_score_thresh > 0.0:\n keep_mask = data[\"stability_score\"] >= self.stability_score_thresh\n data.filter(keep_mask)\n else:\n # One step refinement using previous mask predictions\n in_points = self.predictor._transforms.transform_coords(\n data[\"points\"], normalize=normalize, orig_hw=im_size\n )\n labels = torch.ones(\n in_points.shape[0], dtype=torch.int, device=in_points.device\n )\n masks, ious = self.refine_with_m2m(\n in_points, labels, data[\"low_res_masks\"], self.points_per_batch\n )\n data[\"masks\"] = masks.squeeze(1)\n data[\"iou_preds\"] = ious.squeeze(1)\n\n if self.pred_iou_thresh > 0.0:\n keep_mask = data[\"iou_preds\"] > self.pred_iou_thresh\n data.filter(keep_mask)\n\n data[\"stability_score\"] = calculate_stability_score(\n data[\"masks\"], self.mask_threshold, self.stability_score_offset\n )\n if self.stability_score_thresh > 0.0:\n keep_mask = data[\"stability_score\"] >= self.stability_score_thresh\n data.filter(keep_mask)\n\n # Threshold masks and calculate boxes\n data[\"masks\"] = data[\"masks\"] > self.mask_threshold\n data[\"boxes\"] = batched_mask_to_box(data[\"masks\"])\n\n # Filter boxes that touch crop boundaries\n keep_mask = ~is_box_near_crop_edge(\n data[\"boxes\"], crop_box, [0, 0, orig_w, orig_h]\n )\n if not torch.all(keep_mask):\n data.filter(keep_mask)\n\n # Compress to RLE\n data[\"masks\"] = uncrop_masks(data[\"masks\"], crop_box, orig_h, orig_w)\n data[\"rles\"] = mask_to_rle_pytorch(data[\"masks\"])\n del data[\"masks\"]\n\n return data\n\n @staticmethod\n def postprocess_small_regions(\n mask_data: MaskData, min_area: int, nms_thresh: float\n ) -> MaskData:\n \"\"\"\n Removes small disconnected regions and holes in masks, then reruns\n box NMS to remove any new duplicates.\n\n Edits mask_data in place.\n\n Requires open-cv as a dependency.\n \"\"\"\n if len(mask_data[\"rles\"]) == 0:\n return mask_data\n\n # Filter small disconnected regions and holes\n new_masks = []\n scores = []\n for rle in mask_data[\"rles\"]:\n mask = rle_to_mask(rle)\n\n mask, changed = remove_small_regions(mask, min_area, mode=\"holes\")\n unchanged = not changed\n mask, changed = remove_small_regions(mask, min_area, mode=\"islands\")\n unchanged = unchanged and not changed\n\n new_masks.append(torch.as_tensor(mask).unsqueeze(0))\n # Give score=0 to changed masks and score=1 to unchanged masks\n # so NMS will prefer ones that didn't need postprocessing\n scores.append(float(unchanged))\n\n # Recalculate boxes and remove any new duplicates\n masks = torch.cat(new_masks, dim=0)\n boxes = batched_mask_to_box(masks)\n keep_by_nms = batched_nms(\n boxes.float(),\n torch.as_tensor(scores),\n torch.zeros_like(boxes[:, 0]), # categories\n iou_threshold=nms_thresh,\n )\n\n # Only recalculate RLEs for masks that have changed\n for i_mask in keep_by_nms:\n if scores[i_mask] == 0.0:\n mask_torch = masks[i_mask].unsqueeze(0)\n mask_data[\"rles\"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]\n mask_data[\"boxes\"][i_mask] = boxes[i_mask] # update res directly\n mask_data.filter(keep_by_nms)\n\n return mask_data\n\n def refine_with_m2m(self, points, point_labels, low_res_masks, points_per_batch):\n new_masks = []\n new_iou_preds = []\n\n for cur_points, cur_point_labels, low_res_mask in batch_iterator(\n points_per_batch, points, point_labels, low_res_masks\n ):\n best_masks, best_iou_preds, _ = self.predictor._predict(\n cur_points[:, None, :],\n cur_point_labels[:, None],\n mask_input=low_res_mask[:, None, :],\n multimask_output=False,\n return_logits=True,\n )\n new_masks.append(best_masks)\n new_iou_preds.append(best_iou_preds)\n masks = torch.cat(new_masks, dim=0)\n return masks, torch.cat(new_iou_preds, dim=0)\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/benchmark.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport os\nimport time\n\nimport numpy as np\nimport torch\nfrom tqdm import tqdm\n\nfrom sam2.build_sam import build_sam2_video_predictor\n\n# Only cuda supported\nassert torch.cuda.is_available()\ndevice = torch.device(\"cuda\")\n\ntorch.autocast(device_type=\"cuda\", dtype=torch.bfloat16).__enter__()\nif torch.cuda.get_device_properties(0).major >= 8:\n # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)\n torch.backends.cuda.matmul.allow_tf32 = True\n torch.backends.cudnn.allow_tf32 = True\n\n# Config and checkpoint\nsam2_checkpoint = \"checkpoints/sam2.1_hiera_base_plus.pt\"\nmodel_cfg = \"configs/sam2.1/sam2.1_hiera_b+.yaml\"\n\n# Build video predictor with vos_optimized=True setting\npredictor = build_sam2_video_predictor(\n model_cfg, sam2_checkpoint, device=device, vos_optimized=True\n)\n\n\n# Initialize with video\nvideo_dir = \"notebooks/videos/bedroom\"\n# scan all the JPEG frame names in this directory\nframe_names = [\n p\n for p in os.listdir(video_dir)\n if os.path.splitext(p)[-1] in [\".jpg\", \".jpeg\", \".JPG\", \".JPEG\"]\n]\nframe_names.sort(key=lambda p: int(os.path.splitext(p)[0]))\ninference_state = predictor.init_state(video_path=video_dir)\n\n\n# Number of runs, warmup etc\nwarm_up, runs = 5, 25\nverbose = True\nnum_frames = len(frame_names)\ntotal, count = 0, 0\ntorch.cuda.empty_cache()\n\n# We will select an object with a click.\n# See video_predictor_example.ipynb for more detailed explanation\nann_frame_idx, ann_obj_id = 0, 1\n# Add a positive click at (x, y) = (210, 350)\n# For labels, `1` means positive click\npoints = np.array([[210, 350]], dtype=np.float32)\nlabels = np.array([1], np.int32)\n\n_, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(\n inference_state=inference_state,\n frame_idx=ann_frame_idx,\n obj_id=ann_obj_id,\n points=points,\n labels=labels,\n)\n\n# Warmup and then average FPS over several runs\nwith torch.autocast(\"cuda\", torch.bfloat16):\n with torch.inference_mode():\n for i in tqdm(range(runs), disable=not verbose, desc=\"Benchmarking\"):\n start = time.time()\n # Start tracking\n for (\n out_frame_idx,\n out_obj_ids,\n out_mask_logits,\n ) in predictor.propagate_in_video(inference_state):\n pass\n\n end = time.time()\n total += end - start\n count += 1\n if i == warm_up - 1:\n print(\"Warmup FPS: \", count * num_frames / total)\n total = 0\n count = 0\n\nprint(\"FPS: \", count * num_frames / total)\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/build_sam.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nimport os\n\nimport torch\nfrom hydra import compose\nfrom hydra.utils import instantiate\nfrom omegaconf import OmegaConf\n\nimport sam2\n\n# Check if the user is running Python from the parent directory of the sam2 repo\n# (i.e. the directory where this repo is cloned into) -- this is not supported since\n# it could shadow the sam2 package and cause issues.\nif os.path.isdir(os.path.join(sam2.__path__[0], \"sam2\")):\n # If the user has \"sam2/sam2\" in their path, they are likey importing the repo itself\n # as \"sam2\" rather than importing the \"sam2\" python package (i.e. \"sam2/sam2\" directory).\n # This typically happens because the user is running Python from the parent directory\n # that contains the sam2 repo they cloned.\n raise RuntimeError(\n \"You're likely running Python from the parent directory of the sam2 repository \"\n \"(i.e. the directory where https://github.com/facebookresearch/sam2 is cloned into). \"\n \"This is not supported since the `sam2` Python package could be shadowed by the \"\n \"repository name (the repository is also named `sam2` and contains the Python package \"\n \"in `sam2/sam2`). Please run Python from another directory (e.g. from the repo dir \"\n \"rather than its parent dir, or from your home directory) after installing SAM 2.\"\n )\n\n\nHF_MODEL_ID_TO_FILENAMES = {\n \"facebook/sam2-hiera-tiny\": (\n \"configs/sam2/sam2_hiera_t.yaml\",\n \"sam2_hiera_tiny.pt\",\n ),\n \"facebook/sam2-hiera-small\": (\n \"configs/sam2/sam2_hiera_s.yaml\",\n \"sam2_hiera_small.pt\",\n ),\n \"facebook/sam2-hiera-base-plus\": (\n \"configs/sam2/sam2_hiera_b+.yaml\",\n \"sam2_hiera_base_plus.pt\",\n ),\n \"facebook/sam2-hiera-large\": (\n \"configs/sam2/sam2_hiera_l.yaml\",\n \"sam2_hiera_large.pt\",\n ),\n \"facebook/sam2.1-hiera-tiny\": (\n \"configs/sam2.1/sam2.1_hiera_t.yaml\",\n \"sam2.1_hiera_tiny.pt\",\n ),\n \"facebook/sam2.1-hiera-small\": (\n \"configs/sam2.1/sam2.1_hiera_s.yaml\",\n \"sam2.1_hiera_small.pt\",\n ),\n \"facebook/sam2.1-hiera-base-plus\": (\n \"configs/sam2.1/sam2.1_hiera_b+.yaml\",\n \"sam2.1_hiera_base_plus.pt\",\n ),\n \"facebook/sam2.1-hiera-large\": (\n \"configs/sam2.1/sam2.1_hiera_l.yaml\",\n \"sam2.1_hiera_large.pt\",\n ),\n}\n\n\ndef build_sam2(\n config_file,\n ckpt_path=None,\n device=\"cuda\",\n mode=\"eval\",\n hydra_overrides_extra=[],\n apply_postprocessing=True,\n **kwargs,\n):\n\n if apply_postprocessing:\n hydra_overrides_extra = hydra_overrides_extra.copy()\n hydra_overrides_extra += [\n # dynamically fall back to multi-mask if the single mask is not stable\n \"++model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability=true\",\n \"++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_delta=0.05\",\n \"++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_thresh=0.98\",\n ]\n # Read config and init model\n cfg = compose(config_name=config_file, overrides=hydra_overrides_extra)\n OmegaConf.resolve(cfg)\n model = instantiate(cfg.model, _recursive_=True)\n _load_checkpoint(model, ckpt_path)\n model = model.to(device)\n if mode == \"eval\":\n model.eval()\n return model\n\n\ndef build_sam2_video_predictor(\n config_file,\n ckpt_path=None,\n device=\"cuda\",\n mode=\"eval\",\n hydra_overrides_extra=[],\n apply_postprocessing=True,\n vos_optimized=False,\n **kwargs,\n):\n hydra_overrides = [\n \"++model._target_=sam2.sam2_video_predictor.SAM2VideoPredictor\",\n ]\n if vos_optimized:\n hydra_overrides = [\n \"++model._target_=sam2.sam2_video_predictor.SAM2VideoPredictorVOS\",\n \"++model.compile_image_encoder=True\", # Let sam2_base handle this\n ]\n\n if apply_postprocessing:\n hydra_overrides_extra = hydra_overrides_extra.copy()\n hydra_overrides_extra += [\n # dynamically fall back to multi-mask if the single mask is not stable\n \"++model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability=true\",\n \"++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_delta=0.05\",\n \"++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_thresh=0.98\",\n # the sigmoid mask logits on interacted frames with clicks in the memory encoder so that the encoded masks are exactly as what users see from clicking\n \"++model.binarize_mask_from_pts_for_mem_enc=true\",\n # fill small holes in the low-res masks up to `fill_hole_area` (before resizing them to the original video resolution)\n \"++model.fill_hole_area=8\",\n ]\n hydra_overrides.extend(hydra_overrides_extra)\n\n # Read config and init model\n cfg = compose(config_name=config_file, overrides=hydra_overrides)\n OmegaConf.resolve(cfg)\n model = instantiate(cfg.model, _recursive_=True)\n _load_checkpoint(model, ckpt_path)\n model = model.to(device)\n if mode == \"eval\":\n model.eval()\n return model\n\n\ndef _hf_download(model_id):\n from huggingface_hub import hf_hub_download\n\n config_name, checkpoint_name = HF_MODEL_ID_TO_FILENAMES[model_id]\n ckpt_path = hf_hub_download(repo_id=model_id, filename=checkpoint_name)\n return config_name, ckpt_path\n\n\ndef build_sam2_hf(model_id, **kwargs):\n config_name, ckpt_path = _hf_download(model_id)\n return build_sam2(config_file=config_name, ckpt_path=ckpt_path, **kwargs)\n\n\ndef build_sam2_video_predictor_hf(model_id, **kwargs):\n config_name, ckpt_path = _hf_download(model_id)\n return build_sam2_video_predictor(\n config_file=config_name, ckpt_path=ckpt_path, **kwargs\n )\n\n\ndef _load_checkpoint(model, ckpt_path):\n if ckpt_path is not None:\n sd = torch.load(ckpt_path, map_location=\"cpu\", weights_only=True)[\"model\"]\n missing_keys, unexpected_keys = model.load_state_dict(sd)\n if missing_keys:\n logging.error(missing_keys)\n raise RuntimeError()\n if unexpected_keys:\n logging.error(unexpected_keys)\n raise RuntimeError()\n logging.info(\"Loaded checkpoint sucessfully\")\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2/sam2_hiera_b+.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 112\n num_heads: 2\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [896, 448, 224, 112]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2/sam2_hiera_l.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 144\n num_heads: 2\n stages: [2, 6, 36, 4]\n global_att_blocks: [23, 33, 43]\n window_pos_embed_bkg_spatial_size: [7, 7]\n window_spec: [8, 4, 16, 8]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [1152, 576, 288, 144]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2/sam2_hiera_s.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 96\n num_heads: 1\n stages: [1, 2, 11, 2]\n global_att_blocks: [7, 10, 13]\n window_pos_embed_bkg_spatial_size: [7, 7]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [768, 384, 192, 96]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2/sam2_hiera_t.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 96\n num_heads: 1\n stages: [1, 2, 7, 2]\n global_att_blocks: [5, 7, 9]\n window_pos_embed_bkg_spatial_size: [7, 7]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [768, 384, 192, 96]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n # SAM decoder\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n # HieraT does not currently support compilation, should always be set to False\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2.1/sam2.1_hiera_b+.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 112\n num_heads: 2\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [896, 448, 224, 112]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n no_obj_embed_spatial: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: true\n proj_tpos_enc_in_obj_ptrs: true\n use_signed_tpos_enc_to_obj_ptrs: true\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2.1/sam2.1_hiera_l.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 144\n num_heads: 2\n stages: [2, 6, 36, 4]\n global_att_blocks: [23, 33, 43]\n window_pos_embed_bkg_spatial_size: [7, 7]\n window_spec: [8, 4, 16, 8]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [1152, 576, 288, 144]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n no_obj_embed_spatial: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: true\n proj_tpos_enc_in_obj_ptrs: true\n use_signed_tpos_enc_to_obj_ptrs: true\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2.1/sam2.1_hiera_s.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 96\n num_heads: 1\n stages: [1, 2, 11, 2]\n global_att_blocks: [7, 10, 13]\n window_pos_embed_bkg_spatial_size: [7, 7]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [768, 384, 192, 96]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n no_obj_embed_spatial: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: true\n proj_tpos_enc_in_obj_ptrs: true\n use_signed_tpos_enc_to_obj_ptrs: true\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2.1/sam2.1_hiera_t.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 96\n num_heads: 1\n stages: [1, 2, 7, 2]\n global_att_blocks: [5, 7, 9]\n window_pos_embed_bkg_spatial_size: [7, 7]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [768, 384, 192, 96]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n # SAM decoder\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n no_obj_embed_spatial: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: true\n proj_tpos_enc_in_obj_ptrs: true\n use_signed_tpos_enc_to_obj_ptrs: true\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n # HieraT does not currently support compilation, should always be set to False\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/configs/sam2.1_training/sam2.1_hiera_b+_MOSE_finetune.yaml", "content": "# @package _global_\n\nscratch:\n resolution: 1024\n train_batch_size: 1\n num_train_workers: 10\n num_frames: 8\n max_num_objects: 3\n base_lr: 5.0e-6\n vision_lr: 3.0e-06\n phases_per_epoch: 1\n num_epochs: 40\n\ndataset:\n # PATHS to Dataset\n img_folder: null # PATH to MOSE JPEGImages folder\n gt_folder: null # PATH to MOSE Annotations folder\n file_list_txt: training/assets/MOSE_sample_train_list.txt # Optional PATH to filelist containing a subset of videos to be used for training\n multiplier: 2\n\n# Video transforms\nvos:\n train_transforms:\n - _target_: training.dataset.transforms.ComposeAPI\n transforms:\n - _target_: training.dataset.transforms.RandomHorizontalFlip\n consistent_transform: True\n - _target_: training.dataset.transforms.RandomAffine\n degrees: 25\n shear: 20\n image_interpolation: bilinear\n consistent_transform: True\n - _target_: training.dataset.transforms.RandomResizeAPI\n sizes: ${scratch.resolution}\n square: true\n consistent_transform: True\n - _target_: training.dataset.transforms.ColorJitter\n consistent_transform: True\n brightness: 0.1\n contrast: 0.03\n saturation: 0.03\n hue: null\n - _target_: training.dataset.transforms.RandomGrayscale\n p: 0.05\n consistent_transform: True\n - _target_: training.dataset.transforms.ColorJitter\n consistent_transform: False\n brightness: 0.1\n contrast: 0.05\n saturation: 0.05\n hue: null\n - _target_: training.dataset.transforms.ToTensorAPI\n - _target_: training.dataset.transforms.NormalizeAPI\n mean: [0.485, 0.456, 0.406]\n std: [0.229, 0.224, 0.225]\n\ntrainer:\n _target_: training.trainer.Trainer\n mode: train_only\n max_epochs: ${times:${scratch.num_epochs},${scratch.phases_per_epoch}}\n accelerator: cuda\n seed_value: 123\n\n model:\n _target_: training.model.sam2.SAM2Train\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 112\n num_heads: 2\n drop_path_rate: 0.1\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [896, 448, 224, 112]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: ${scratch.resolution}\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n no_obj_embed_spatial: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: true\n proj_tpos_enc_in_obj_ptrs: true\n use_signed_tpos_enc_to_obj_ptrs: true\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n # compile_image_encoder: False\n\n ####### Training specific params #######\n # box/point input and corrections\n prob_to_use_pt_input_for_train: 0.5\n prob_to_use_pt_input_for_eval: 0.0\n prob_to_use_box_input_for_train: 0.5 # 0.5*0.5 = 0.25 prob to use box instead of points\n prob_to_use_box_input_for_eval: 0.0\n prob_to_sample_from_gt_for_train: 0.1 # with a small prob, sampling correction points from GT mask instead of prediction errors\n num_frames_to_correct_for_train: 2 # iteratively sample on random 1~2 frames (always include the first frame)\n num_frames_to_correct_for_eval: 1 # only iteratively sample on first frame\n rand_frames_to_correct_for_train: True # random #init-cond-frame ~ 2\n add_all_frames_to_correct_as_cond: True # when a frame receives a correction click, it becomes a conditioning frame (even if it's not initially a conditioning frame)\n # maximum 2 initial conditioning frames\n num_init_cond_frames_for_train: 2\n rand_init_cond_frames_for_train: True # random 1~2\n num_correction_pt_per_frame: 7\n use_act_ckpt_iterative_pt_sampling: false\n \n\n \n num_init_cond_frames_for_eval: 1 # only mask on the first frame\n forward_backbone_per_frame_for_eval: True\n \n\n data:\n train:\n _target_: training.dataset.sam2_datasets.TorchTrainMixedDataset\n phases_per_epoch: ${scratch.phases_per_epoch}\n batch_sizes:\n - ${scratch.train_batch_size}\n\n datasets:\n - _target_: training.dataset.utils.RepeatFactorWrapper\n dataset:\n _target_: training.dataset.utils.ConcatDataset\n datasets:\n - _target_: training.dataset.vos_dataset.VOSDataset\n transforms: ${vos.train_transforms}\n training: true\n video_dataset:\n _target_: training.dataset.vos_raw_dataset.PNGRawDataset\n img_folder: ${dataset.img_folder}\n gt_folder: ${dataset.gt_folder}\n file_list_txt: ${dataset.file_list_txt}\n sampler:\n _target_: training.dataset.vos_sampler.RandomUniformSampler\n num_frames: ${scratch.num_frames}\n max_num_objects: ${scratch.max_num_objects}\n multiplier: ${dataset.multiplier}\n shuffle: True\n num_workers: ${scratch.num_train_workers}\n pin_memory: True\n drop_last: True\n collate_fn:\n _target_: training.utils.data_utils.collate_fn\n _partial_: true\n dict_key: all\n\n optim:\n amp:\n enabled: True\n amp_dtype: bfloat16\n\n optimizer:\n _target_: torch.optim.AdamW\n\n gradient_clip:\n _target_: training.optimizer.GradientClipper\n max_norm: 0.1\n norm_type: 2\n\n param_group_modifiers:\n - _target_: training.optimizer.layer_decay_param_modifier\n _partial_: True\n layer_decay_value: 0.9\n apply_to: 'image_encoder.trunk'\n overrides:\n - pattern: '*pos_embed*'\n value: 1.0\n\n options:\n lr:\n - scheduler:\n _target_: fvcore.common.param_scheduler.CosineParamScheduler\n start_value: ${scratch.base_lr}\n end_value: ${divide:${scratch.base_lr},10}\n - scheduler:\n _target_: fvcore.common.param_scheduler.CosineParamScheduler\n start_value: ${scratch.vision_lr}\n end_value: ${divide:${scratch.vision_lr},10}\n param_names:\n - 'image_encoder.*'\n weight_decay:\n - scheduler:\n _target_: fvcore.common.param_scheduler.ConstantParamScheduler\n value: 0.1\n - scheduler:\n _target_: fvcore.common.param_scheduler.ConstantParamScheduler\n value: 0.0\n param_names:\n - '*bias*'\n module_cls_names: ['torch.nn.LayerNorm']\n\n loss:\n all:\n _target_: training.loss_fns.MultiStepMultiMasksAndIous\n weight_dict:\n loss_mask: 20\n loss_dice: 1\n loss_iou: 1\n loss_class: 1\n supervise_all_iou: true\n iou_use_l1_loss: true\n pred_obj_scores: true\n focal_gamma_obj_score: 0.0\n focal_alpha_obj_score: -1.0\n\n distributed:\n backend: nccl\n find_unused_parameters: True\n\n logging:\n tensorboard_writer:\n _target_: training.utils.logger.make_tensorboard_logger\n log_dir: ${launcher.experiment_log_dir}/tensorboard\n flush_secs: 120\n should_log: True\n log_dir: ${launcher.experiment_log_dir}/logs\n log_freq: 10\n\n # initialize from a SAM 2 checkpoint\n checkpoint:\n save_dir: ${launcher.experiment_log_dir}/checkpoints\n save_freq: 0 # 0 only last checkpoint is saved.\n model_weight_initializer:\n _partial_: True\n _target_: training.utils.checkpoint_utils.load_state_dict_into_model\n strict: True\n ignore_unexpected_keys: null\n ignore_missing_keys: null\n\n state_dict:\n _target_: training.utils.checkpoint_utils.load_checkpoint_and_apply_kernels\n checkpoint_path: ./checkpoints/sam2.1_hiera_base_plus.pt # PATH to SAM 2.1 checkpoint\n ckpt_state_dict_keys: ['model']\n\nlauncher:\n num_nodes: 1\n gpus_per_node: 8\n experiment_log_dir: null # Path to log directory, defaults to ./sam2_logs/${config_name}\n\n# SLURM args if running on a cluster\nsubmitit:\n partition: null\n account: null\n qos: null\n cpus_per_task: 10\n use_cluster: false\n timeout_hour: 24\n name: null\n port_range: [10000, 65000]\n\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/csrc/connected_components.cu", "content": "// Copyright (c) Meta Platforms, Inc. and affiliates.\n// All rights reserved.\n\n// This source code is licensed under the license found in the\n// LICENSE file in the root directory of this source tree.\n\n// adapted from https://github.com/zsef123/Connected_components_PyTorch\n// with license found in the LICENSE_cctorch file in the root directory.\n#include \n#include \n#include \n#include \n#include \n#include \n\n// 2d\n#define BLOCK_ROWS 16\n#define BLOCK_COLS 16\n\nnamespace cc2d {\n\ntemplate \n__device__ __forceinline__ unsigned char hasBit(T bitmap, unsigned char pos) {\n return (bitmap >> pos) & 1;\n}\n\n__device__ int32_t find(const int32_t* s_buf, int32_t n) {\n while (s_buf[n] != n)\n n = s_buf[n];\n return n;\n}\n\n__device__ int32_t find_n_compress(int32_t* s_buf, int32_t n) {\n const int32_t id = n;\n while (s_buf[n] != n) {\n n = s_buf[n];\n s_buf[id] = n;\n }\n return n;\n}\n\n__device__ void union_(int32_t* s_buf, int32_t a, int32_t b) {\n bool done;\n do {\n a = find(s_buf, a);\n b = find(s_buf, b);\n\n if (a < b) {\n int32_t old = atomicMin(s_buf + b, a);\n done = (old == b);\n b = old;\n } else if (b < a) {\n int32_t old = atomicMin(s_buf + a, b);\n done = (old == a);\n a = old;\n } else\n done = true;\n\n } while (!done);\n}\n\n__global__ void\ninit_labeling(int32_t* label, const uint32_t W, const uint32_t H) {\n const uint32_t row = (blockIdx.y * blockDim.y + threadIdx.y) * 2;\n const uint32_t col = (blockIdx.x * blockDim.x + threadIdx.x) * 2;\n const uint32_t idx = row * W + col;\n\n if (row < H && col < W)\n label[idx] = idx;\n}\n\n__global__ void\nmerge(uint8_t* img, int32_t* label, const uint32_t W, const uint32_t H) {\n const uint32_t row = (blockIdx.y * blockDim.y + threadIdx.y) * 2;\n const uint32_t col = (blockIdx.x * blockDim.x + threadIdx.x) * 2;\n const uint32_t idx = row * W + col;\n\n if (row >= H || col >= W)\n return;\n\n uint32_t P = 0;\n\n if (img[idx])\n P |= 0x777;\n if (row + 1 < H && img[idx + W])\n P |= 0x777 << 4;\n if (col + 1 < W && img[idx + 1])\n P |= 0x777 << 1;\n\n if (col == 0)\n P &= 0xEEEE;\n if (col + 1 >= W)\n P &= 0x3333;\n else if (col + 2 >= W)\n P &= 0x7777;\n\n if (row == 0)\n P &= 0xFFF0;\n if (row + 1 >= H)\n P &= 0xFF;\n\n if (P > 0) {\n // If need check about top-left pixel(if flag the first bit) and hit the\n // top-left pixel\n if (hasBit(P, 0) && img[idx - W - 1]) {\n union_(label, idx, idx - 2 * W - 2); // top left block\n }\n\n if ((hasBit(P, 1) && img[idx - W]) || (hasBit(P, 2) && img[idx - W + 1]))\n union_(label, idx, idx - 2 * W); // top bottom block\n\n if (hasBit(P, 3) && img[idx + 2 - W])\n union_(label, idx, idx - 2 * W + 2); // top right block\n\n if ((hasBit(P, 4) && img[idx - 1]) || (hasBit(P, 8) && img[idx + W - 1]))\n union_(label, idx, idx - 2); // just left block\n }\n}\n\n__global__ void compression(int32_t* label, const int32_t W, const int32_t H) {\n const uint32_t row = (blockIdx.y * blockDim.y + threadIdx.y) * 2;\n const uint32_t col = (blockIdx.x * blockDim.x + threadIdx.x) * 2;\n const uint32_t idx = row * W + col;\n\n if (row < H && col < W)\n find_n_compress(label, idx);\n}\n\n__global__ void final_labeling(\n const uint8_t* img,\n int32_t* label,\n const int32_t W,\n const int32_t H) {\n const uint32_t row = (blockIdx.y * blockDim.y + threadIdx.y) * 2;\n const uint32_t col = (blockIdx.x * blockDim.x + threadIdx.x) * 2;\n const uint32_t idx = row * W + col;\n\n if (row >= H || col >= W)\n return;\n\n int32_t y = label[idx] + 1;\n\n if (img[idx])\n label[idx] = y;\n else\n label[idx] = 0;\n\n if (col + 1 < W) {\n if (img[idx + 1])\n label[idx + 1] = y;\n else\n label[idx + 1] = 0;\n\n if (row + 1 < H) {\n if (img[idx + W + 1])\n label[idx + W + 1] = y;\n else\n label[idx + W + 1] = 0;\n }\n }\n\n if (row + 1 < H) {\n if (img[idx + W])\n label[idx + W] = y;\n else\n label[idx + W] = 0;\n }\n}\n\n__global__ void init_counting(\n const int32_t* label,\n int32_t* count_init,\n const int32_t W,\n const int32_t H) {\n const uint32_t row = (blockIdx.y * blockDim.y + threadIdx.y);\n const uint32_t col = (blockIdx.x * blockDim.x + threadIdx.x);\n const uint32_t idx = row * W + col;\n\n if (row >= H || col >= W)\n return;\n\n int32_t y = label[idx];\n if (y > 0) {\n int32_t count_idx = y - 1;\n atomicAdd(count_init + count_idx, 1);\n }\n}\n\n__global__ void final_counting(\n const int32_t* label,\n const int32_t* count_init,\n int32_t* count_final,\n const int32_t W,\n const int32_t H) {\n const uint32_t row = (blockIdx.y * blockDim.y + threadIdx.y);\n const uint32_t col = (blockIdx.x * blockDim.x + threadIdx.x);\n const uint32_t idx = row * W + col;\n\n if (row >= H || col >= W)\n return;\n\n int32_t y = label[idx];\n if (y > 0) {\n int32_t count_idx = y - 1;\n count_final[idx] = count_init[count_idx];\n } else {\n count_final[idx] = 0;\n }\n}\n\n} // namespace cc2d\n\nstd::vector get_connected_componnets(\n const torch::Tensor& inputs) {\n AT_ASSERTM(inputs.is_cuda(), \"inputs must be a CUDA tensor\");\n AT_ASSERTM(inputs.ndimension() == 4, \"inputs must be [N, 1, H, W] shape\");\n AT_ASSERTM(\n inputs.scalar_type() == torch::kUInt8, \"inputs must be a uint8 type\");\n\n const uint32_t N = inputs.size(0);\n const uint32_t C = inputs.size(1);\n const uint32_t H = inputs.size(2);\n const uint32_t W = inputs.size(3);\n\n AT_ASSERTM(C == 1, \"inputs must be [N, 1, H, W] shape\");\n AT_ASSERTM((H % 2) == 0, \"height must be an even number\");\n AT_ASSERTM((W % 2) == 0, \"width must be an even number\");\n\n // label must be uint32_t\n auto label_options =\n torch::TensorOptions().dtype(torch::kInt32).device(inputs.device());\n torch::Tensor labels = torch::zeros({N, C, H, W}, label_options);\n torch::Tensor counts_init = torch::zeros({N, C, H, W}, label_options);\n torch::Tensor counts_final = torch::zeros({N, C, H, W}, label_options);\n\n dim3 grid = dim3(\n ((W + 1) / 2 + BLOCK_COLS - 1) / BLOCK_COLS,\n ((H + 1) / 2 + BLOCK_ROWS - 1) / BLOCK_ROWS);\n dim3 block = dim3(BLOCK_COLS, BLOCK_ROWS);\n dim3 grid_count =\n dim3((W + BLOCK_COLS) / BLOCK_COLS, (H + BLOCK_ROWS) / BLOCK_ROWS);\n dim3 block_count = dim3(BLOCK_COLS, BLOCK_ROWS);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n for (int n = 0; n < N; n++) {\n uint32_t offset = n * H * W;\n\n cc2d::init_labeling<<>>(\n labels.data_ptr() + offset, W, H);\n cc2d::merge<<>>(\n inputs.data_ptr() + offset,\n labels.data_ptr() + offset,\n W,\n H);\n cc2d::compression<<>>(\n labels.data_ptr() + offset, W, H);\n cc2d::final_labeling<<>>(\n inputs.data_ptr() + offset,\n labels.data_ptr() + offset,\n W,\n H);\n\n // get the counting of each pixel\n cc2d::init_counting<<>>(\n labels.data_ptr() + offset,\n counts_init.data_ptr() + offset,\n W,\n H);\n cc2d::final_counting<<>>(\n labels.data_ptr() + offset,\n counts_init.data_ptr() + offset,\n counts_final.data_ptr() + offset,\n W,\n H);\n }\n\n // returned values are [labels, counts]\n std::vector outputs;\n outputs.push_back(labels);\n outputs.push_back(counts_final);\n return outputs;\n}\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\n \"get_connected_componnets\",\n &get_connected_componnets,\n \"get_connected_componnets\");\n}\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/backbones/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/backbones/hieradet.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nfrom functools import partial\nfrom typing import List, Tuple, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom iopath.common.file_io import g_pathmgr\n\nfrom sam2.modeling.backbones.utils import (\n PatchEmbed,\n window_partition,\n window_unpartition,\n)\n\nfrom sam2.modeling.sam2_utils import DropPath, MLP\n\n\ndef do_pool(x: torch.Tensor, pool: nn.Module, norm: nn.Module = None) -> torch.Tensor:\n if pool is None:\n return x\n # (B, H, W, C) -> (B, C, H, W)\n x = x.permute(0, 3, 1, 2)\n x = pool(x)\n # (B, C, H', W') -> (B, H', W', C)\n x = x.permute(0, 2, 3, 1)\n if norm:\n x = norm(x)\n\n return x\n\n\nclass MultiScaleAttention(nn.Module):\n def __init__(\n self,\n dim: int,\n dim_out: int,\n num_heads: int,\n q_pool: nn.Module = None,\n ):\n super().__init__()\n\n self.dim = dim\n self.dim_out = dim_out\n self.num_heads = num_heads\n self.q_pool = q_pool\n self.qkv = nn.Linear(dim, dim_out * 3)\n self.proj = nn.Linear(dim_out, dim_out)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n B, H, W, _ = x.shape\n # qkv with shape (B, H * W, 3, nHead, C)\n qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1)\n # q, k, v with shape (B, H * W, nheads, C)\n q, k, v = torch.unbind(qkv, 2)\n\n # Q pooling (for downsample at stage changes)\n if self.q_pool:\n q = do_pool(q.reshape(B, H, W, -1), self.q_pool)\n H, W = q.shape[1:3] # downsampled shape\n q = q.reshape(B, H * W, self.num_heads, -1)\n\n # Torch's SDPA expects [B, nheads, H*W, C] so we transpose\n x = F.scaled_dot_product_attention(\n q.transpose(1, 2),\n k.transpose(1, 2),\n v.transpose(1, 2),\n )\n # Transpose back\n x = x.transpose(1, 2)\n x = x.reshape(B, H, W, -1)\n\n x = self.proj(x)\n\n return x\n\n\nclass MultiScaleBlock(nn.Module):\n def __init__(\n self,\n dim: int,\n dim_out: int,\n num_heads: int,\n mlp_ratio: float = 4.0,\n drop_path: float = 0.0,\n norm_layer: Union[nn.Module, str] = \"LayerNorm\",\n q_stride: Tuple[int, int] = None,\n act_layer: nn.Module = nn.GELU,\n window_size: int = 0,\n ):\n super().__init__()\n\n if isinstance(norm_layer, str):\n norm_layer = partial(getattr(nn, norm_layer), eps=1e-6)\n\n self.dim = dim\n self.dim_out = dim_out\n self.norm1 = norm_layer(dim)\n\n self.window_size = window_size\n\n self.pool, self.q_stride = None, q_stride\n if self.q_stride:\n self.pool = nn.MaxPool2d(\n kernel_size=q_stride, stride=q_stride, ceil_mode=False\n )\n\n self.attn = MultiScaleAttention(\n dim,\n dim_out,\n num_heads=num_heads,\n q_pool=self.pool,\n )\n self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n self.norm2 = norm_layer(dim_out)\n self.mlp = MLP(\n dim_out,\n int(dim_out * mlp_ratio),\n dim_out,\n num_layers=2,\n activation=act_layer,\n )\n\n if dim != dim_out:\n self.proj = nn.Linear(dim, dim_out)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n shortcut = x # B, H, W, C\n x = self.norm1(x)\n\n # Skip connection\n if self.dim != self.dim_out:\n shortcut = do_pool(self.proj(x), self.pool)\n\n # Window partition\n window_size = self.window_size\n if window_size > 0:\n H, W = x.shape[1], x.shape[2]\n x, pad_hw = window_partition(x, window_size)\n\n # Window Attention + Q Pooling (if stage change)\n x = self.attn(x)\n if self.q_stride:\n # Shapes have changed due to Q pooling\n window_size = self.window_size // self.q_stride[0]\n H, W = shortcut.shape[1:3]\n\n pad_h = (window_size - H % window_size) % window_size\n pad_w = (window_size - W % window_size) % window_size\n pad_hw = (H + pad_h, W + pad_w)\n\n # Reverse window partition\n if self.window_size > 0:\n x = window_unpartition(x, window_size, pad_hw, (H, W))\n\n x = shortcut + self.drop_path(x)\n # MLP\n x = x + self.drop_path(self.mlp(self.norm2(x)))\n return x\n\n\nclass Hiera(nn.Module):\n \"\"\"\n Reference: https://arxiv.org/abs/2306.00989\n \"\"\"\n\n def __init__(\n self,\n embed_dim: int = 96, # initial embed dim\n num_heads: int = 1, # initial number of heads\n drop_path_rate: float = 0.0, # stochastic depth\n q_pool: int = 3, # number of q_pool stages\n q_stride: Tuple[int, int] = (2, 2), # downsample stride bet. stages\n stages: Tuple[int, ...] = (2, 3, 16, 3), # blocks per stage\n dim_mul: float = 2.0, # dim_mul factor at stage shift\n head_mul: float = 2.0, # head_mul factor at stage shift\n window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14),\n # window size per stage, when not using global att.\n window_spec: Tuple[int, ...] = (\n 8,\n 4,\n 14,\n 7,\n ),\n # global attn in these blocks\n global_att_blocks: Tuple[int, ...] = (\n 12,\n 16,\n 20,\n ),\n weights_path=None,\n return_interm_layers=True, # return feats from every stage\n ):\n super().__init__()\n\n assert len(stages) == len(window_spec)\n self.window_spec = window_spec\n\n depth = sum(stages)\n self.q_stride = q_stride\n self.stage_ends = [sum(stages[:i]) - 1 for i in range(1, len(stages) + 1)]\n assert 0 <= q_pool <= len(self.stage_ends[:-1])\n self.q_pool_blocks = [x + 1 for x in self.stage_ends[:-1]][:q_pool]\n self.return_interm_layers = return_interm_layers\n\n self.patch_embed = PatchEmbed(\n embed_dim=embed_dim,\n )\n # Which blocks have global att?\n self.global_att_blocks = global_att_blocks\n\n # Windowed positional embedding (https://arxiv.org/abs/2311.05613)\n self.window_pos_embed_bkg_spatial_size = window_pos_embed_bkg_spatial_size\n self.pos_embed = nn.Parameter(\n torch.zeros(1, embed_dim, *self.window_pos_embed_bkg_spatial_size)\n )\n self.pos_embed_window = nn.Parameter(\n torch.zeros(1, embed_dim, self.window_spec[0], self.window_spec[0])\n )\n\n dpr = [\n x.item() for x in torch.linspace(0, drop_path_rate, depth)\n ] # stochastic depth decay rule\n\n cur_stage = 1\n self.blocks = nn.ModuleList()\n\n for i in range(depth):\n dim_out = embed_dim\n # lags by a block, so first block of\n # next stage uses an initial window size\n # of previous stage and final window size of current stage\n window_size = self.window_spec[cur_stage - 1]\n\n if self.global_att_blocks is not None:\n window_size = 0 if i in self.global_att_blocks else window_size\n\n if i - 1 in self.stage_ends:\n dim_out = int(embed_dim * dim_mul)\n num_heads = int(num_heads * head_mul)\n cur_stage += 1\n\n block = MultiScaleBlock(\n dim=embed_dim,\n dim_out=dim_out,\n num_heads=num_heads,\n drop_path=dpr[i],\n q_stride=self.q_stride if i in self.q_pool_blocks else None,\n window_size=window_size,\n )\n\n embed_dim = dim_out\n self.blocks.append(block)\n\n self.channel_list = (\n [self.blocks[i].dim_out for i in self.stage_ends[::-1]]\n if return_interm_layers\n else [self.blocks[-1].dim_out]\n )\n\n if weights_path is not None:\n with g_pathmgr.open(weights_path, \"rb\") as f:\n chkpt = torch.load(f, map_location=\"cpu\")\n logging.info(\"loading Hiera\", self.load_state_dict(chkpt, strict=False))\n\n def _get_pos_embed(self, hw: Tuple[int, int]) -> torch.Tensor:\n h, w = hw\n window_embed = self.pos_embed_window\n pos_embed = F.interpolate(self.pos_embed, size=(h, w), mode=\"bicubic\")\n pos_embed = pos_embed + window_embed.tile(\n [x // y for x, y in zip(pos_embed.shape, window_embed.shape)]\n )\n pos_embed = pos_embed.permute(0, 2, 3, 1)\n return pos_embed\n\n def forward(self, x: torch.Tensor) -> List[torch.Tensor]:\n x = self.patch_embed(x)\n # x: (B, H, W, C)\n\n # Add pos embed\n x = x + self._get_pos_embed(x.shape[1:3])\n\n outputs = []\n for i, blk in enumerate(self.blocks):\n x = blk(x)\n if (i == self.stage_ends[-1]) or (\n i in self.stage_ends and self.return_interm_layers\n ):\n feats = x.permute(0, 3, 1, 2)\n outputs.append(feats)\n\n return outputs\n\n def get_layer_id(self, layer_name):\n # https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33\n num_layers = self.get_num_layers()\n\n if layer_name.find(\"rel_pos\") != -1:\n return num_layers + 1\n elif layer_name.find(\"pos_embed\") != -1:\n return 0\n elif layer_name.find(\"patch_embed\") != -1:\n return 0\n elif layer_name.find(\"blocks\") != -1:\n return int(layer_name.split(\"blocks\")[1].split(\".\")[1]) + 1\n else:\n return num_layers + 1\n\n def get_num_layers(self) -> int:\n return len(self.blocks)\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/backbones/image_encoder.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import List, Optional\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass ImageEncoder(nn.Module):\n def __init__(\n self,\n trunk: nn.Module,\n neck: nn.Module,\n scalp: int = 0,\n ):\n super().__init__()\n self.trunk = trunk\n self.neck = neck\n self.scalp = scalp\n assert (\n self.trunk.channel_list == self.neck.backbone_channel_list\n ), f\"Channel dims of trunk and neck do not match. Trunk: {self.trunk.channel_list}, neck: {self.neck.backbone_channel_list}\"\n\n def forward(self, sample: torch.Tensor):\n # Forward through backbone\n features, pos = self.neck(self.trunk(sample))\n if self.scalp > 0:\n # Discard the lowest resolution features\n features, pos = features[: -self.scalp], pos[: -self.scalp]\n\n src = features[-1]\n output = {\n \"vision_features\": src,\n \"vision_pos_enc\": pos,\n \"backbone_fpn\": features,\n }\n return output\n\n\nclass FpnNeck(nn.Module):\n \"\"\"\n A modified variant of Feature Pyramid Network (FPN) neck\n (we remove output conv and also do bicubic interpolation similar to ViT\n pos embed interpolation)\n \"\"\"\n\n def __init__(\n self,\n position_encoding: nn.Module,\n d_model: int,\n backbone_channel_list: List[int],\n kernel_size: int = 1,\n stride: int = 1,\n padding: int = 0,\n fpn_interp_model: str = \"bilinear\",\n fuse_type: str = \"sum\",\n fpn_top_down_levels: Optional[List[int]] = None,\n ):\n \"\"\"Initialize the neck\n :param trunk: the backbone\n :param position_encoding: the positional encoding to use\n :param d_model: the dimension of the model\n :param neck_norm: the normalization to use\n \"\"\"\n super().__init__()\n self.position_encoding = position_encoding\n self.convs = nn.ModuleList()\n self.backbone_channel_list = backbone_channel_list\n self.d_model = d_model\n for dim in backbone_channel_list:\n current = nn.Sequential()\n current.add_module(\n \"conv\",\n nn.Conv2d(\n in_channels=dim,\n out_channels=d_model,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n ),\n )\n\n self.convs.append(current)\n self.fpn_interp_model = fpn_interp_model\n assert fuse_type in [\"sum\", \"avg\"]\n self.fuse_type = fuse_type\n\n # levels to have top-down features in its outputs\n # e.g. if fpn_top_down_levels is [2, 3], then only outputs of level 2 and 3\n # have top-down propagation, while outputs of level 0 and level 1 have only\n # lateral features from the same backbone level.\n if fpn_top_down_levels is None:\n # default is to have top-down features on all levels\n fpn_top_down_levels = range(len(self.convs))\n self.fpn_top_down_levels = list(fpn_top_down_levels)\n\n def forward(self, xs: List[torch.Tensor]):\n\n out = [None] * len(self.convs)\n pos = [None] * len(self.convs)\n assert len(xs) == len(self.convs)\n # fpn forward pass\n # see https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/fpn.py\n prev_features = None\n # forward in top-down order (from low to high resolution)\n n = len(self.convs) - 1\n for i in range(n, -1, -1):\n x = xs[i]\n lateral_features = self.convs[n - i](x)\n if i in self.fpn_top_down_levels and prev_features is not None:\n top_down_features = F.interpolate(\n prev_features.to(dtype=torch.float32),\n scale_factor=2.0,\n mode=self.fpn_interp_model,\n align_corners=(\n None if self.fpn_interp_model == \"nearest\" else False\n ),\n antialias=False,\n )\n prev_features = lateral_features + top_down_features\n if self.fuse_type == \"avg\":\n prev_features /= 2\n else:\n prev_features = lateral_features\n x_out = prev_features\n out[i] = x_out\n pos[i] = self.position_encoding(x_out).to(x_out.dtype)\n\n return out, pos\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/backbones/utils.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\"\"\"Some utilities for backbones, in particular for windowing\"\"\"\n\nfrom typing import Tuple\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef window_partition(x, window_size):\n \"\"\"\n Partition into non-overlapping windows with padding if needed.\n Args:\n x (tensor): input tokens with [B, H, W, C].\n window_size (int): window size.\n Returns:\n windows: windows after partition with [B * num_windows, window_size, window_size, C].\n (Hp, Wp): padded height and width before partition\n \"\"\"\n B, H, W, C = x.shape\n\n pad_h = (window_size - H % window_size) % window_size\n pad_w = (window_size - W % window_size) % window_size\n if pad_h > 0 or pad_w > 0:\n x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))\n Hp, Wp = H + pad_h, W + pad_w\n\n x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)\n windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C)\n return windows, (Hp, Wp)\n\n\ndef window_unpartition(windows, window_size, pad_hw, hw):\n \"\"\"\n Window unpartition into original sequences and removing padding.\n Args:\n x (tensor): input tokens with [B * num_windows, window_size, window_size, C].\n window_size (int): window size.\n pad_hw (Tuple): padded height and width (Hp, Wp).\n hw (Tuple): original height and width (H, W) before padding.\n Returns:\n x: unpartitioned sequences with [B, H, W, C].\n \"\"\"\n Hp, Wp = pad_hw\n H, W = hw\n B = windows.shape[0] // (Hp * Wp // window_size // window_size)\n x = windows.reshape(\n B, Hp // window_size, Wp // window_size, window_size, window_size, -1\n )\n x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1)\n\n if Hp > H or Wp > W:\n x = x[:, :H, :W, :]\n return x\n\n\nclass PatchEmbed(nn.Module):\n \"\"\"\n Image to Patch Embedding.\n \"\"\"\n\n def __init__(\n self,\n kernel_size: Tuple[int, ...] = (7, 7),\n stride: Tuple[int, ...] = (4, 4),\n padding: Tuple[int, ...] = (3, 3),\n in_chans: int = 3,\n embed_dim: int = 768,\n ):\n \"\"\"\n Args:\n kernel_size (Tuple): kernel size of the projection layer.\n stride (Tuple): stride of the projection layer.\n padding (Tuple): padding size of the projection layer.\n in_chans (int): Number of input image channels.\n embed_dim (int): embed_dim (int): Patch embedding dimension.\n \"\"\"\n super().__init__()\n self.proj = nn.Conv2d(\n in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding\n )\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n x = self.proj(x)\n # B C H W -> B H W C\n x = x.permute(0, 2, 3, 1)\n return x\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/memory_attention.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Optional\n\nimport torch\nfrom torch import nn, Tensor\n\nfrom sam2.modeling.sam.transformer import RoPEAttention\n\nfrom sam2.modeling.sam2_utils import get_activation_fn, get_clones\n\n\nclass MemoryAttentionLayer(nn.Module):\n\n def __init__(\n self,\n activation: str,\n cross_attention: nn.Module,\n d_model: int,\n dim_feedforward: int,\n dropout: float,\n pos_enc_at_attn: bool,\n pos_enc_at_cross_attn_keys: bool,\n pos_enc_at_cross_attn_queries: bool,\n self_attention: nn.Module,\n ):\n super().__init__()\n self.d_model = d_model\n self.dim_feedforward = dim_feedforward\n self.dropout_value = dropout\n self.self_attn = self_attention\n self.cross_attn_image = cross_attention\n\n # Implementation of Feedforward model\n self.linear1 = nn.Linear(d_model, dim_feedforward)\n self.dropout = nn.Dropout(dropout)\n self.linear2 = nn.Linear(dim_feedforward, d_model)\n\n self.norm1 = nn.LayerNorm(d_model)\n self.norm2 = nn.LayerNorm(d_model)\n self.norm3 = nn.LayerNorm(d_model)\n self.dropout1 = nn.Dropout(dropout)\n self.dropout2 = nn.Dropout(dropout)\n self.dropout3 = nn.Dropout(dropout)\n\n self.activation_str = activation\n self.activation = get_activation_fn(activation)\n\n # Where to add pos enc\n self.pos_enc_at_attn = pos_enc_at_attn\n self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries\n self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys\n\n def _forward_sa(self, tgt, query_pos):\n # Self-Attention\n tgt2 = self.norm1(tgt)\n q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2\n tgt2 = self.self_attn(q, k, v=tgt2)\n tgt = tgt + self.dropout1(tgt2)\n return tgt\n\n def _forward_ca(self, tgt, memory, query_pos, pos, num_k_exclude_rope=0):\n kwds = {}\n if num_k_exclude_rope > 0:\n assert isinstance(self.cross_attn_image, RoPEAttention)\n kwds = {\"num_k_exclude_rope\": num_k_exclude_rope}\n\n # Cross-Attention\n tgt2 = self.norm2(tgt)\n tgt2 = self.cross_attn_image(\n q=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2,\n k=memory + pos if self.pos_enc_at_cross_attn_keys else memory,\n v=memory,\n **kwds,\n )\n tgt = tgt + self.dropout2(tgt2)\n return tgt\n\n def forward(\n self,\n tgt,\n memory,\n pos: Optional[Tensor] = None,\n query_pos: Optional[Tensor] = None,\n num_k_exclude_rope: int = 0,\n ) -> torch.Tensor:\n\n # Self-Attn, Cross-Attn\n tgt = self._forward_sa(tgt, query_pos)\n tgt = self._forward_ca(tgt, memory, query_pos, pos, num_k_exclude_rope)\n # MLP\n tgt2 = self.norm3(tgt)\n tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))\n tgt = tgt + self.dropout3(tgt2)\n return tgt\n\n\nclass MemoryAttention(nn.Module):\n def __init__(\n self,\n d_model: int,\n pos_enc_at_input: bool,\n layer: nn.Module,\n num_layers: int,\n batch_first: bool = True, # Do layers expect batch first input?\n ):\n super().__init__()\n self.d_model = d_model\n self.layers = get_clones(layer, num_layers)\n self.num_layers = num_layers\n self.norm = nn.LayerNorm(d_model)\n self.pos_enc_at_input = pos_enc_at_input\n self.batch_first = batch_first\n\n def forward(\n self,\n curr: torch.Tensor, # self-attention inputs\n memory: torch.Tensor, # cross-attention inputs\n curr_pos: Optional[Tensor] = None, # pos_enc for self-attention inputs\n memory_pos: Optional[Tensor] = None, # pos_enc for cross-attention inputs\n num_obj_ptr_tokens: int = 0, # number of object pointer *tokens*\n ):\n if isinstance(curr, list):\n assert isinstance(curr_pos, list)\n assert len(curr) == len(curr_pos) == 1\n curr, curr_pos = (\n curr[0],\n curr_pos[0],\n )\n\n assert (\n curr.shape[1] == memory.shape[1]\n ), \"Batch size must be the same for curr and memory\"\n\n output = curr\n if self.pos_enc_at_input and curr_pos is not None:\n output = output + 0.1 * curr_pos\n\n if self.batch_first:\n # Convert to batch first\n output = output.transpose(0, 1)\n curr_pos = curr_pos.transpose(0, 1)\n memory = memory.transpose(0, 1)\n memory_pos = memory_pos.transpose(0, 1)\n\n for layer in self.layers:\n kwds = {}\n if isinstance(layer.cross_attn_image, RoPEAttention):\n kwds = {\"num_k_exclude_rope\": num_obj_ptr_tokens}\n\n output = layer(\n tgt=output,\n memory=memory,\n pos=memory_pos,\n query_pos=curr_pos,\n **kwds,\n )\n normed_output = self.norm(output)\n\n if self.batch_first:\n # Convert back to seq first\n normed_output = normed_output.transpose(0, 1)\n curr_pos = curr_pos.transpose(0, 1)\n\n return normed_output\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/memory_encoder.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom typing import Tuple\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom sam2.modeling.sam2_utils import DropPath, get_clones, LayerNorm2d\n\n\nclass MaskDownSampler(nn.Module):\n \"\"\"\n Progressively downsample a mask by total_stride, each time by stride.\n Note that LayerNorm is applied per *token*, like in ViT.\n\n With each downsample (by a factor stride**2), channel capacity increases by the same factor.\n In the end, we linearly project to embed_dim channels.\n \"\"\"\n\n def __init__(\n self,\n embed_dim=256,\n kernel_size=4,\n stride=4,\n padding=0,\n total_stride=16,\n activation=nn.GELU,\n ):\n super().__init__()\n num_layers = int(math.log2(total_stride) // math.log2(stride))\n assert stride**num_layers == total_stride\n self.encoder = nn.Sequential()\n mask_in_chans, mask_out_chans = 1, 1\n for _ in range(num_layers):\n mask_out_chans = mask_in_chans * (stride**2)\n self.encoder.append(\n nn.Conv2d(\n mask_in_chans,\n mask_out_chans,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n )\n )\n self.encoder.append(LayerNorm2d(mask_out_chans))\n self.encoder.append(activation())\n mask_in_chans = mask_out_chans\n\n self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1))\n\n def forward(self, x):\n return self.encoder(x)\n\n\n# Lightly adapted from ConvNext (https://github.com/facebookresearch/ConvNeXt)\nclass CXBlock(nn.Module):\n r\"\"\"ConvNeXt Block. There are two equivalent implementations:\n (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)\n (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back\n We use (2) as we find it slightly faster in PyTorch\n\n Args:\n dim (int): Number of input channels.\n drop_path (float): Stochastic depth rate. Default: 0.0\n layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.\n \"\"\"\n\n def __init__(\n self,\n dim,\n kernel_size=7,\n padding=3,\n drop_path=0.0,\n layer_scale_init_value=1e-6,\n use_dwconv=True,\n ):\n super().__init__()\n self.dwconv = nn.Conv2d(\n dim,\n dim,\n kernel_size=kernel_size,\n padding=padding,\n groups=dim if use_dwconv else 1,\n ) # depthwise conv\n self.norm = LayerNorm2d(dim, eps=1e-6)\n self.pwconv1 = nn.Linear(\n dim, 4 * dim\n ) # pointwise/1x1 convs, implemented with linear layers\n self.act = nn.GELU()\n self.pwconv2 = nn.Linear(4 * dim, dim)\n self.gamma = (\n nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)\n if layer_scale_init_value > 0\n else None\n )\n self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n\n def forward(self, x):\n input = x\n x = self.dwconv(x)\n x = self.norm(x)\n x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)\n x = self.pwconv1(x)\n x = self.act(x)\n x = self.pwconv2(x)\n if self.gamma is not None:\n x = self.gamma * x\n x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)\n\n x = input + self.drop_path(x)\n return x\n\n\nclass Fuser(nn.Module):\n def __init__(self, layer, num_layers, dim=None, input_projection=False):\n super().__init__()\n self.proj = nn.Identity()\n self.layers = get_clones(layer, num_layers)\n\n if input_projection:\n assert dim is not None\n self.proj = nn.Conv2d(dim, dim, kernel_size=1)\n\n def forward(self, x):\n # normally x: (N, C, H, W)\n x = self.proj(x)\n for layer in self.layers:\n x = layer(x)\n return x\n\n\nclass MemoryEncoder(nn.Module):\n def __init__(\n self,\n out_dim,\n mask_downsampler,\n fuser,\n position_encoding,\n in_dim=256, # in_dim of pix_feats\n ):\n super().__init__()\n\n self.mask_downsampler = mask_downsampler\n\n self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1)\n self.fuser = fuser\n self.position_encoding = position_encoding\n self.out_proj = nn.Identity()\n if out_dim != in_dim:\n self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1)\n\n def forward(\n self,\n pix_feat: torch.Tensor,\n masks: torch.Tensor,\n skip_mask_sigmoid: bool = False,\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n ## Process masks\n # sigmoid, so that less domain shift from gt masks which are bool\n if not skip_mask_sigmoid:\n masks = F.sigmoid(masks)\n masks = self.mask_downsampler(masks)\n\n ## Fuse pix_feats and downsampled masks\n # in case the visual features are on CPU, cast them to CUDA\n pix_feat = pix_feat.to(masks.device)\n\n x = self.pix_feat_proj(pix_feat)\n x = x + masks\n x = self.fuser(x)\n x = self.out_proj(x)\n\n pos = self.position_encoding(x).to(x.dtype)\n\n return {\"vision_features\": x, \"vision_pos_enc\": [pos]}\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/position_encoding.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom typing import Any, Optional, Tuple\n\nimport numpy as np\n\nimport torch\nfrom torch import nn\n\n\nclass PositionEmbeddingSine(nn.Module):\n \"\"\"\n This is a more standard version of the position embedding, very similar to the one\n used by the Attention Is All You Need paper, generalized to work on images.\n \"\"\"\n\n def __init__(\n self,\n num_pos_feats,\n temperature: int = 10000,\n normalize: bool = True,\n scale: Optional[float] = None,\n # Following settings only relevant\n # for warmping up cache for compilation\n warmup_cache: bool = True,\n image_size: int = 1024,\n strides: Tuple[int] = (4, 8, 16, 32),\n ):\n super().__init__()\n assert num_pos_feats % 2 == 0, \"Expecting even model width\"\n self.num_pos_feats = num_pos_feats // 2\n self.temperature = temperature\n self.normalize = normalize\n if scale is not None and normalize is False:\n raise ValueError(\"normalize should be True if scale is passed\")\n if scale is None:\n scale = 2 * math.pi\n self.scale = scale\n\n self.cache = {}\n if warmup_cache and torch.cuda.is_available():\n # Warmup cache for cuda, to help with compilation\n device = torch.device(\"cuda\")\n for stride in strides:\n cache_key = (image_size // stride, image_size // stride)\n self._pe(1, device, *cache_key)\n\n def _encode_xy(self, x, y):\n # The positions are expected to be normalized\n assert len(x) == len(y) and x.ndim == y.ndim == 1\n x_embed = x * self.scale\n y_embed = y * self.scale\n\n dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)\n dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)\n\n pos_x = x_embed[:, None] / dim_t\n pos_y = y_embed[:, None] / dim_t\n pos_x = torch.stack(\n (pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2\n ).flatten(1)\n pos_y = torch.stack(\n (pos_y[:, 0::2].sin(), pos_y[:, 1::2].cos()), dim=2\n ).flatten(1)\n return pos_x, pos_y\n\n @torch.no_grad()\n def encode_boxes(self, x, y, w, h):\n pos_x, pos_y = self._encode_xy(x, y)\n pos = torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim=1)\n return pos\n\n encode = encode_boxes # Backwards compatibility\n\n @torch.no_grad()\n def encode_points(self, x, y, labels):\n (bx, nx), (by, ny), (bl, nl) = x.shape, y.shape, labels.shape\n assert bx == by and nx == ny and bx == bl and nx == nl\n pos_x, pos_y = self._encode_xy(x.flatten(), y.flatten())\n pos_x, pos_y = pos_x.reshape(bx, nx, -1), pos_y.reshape(by, ny, -1)\n pos = torch.cat((pos_y, pos_x, labels[:, :, None]), dim=2)\n return pos\n\n @torch.no_grad()\n def _pe(self, B, device, *cache_key):\n H, W = cache_key\n if cache_key in self.cache:\n return self.cache[cache_key].to(device)[None].repeat(B, 1, 1, 1)\n\n y_embed = (\n torch.arange(1, H + 1, dtype=torch.float32, device=device)\n .view(1, -1, 1)\n .repeat(B, 1, W)\n )\n x_embed = (\n torch.arange(1, W + 1, dtype=torch.float32, device=device)\n .view(1, 1, -1)\n .repeat(B, H, 1)\n )\n\n if self.normalize:\n eps = 1e-6\n y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale\n x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale\n\n dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=device)\n dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)\n\n pos_x = x_embed[:, :, :, None] / dim_t\n pos_y = y_embed[:, :, :, None] / dim_t\n pos_x = torch.stack(\n (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4\n ).flatten(3)\n pos_y = torch.stack(\n (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4\n ).flatten(3)\n pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)\n self.cache[cache_key] = pos[0]\n return pos\n\n @torch.no_grad()\n def forward(self, x: torch.Tensor):\n B = x.shape[0]\n cache_key = (x.shape[-2], x.shape[-1])\n return self._pe(B, x.device, *cache_key)\n\n\nclass PositionEmbeddingRandom(nn.Module):\n \"\"\"\n Positional encoding using random spatial frequencies.\n \"\"\"\n\n def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:\n super().__init__()\n if scale is None or scale <= 0.0:\n scale = 1.0\n self.register_buffer(\n \"positional_encoding_gaussian_matrix\",\n scale * torch.randn((2, num_pos_feats)),\n )\n\n def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:\n \"\"\"Positionally encode points that are normalized to [0,1].\"\"\"\n # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape\n coords = 2 * coords - 1\n coords = coords @ self.positional_encoding_gaussian_matrix\n coords = 2 * np.pi * coords\n # outputs d_1 x ... x d_n x C shape\n return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)\n\n def forward(self, size: Tuple[int, int]) -> torch.Tensor:\n \"\"\"Generate positional encoding for a grid of the specified size.\"\"\"\n h, w = size\n device: Any = self.positional_encoding_gaussian_matrix.device\n grid = torch.ones((h, w), device=device, dtype=torch.float32)\n y_embed = grid.cumsum(dim=0) - 0.5\n x_embed = grid.cumsum(dim=1) - 0.5\n y_embed = y_embed / h\n x_embed = x_embed / w\n\n pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))\n return pe.permute(2, 0, 1) # C x H x W\n\n def forward_with_coords(\n self, coords_input: torch.Tensor, image_size: Tuple[int, int]\n ) -> torch.Tensor:\n \"\"\"Positionally encode points that are not normalized to [0,1].\"\"\"\n coords = coords_input.clone()\n coords[:, :, 0] = coords[:, :, 0] / image_size[1]\n coords[:, :, 1] = coords[:, :, 1] / image_size[0]\n return self._pe_encoding(coords.to(torch.float)) # B x N x C\n\n\n# Rotary Positional Encoding, adapted from:\n# 1. https://github.com/meta-llama/codellama/blob/main/llama/model.py\n# 2. https://github.com/naver-ai/rope-vit\n# 3. https://github.com/lucidrains/rotary-embedding-torch\n\n\ndef init_t_xy(end_x: int, end_y: int):\n t = torch.arange(end_x * end_y, dtype=torch.float32)\n t_x = (t % end_x).float()\n t_y = torch.div(t, end_x, rounding_mode=\"floor\").float()\n return t_x, t_y\n\n\ndef compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):\n freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))\n freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))\n\n t_x, t_y = init_t_xy(end_x, end_y)\n freqs_x = torch.outer(t_x, freqs_x)\n freqs_y = torch.outer(t_y, freqs_y)\n freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)\n freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)\n return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)\n\n\ndef reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):\n ndim = x.ndim\n assert 0 <= 1 < ndim\n assert freqs_cis.shape == (x.shape[-2], x.shape[-1])\n shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)]\n return freqs_cis.view(*shape)\n\n\ndef apply_rotary_enc(\n xq: torch.Tensor,\n xk: torch.Tensor,\n freqs_cis: torch.Tensor,\n repeat_freqs_k: bool = False,\n):\n xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n xk_ = (\n torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n if xk.shape[-2] != 0\n else None\n )\n freqs_cis = reshape_for_broadcast(freqs_cis, xq_)\n xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)\n if xk_ is None:\n # no keys to rotate, due to dropout\n return xq_out.type_as(xq).to(xq.device), xk\n # repeat freqs along seq_len dim to match k seq_len\n if repeat_freqs_k:\n r = xk_.shape[-2] // xq_.shape[-2]\n if freqs_cis.is_cuda:\n freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)\n else:\n # torch.repeat on complex numbers may not be supported on non-CUDA devices\n # (freqs_cis has 4 dims and we repeat on dim 2) so we use expand + flatten\n freqs_cis = freqs_cis.unsqueeze(2).expand(-1, -1, r, -1, -1).flatten(2, 3)\n xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)\n return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/sam/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/sam/mask_decoder.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import List, Optional, Tuple, Type\n\nimport torch\nfrom torch import nn\n\nfrom sam2.modeling.sam2_utils import LayerNorm2d, MLP\n\n\nclass MaskDecoder(nn.Module):\n def __init__(\n self,\n *,\n transformer_dim: int,\n transformer: nn.Module,\n num_multimask_outputs: int = 3,\n activation: Type[nn.Module] = nn.GELU,\n iou_head_depth: int = 3,\n iou_head_hidden_dim: int = 256,\n use_high_res_features: bool = False,\n iou_prediction_use_sigmoid=False,\n dynamic_multimask_via_stability=False,\n dynamic_multimask_stability_delta=0.05,\n dynamic_multimask_stability_thresh=0.98,\n pred_obj_scores: bool = False,\n pred_obj_scores_mlp: bool = False,\n use_multimask_token_for_obj_ptr: bool = False,\n ) -> None:\n \"\"\"\n Predicts masks given an image and prompt embeddings, using a\n transformer architecture.\n\n Arguments:\n transformer_dim (int): the channel dimension of the transformer\n transformer (nn.Module): the transformer used to predict masks\n num_multimask_outputs (int): the number of masks to predict\n when disambiguating masks\n activation (nn.Module): the type of activation to use when\n upscaling masks\n iou_head_depth (int): the depth of the MLP used to predict\n mask quality\n iou_head_hidden_dim (int): the hidden dimension of the MLP\n used to predict mask quality\n \"\"\"\n super().__init__()\n self.transformer_dim = transformer_dim\n self.transformer = transformer\n\n self.num_multimask_outputs = num_multimask_outputs\n\n self.iou_token = nn.Embedding(1, transformer_dim)\n self.num_mask_tokens = num_multimask_outputs + 1\n self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)\n\n self.pred_obj_scores = pred_obj_scores\n if self.pred_obj_scores:\n self.obj_score_token = nn.Embedding(1, transformer_dim)\n self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr\n\n self.output_upscaling = nn.Sequential(\n nn.ConvTranspose2d(\n transformer_dim, transformer_dim // 4, kernel_size=2, stride=2\n ),\n LayerNorm2d(transformer_dim // 4),\n activation(),\n nn.ConvTranspose2d(\n transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2\n ),\n activation(),\n )\n self.use_high_res_features = use_high_res_features\n if use_high_res_features:\n self.conv_s0 = nn.Conv2d(\n transformer_dim, transformer_dim // 8, kernel_size=1, stride=1\n )\n self.conv_s1 = nn.Conv2d(\n transformer_dim, transformer_dim // 4, kernel_size=1, stride=1\n )\n\n self.output_hypernetworks_mlps = nn.ModuleList(\n [\n MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)\n for i in range(self.num_mask_tokens)\n ]\n )\n\n self.iou_prediction_head = MLP(\n transformer_dim,\n iou_head_hidden_dim,\n self.num_mask_tokens,\n iou_head_depth,\n sigmoid_output=iou_prediction_use_sigmoid,\n )\n if self.pred_obj_scores:\n self.pred_obj_score_head = nn.Linear(transformer_dim, 1)\n if pred_obj_scores_mlp:\n self.pred_obj_score_head = MLP(transformer_dim, transformer_dim, 1, 3)\n\n # When outputting a single mask, optionally we can dynamically fall back to the best\n # multimask output token if the single mask output token gives low stability scores.\n self.dynamic_multimask_via_stability = dynamic_multimask_via_stability\n self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta\n self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh\n\n def forward(\n self,\n image_embeddings: torch.Tensor,\n image_pe: torch.Tensor,\n sparse_prompt_embeddings: torch.Tensor,\n dense_prompt_embeddings: torch.Tensor,\n multimask_output: bool,\n repeat_image: bool,\n high_res_features: Optional[List[torch.Tensor]] = None,\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Predict masks given image and prompt embeddings.\n\n Arguments:\n image_embeddings (torch.Tensor): the embeddings from the image encoder\n image_pe (torch.Tensor): positional encoding with the shape of image_embeddings\n sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes\n dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs\n multimask_output (bool): Whether to return multiple masks or a single\n mask.\n\n Returns:\n torch.Tensor: batched predicted masks\n torch.Tensor: batched predictions of mask quality\n torch.Tensor: batched SAM token for mask output\n \"\"\"\n masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks(\n image_embeddings=image_embeddings,\n image_pe=image_pe,\n sparse_prompt_embeddings=sparse_prompt_embeddings,\n dense_prompt_embeddings=dense_prompt_embeddings,\n repeat_image=repeat_image,\n high_res_features=high_res_features,\n )\n\n # Select the correct mask or masks for output\n if multimask_output:\n masks = masks[:, 1:, :, :]\n iou_pred = iou_pred[:, 1:]\n elif self.dynamic_multimask_via_stability and not self.training:\n masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred)\n else:\n masks = masks[:, 0:1, :, :]\n iou_pred = iou_pred[:, 0:1]\n\n if multimask_output and self.use_multimask_token_for_obj_ptr:\n sam_tokens_out = mask_tokens_out[:, 1:] # [b, 3, c] shape\n else:\n # Take the mask output token. Here we *always* use the token for single mask output.\n # At test time, even if we track after 1-click (and using multimask_output=True),\n # we still take the single mask token here. The rationale is that we always track\n # after multiple clicks during training, so the past tokens seen during training\n # are always the single mask token (and we'll let it be the object-memory token).\n sam_tokens_out = mask_tokens_out[:, 0:1] # [b, 1, c] shape\n\n # Prepare output\n return masks, iou_pred, sam_tokens_out, object_score_logits\n\n def predict_masks(\n self,\n image_embeddings: torch.Tensor,\n image_pe: torch.Tensor,\n sparse_prompt_embeddings: torch.Tensor,\n dense_prompt_embeddings: torch.Tensor,\n repeat_image: bool,\n high_res_features: Optional[List[torch.Tensor]] = None,\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Predicts masks. See 'forward' for more details.\"\"\"\n # Concatenate output tokens\n s = 0\n if self.pred_obj_scores:\n output_tokens = torch.cat(\n [\n self.obj_score_token.weight,\n self.iou_token.weight,\n self.mask_tokens.weight,\n ],\n dim=0,\n )\n s = 1\n else:\n output_tokens = torch.cat(\n [self.iou_token.weight, self.mask_tokens.weight], dim=0\n )\n output_tokens = output_tokens.unsqueeze(0).expand(\n sparse_prompt_embeddings.size(0), -1, -1\n )\n tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)\n\n # Expand per-image data in batch direction to be per-mask\n if repeat_image:\n src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)\n else:\n assert image_embeddings.shape[0] == tokens.shape[0]\n src = image_embeddings\n src = src + dense_prompt_embeddings\n assert (\n image_pe.size(0) == 1\n ), \"image_pe should have size 1 in batch dim (from `get_dense_pe()`)\"\n pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)\n b, c, h, w = src.shape\n\n # Run the transformer\n hs, src = self.transformer(src, pos_src, tokens)\n iou_token_out = hs[:, s, :]\n mask_tokens_out = hs[:, s + 1 : (s + 1 + self.num_mask_tokens), :]\n\n # Upscale mask embeddings and predict masks using the mask tokens\n src = src.transpose(1, 2).view(b, c, h, w)\n if not self.use_high_res_features:\n upscaled_embedding = self.output_upscaling(src)\n else:\n dc1, ln1, act1, dc2, act2 = self.output_upscaling\n feat_s0, feat_s1 = high_res_features\n upscaled_embedding = act1(ln1(dc1(src) + feat_s1))\n upscaled_embedding = act2(dc2(upscaled_embedding) + feat_s0)\n\n hyper_in_list: List[torch.Tensor] = []\n for i in range(self.num_mask_tokens):\n hyper_in_list.append(\n self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])\n )\n hyper_in = torch.stack(hyper_in_list, dim=1)\n b, c, h, w = upscaled_embedding.shape\n masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)\n\n # Generate mask quality predictions\n iou_pred = self.iou_prediction_head(iou_token_out)\n if self.pred_obj_scores:\n assert s == 1\n object_score_logits = self.pred_obj_score_head(hs[:, 0, :])\n else:\n # Obj scores logits - default to 10.0, i.e. assuming the object is present, sigmoid(10)=1\n object_score_logits = 10.0 * iou_pred.new_ones(iou_pred.shape[0], 1)\n\n return masks, iou_pred, mask_tokens_out, object_score_logits\n\n def _get_stability_scores(self, mask_logits):\n \"\"\"\n Compute stability scores of the mask logits based on the IoU between upper and\n lower thresholds.\n \"\"\"\n mask_logits = mask_logits.flatten(-2)\n stability_delta = self.dynamic_multimask_stability_delta\n area_i = torch.sum(mask_logits > stability_delta, dim=-1).float()\n area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float()\n stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0)\n return stability_scores\n\n def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):\n \"\"\"\n When outputting a single mask, if the stability score from the current single-mask\n output (based on output token 0) falls below a threshold, we instead select from\n multi-mask outputs (based on output token 1~3) the mask with the highest predicted\n IoU score. This is intended to ensure a valid mask for both clicking and tracking.\n \"\"\"\n # The best mask from multimask output tokens (1~3)\n multimask_logits = all_mask_logits[:, 1:, :, :]\n multimask_iou_scores = all_iou_scores[:, 1:]\n best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1)\n batch_inds = torch.arange(\n multimask_iou_scores.size(0), device=all_iou_scores.device\n )\n best_multimask_logits = multimask_logits[batch_inds, best_scores_inds]\n best_multimask_logits = best_multimask_logits.unsqueeze(1)\n best_multimask_iou_scores = multimask_iou_scores[batch_inds, best_scores_inds]\n best_multimask_iou_scores = best_multimask_iou_scores.unsqueeze(1)\n\n # The mask from singlemask output token 0 and its stability score\n singlemask_logits = all_mask_logits[:, 0:1, :, :]\n singlemask_iou_scores = all_iou_scores[:, 0:1]\n stability_scores = self._get_stability_scores(singlemask_logits)\n is_stable = stability_scores >= self.dynamic_multimask_stability_thresh\n\n # Dynamically fall back to best multimask output upon low stability scores.\n mask_logits_out = torch.where(\n is_stable[..., None, None].expand_as(singlemask_logits),\n singlemask_logits,\n best_multimask_logits,\n )\n iou_scores_out = torch.where(\n is_stable.expand_as(singlemask_iou_scores),\n singlemask_iou_scores,\n best_multimask_iou_scores,\n )\n return mask_logits_out, iou_scores_out\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/sam/prompt_encoder.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Optional, Tuple, Type\n\nimport torch\nfrom torch import nn\n\nfrom sam2.modeling.position_encoding import PositionEmbeddingRandom\n\nfrom sam2.modeling.sam2_utils import LayerNorm2d\n\n\nclass PromptEncoder(nn.Module):\n def __init__(\n self,\n embed_dim: int,\n image_embedding_size: Tuple[int, int],\n input_image_size: Tuple[int, int],\n mask_in_chans: int,\n activation: Type[nn.Module] = nn.GELU,\n ) -> None:\n \"\"\"\n Encodes prompts for input to SAM's mask decoder.\n\n Arguments:\n embed_dim (int): The prompts' embedding dimension\n image_embedding_size (tuple(int, int)): The spatial size of the\n image embedding, as (H, W).\n input_image_size (int): The padded size of the image as input\n to the image encoder, as (H, W).\n mask_in_chans (int): The number of hidden channels used for\n encoding input masks.\n activation (nn.Module): The activation to use when encoding\n input masks.\n \"\"\"\n super().__init__()\n self.embed_dim = embed_dim\n self.input_image_size = input_image_size\n self.image_embedding_size = image_embedding_size\n self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)\n\n self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners\n point_embeddings = [\n nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)\n ]\n self.point_embeddings = nn.ModuleList(point_embeddings)\n self.not_a_point_embed = nn.Embedding(1, embed_dim)\n\n self.mask_input_size = (\n 4 * image_embedding_size[0],\n 4 * image_embedding_size[1],\n )\n self.mask_downscaling = nn.Sequential(\n nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),\n LayerNorm2d(mask_in_chans // 4),\n activation(),\n nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),\n LayerNorm2d(mask_in_chans),\n activation(),\n nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),\n )\n self.no_mask_embed = nn.Embedding(1, embed_dim)\n\n def get_dense_pe(self) -> torch.Tensor:\n \"\"\"\n Returns the positional encoding used to encode point prompts,\n applied to a dense set of points the shape of the image encoding.\n\n Returns:\n torch.Tensor: Positional encoding with shape\n 1x(embed_dim)x(embedding_h)x(embedding_w)\n \"\"\"\n return self.pe_layer(self.image_embedding_size).unsqueeze(0)\n\n def _embed_points(\n self,\n points: torch.Tensor,\n labels: torch.Tensor,\n pad: bool,\n ) -> torch.Tensor:\n \"\"\"Embeds point prompts.\"\"\"\n points = points + 0.5 # Shift to center of pixel\n if pad:\n padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)\n padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)\n points = torch.cat([points, padding_point], dim=1)\n labels = torch.cat([labels, padding_label], dim=1)\n point_embedding = self.pe_layer.forward_with_coords(\n points, self.input_image_size\n )\n\n point_embedding = torch.where(\n (labels == -1).unsqueeze(-1),\n torch.zeros_like(point_embedding) + self.not_a_point_embed.weight,\n point_embedding,\n )\n point_embedding = torch.where(\n (labels == 0).unsqueeze(-1),\n point_embedding + self.point_embeddings[0].weight,\n point_embedding,\n )\n point_embedding = torch.where(\n (labels == 1).unsqueeze(-1),\n point_embedding + self.point_embeddings[1].weight,\n point_embedding,\n )\n point_embedding = torch.where(\n (labels == 2).unsqueeze(-1),\n point_embedding + self.point_embeddings[2].weight,\n point_embedding,\n )\n point_embedding = torch.where(\n (labels == 3).unsqueeze(-1),\n point_embedding + self.point_embeddings[3].weight,\n point_embedding,\n )\n return point_embedding\n\n def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:\n \"\"\"Embeds box prompts.\"\"\"\n boxes = boxes + 0.5 # Shift to center of pixel\n coords = boxes.reshape(-1, 2, 2)\n corner_embedding = self.pe_layer.forward_with_coords(\n coords, self.input_image_size\n )\n corner_embedding[:, 0, :] += self.point_embeddings[2].weight\n corner_embedding[:, 1, :] += self.point_embeddings[3].weight\n return corner_embedding\n\n def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:\n \"\"\"Embeds mask inputs.\"\"\"\n mask_embedding = self.mask_downscaling(masks)\n return mask_embedding\n\n def _get_batch_size(\n self,\n points: Optional[Tuple[torch.Tensor, torch.Tensor]],\n boxes: Optional[torch.Tensor],\n masks: Optional[torch.Tensor],\n ) -> int:\n \"\"\"\n Gets the batch size of the output given the batch size of the input prompts.\n \"\"\"\n if points is not None:\n return points[0].shape[0]\n elif boxes is not None:\n return boxes.shape[0]\n elif masks is not None:\n return masks.shape[0]\n else:\n return 1\n\n def _get_device(self) -> torch.device:\n return self.point_embeddings[0].weight.device\n\n def forward(\n self,\n points: Optional[Tuple[torch.Tensor, torch.Tensor]],\n boxes: Optional[torch.Tensor],\n masks: Optional[torch.Tensor],\n ) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Embeds different types of prompts, returning both sparse and dense\n embeddings.\n\n Arguments:\n points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates\n and labels to embed.\n boxes (torch.Tensor or none): boxes to embed\n masks (torch.Tensor or none): masks to embed\n\n Returns:\n torch.Tensor: sparse embeddings for the points and boxes, with shape\n BxNx(embed_dim), where N is determined by the number of input points\n and boxes.\n torch.Tensor: dense embeddings for the masks, in the shape\n Bx(embed_dim)x(embed_H)x(embed_W)\n \"\"\"\n bs = self._get_batch_size(points, boxes, masks)\n sparse_embeddings = torch.empty(\n (bs, 0, self.embed_dim), device=self._get_device()\n )\n if points is not None:\n coords, labels = points\n point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))\n sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)\n if boxes is not None:\n box_embeddings = self._embed_boxes(boxes)\n sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)\n\n if masks is not None:\n dense_embeddings = self._embed_masks(masks)\n else:\n dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(\n bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]\n )\n\n return sparse_embeddings, dense_embeddings\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/sam/transformer.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom functools import partial\nfrom typing import Tuple, Type\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn, Tensor\n\nfrom sam2.modeling.position_encoding import apply_rotary_enc, compute_axial_cis\nfrom sam2.modeling.sam2_utils import MLP\n\n\nclass TwoWayTransformer(nn.Module):\n def __init__(\n self,\n depth: int,\n embedding_dim: int,\n num_heads: int,\n mlp_dim: int,\n activation: Type[nn.Module] = nn.ReLU,\n attention_downsample_rate: int = 2,\n ) -> None:\n \"\"\"\n A transformer decoder that attends to an input image using\n queries whose positional embedding is supplied.\n\n Args:\n depth (int): number of layers in the transformer\n embedding_dim (int): the channel dimension for the input embeddings\n num_heads (int): the number of heads for multihead attention. Must\n divide embedding_dim\n mlp_dim (int): the channel dimension internal to the MLP block\n activation (nn.Module): the activation to use in the MLP block\n \"\"\"\n super().__init__()\n self.depth = depth\n self.embedding_dim = embedding_dim\n self.num_heads = num_heads\n self.mlp_dim = mlp_dim\n self.layers = nn.ModuleList()\n\n for i in range(depth):\n self.layers.append(\n TwoWayAttentionBlock(\n embedding_dim=embedding_dim,\n num_heads=num_heads,\n mlp_dim=mlp_dim,\n activation=activation,\n attention_downsample_rate=attention_downsample_rate,\n skip_first_layer_pe=(i == 0),\n )\n )\n\n self.final_attn_token_to_image = Attention(\n embedding_dim, num_heads, downsample_rate=attention_downsample_rate\n )\n self.norm_final_attn = nn.LayerNorm(embedding_dim)\n\n def forward(\n self,\n image_embedding: Tensor,\n image_pe: Tensor,\n point_embedding: Tensor,\n ) -> Tuple[Tensor, Tensor]:\n \"\"\"\n Args:\n image_embedding (torch.Tensor): image to attend to. Should be shape\n B x embedding_dim x h x w for any h and w.\n image_pe (torch.Tensor): the positional encoding to add to the image. Must\n have the same shape as image_embedding.\n point_embedding (torch.Tensor): the embedding to add to the query points.\n Must have shape B x N_points x embedding_dim for any N_points.\n\n Returns:\n torch.Tensor: the processed point_embedding\n torch.Tensor: the processed image_embedding\n \"\"\"\n # BxCxHxW -> BxHWxC == B x N_image_tokens x C\n bs, c, h, w = image_embedding.shape\n image_embedding = image_embedding.flatten(2).permute(0, 2, 1)\n image_pe = image_pe.flatten(2).permute(0, 2, 1)\n\n # Prepare queries\n queries = point_embedding\n keys = image_embedding\n\n # Apply transformer blocks and final layernorm\n for layer in self.layers:\n queries, keys = layer(\n queries=queries,\n keys=keys,\n query_pe=point_embedding,\n key_pe=image_pe,\n )\n\n # Apply the final attention layer from the points to the image\n q = queries + point_embedding\n k = keys + image_pe\n attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)\n queries = queries + attn_out\n queries = self.norm_final_attn(queries)\n\n return queries, keys\n\n\nclass TwoWayAttentionBlock(nn.Module):\n def __init__(\n self,\n embedding_dim: int,\n num_heads: int,\n mlp_dim: int = 2048,\n activation: Type[nn.Module] = nn.ReLU,\n attention_downsample_rate: int = 2,\n skip_first_layer_pe: bool = False,\n ) -> None:\n \"\"\"\n A transformer block with four layers: (1) self-attention of sparse\n inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp\n block on sparse inputs, and (4) cross attention of dense inputs to sparse\n inputs.\n\n Arguments:\n embedding_dim (int): the channel dimension of the embeddings\n num_heads (int): the number of heads in the attention layers\n mlp_dim (int): the hidden dimension of the mlp block\n activation (nn.Module): the activation of the mlp block\n skip_first_layer_pe (bool): skip the PE on the first layer\n \"\"\"\n super().__init__()\n self.self_attn = Attention(embedding_dim, num_heads)\n self.norm1 = nn.LayerNorm(embedding_dim)\n\n self.cross_attn_token_to_image = Attention(\n embedding_dim, num_heads, downsample_rate=attention_downsample_rate\n )\n self.norm2 = nn.LayerNorm(embedding_dim)\n\n self.mlp = MLP(\n embedding_dim, mlp_dim, embedding_dim, num_layers=2, activation=activation\n )\n self.norm3 = nn.LayerNorm(embedding_dim)\n\n self.norm4 = nn.LayerNorm(embedding_dim)\n self.cross_attn_image_to_token = Attention(\n embedding_dim, num_heads, downsample_rate=attention_downsample_rate\n )\n\n self.skip_first_layer_pe = skip_first_layer_pe\n\n def forward(\n self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor\n ) -> Tuple[Tensor, Tensor]:\n # Self attention block\n if self.skip_first_layer_pe:\n queries = self.self_attn(q=queries, k=queries, v=queries)\n else:\n q = queries + query_pe\n attn_out = self.self_attn(q=q, k=q, v=queries)\n queries = queries + attn_out\n queries = self.norm1(queries)\n\n # Cross attention block, tokens attending to image embedding\n q = queries + query_pe\n k = keys + key_pe\n attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)\n queries = queries + attn_out\n queries = self.norm2(queries)\n\n # MLP block\n mlp_out = self.mlp(queries)\n queries = queries + mlp_out\n queries = self.norm3(queries)\n\n # Cross attention block, image embedding attending to tokens\n q = queries + query_pe\n k = keys + key_pe\n attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)\n keys = keys + attn_out\n keys = self.norm4(keys)\n\n return queries, keys\n\n\nclass Attention(nn.Module):\n \"\"\"\n An attention layer that allows for downscaling the size of the embedding\n after projection to queries, keys, and values.\n \"\"\"\n\n def __init__(\n self,\n embedding_dim: int,\n num_heads: int,\n downsample_rate: int = 1,\n dropout: float = 0.0,\n kv_in_dim: int = None,\n ) -> None:\n super().__init__()\n self.embedding_dim = embedding_dim\n self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim\n self.internal_dim = embedding_dim // downsample_rate\n self.num_heads = num_heads\n assert (\n self.internal_dim % num_heads == 0\n ), \"num_heads must divide embedding_dim.\"\n\n self.q_proj = nn.Linear(embedding_dim, self.internal_dim)\n self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)\n self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)\n self.out_proj = nn.Linear(self.internal_dim, embedding_dim)\n\n self.dropout_p = dropout\n\n def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:\n b, n, c = x.shape\n x = x.reshape(b, n, num_heads, c // num_heads)\n return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head\n\n def _recombine_heads(self, x: Tensor) -> Tensor:\n b, n_heads, n_tokens, c_per_head = x.shape\n x = x.transpose(1, 2)\n return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C\n\n def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:\n # Input projections\n q = self.q_proj(q)\n k = self.k_proj(k)\n v = self.v_proj(v)\n\n # Separate into heads\n q = self._separate_heads(q, self.num_heads)\n k = self._separate_heads(k, self.num_heads)\n v = self._separate_heads(v, self.num_heads)\n\n dropout_p = self.dropout_p if self.training else 0.0\n # Attention\n out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)\n\n out = self._recombine_heads(out)\n out = self.out_proj(out)\n\n return out\n\n\nclass RoPEAttention(Attention):\n \"\"\"Attention with rotary position encoding.\"\"\"\n\n def __init__(\n self,\n *args,\n rope_theta=10000.0,\n # whether to repeat q rope to match k length\n # this is needed for cross-attention to memories\n rope_k_repeat=False,\n feat_sizes=(64, 64), # [w, h] for stride 16 feats at 1024 resolution\n **kwargs,\n ):\n super().__init__(*args, **kwargs)\n\n self.compute_cis = partial(\n compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta\n )\n freqs_cis = self.compute_cis(end_x=feat_sizes[0], end_y=feat_sizes[1])\n self.freqs_cis = (\n freqs_cis.to(\"cuda\") if torch.cuda.is_available() else freqs_cis\n )\n self.rope_k_repeat = rope_k_repeat\n\n def forward(\n self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0\n ) -> Tensor:\n # Input projections\n q = self.q_proj(q)\n k = self.k_proj(k)\n v = self.v_proj(v)\n\n # Separate into heads\n q = self._separate_heads(q, self.num_heads)\n k = self._separate_heads(k, self.num_heads)\n v = self._separate_heads(v, self.num_heads)\n\n # Apply rotary position encoding\n w = h = math.sqrt(q.shape[-2])\n self.freqs_cis = self.freqs_cis.to(q.device)\n if self.freqs_cis.shape[0] != q.shape[-2]:\n self.freqs_cis = self.compute_cis(end_x=w, end_y=h).to(q.device)\n if q.shape[-2] != k.shape[-2]:\n assert self.rope_k_repeat\n\n num_k_rope = k.size(-2) - num_k_exclude_rope\n q, k[:, :, :num_k_rope] = apply_rotary_enc(\n q,\n k[:, :, :num_k_rope],\n freqs_cis=self.freqs_cis,\n repeat_freqs_k=self.rope_k_repeat,\n )\n\n dropout_p = self.dropout_p if self.training else 0.0\n # Attention\n out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)\n\n out = self._recombine_heads(out)\n out = self.out_proj(out)\n\n return out\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/sam2_base.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport torch\nimport torch.distributed\nimport torch.nn.functional as F\n\nfrom torch.nn.init import trunc_normal_\n\nfrom sam2.modeling.sam.mask_decoder import MaskDecoder\nfrom sam2.modeling.sam.prompt_encoder import PromptEncoder\nfrom sam2.modeling.sam.transformer import TwoWayTransformer\nfrom sam2.modeling.sam2_utils import get_1d_sine_pe, MLP, select_closest_cond_frames\n\n# a large negative value as a placeholder score for missing objects\nNO_OBJ_SCORE = -1024.0\n\n\nclass SAM2Base(torch.nn.Module):\n def __init__(\n self,\n image_encoder,\n memory_attention,\n memory_encoder,\n num_maskmem=7, # default 1 input frame + 6 previous frames\n image_size=512,\n backbone_stride=16, # stride of the image backbone output\n sigmoid_scale_for_mem_enc=1.0, # scale factor for mask sigmoid prob\n sigmoid_bias_for_mem_enc=0.0, # bias factor for mask sigmoid prob\n # During evaluation, whether to binarize the sigmoid mask logits on interacted frames with clicks\n binarize_mask_from_pts_for_mem_enc=False,\n use_mask_input_as_output_without_sam=False, # on frames with mask input, whether to directly output the input mask without using a SAM prompt encoder + mask decoder\n # The maximum number of conditioning frames to participate in the memory attention (-1 means no limit; if there are more conditioning frames than this limit,\n # we only cross-attend to the temporally closest `max_cond_frames_in_attn` conditioning frames in the encoder when tracking each frame). This gives the model\n # a temporal locality when handling a large number of annotated frames (since closer frames should be more important) and also avoids GPU OOM.\n max_cond_frames_in_attn=-1,\n # on the first frame, whether to directly add the no-memory embedding to the image feature\n # (instead of using the transformer encoder)\n directly_add_no_mem_embed=False,\n # whether to use high-resolution feature maps in the SAM mask decoder\n use_high_res_features_in_sam=False,\n # whether to output multiple (3) masks for the first click on initial conditioning frames\n multimask_output_in_sam=False,\n # the minimum and maximum number of clicks to use multimask_output_in_sam (only relevant when `multimask_output_in_sam=True`;\n # default is 1 for both, meaning that only the first click gives multimask output; also note that a box counts as two points)\n multimask_min_pt_num=1,\n multimask_max_pt_num=1,\n # whether to also use multimask output for tracking (not just for the first click on initial conditioning frames; only relevant when `multimask_output_in_sam=True`)\n multimask_output_for_tracking=False,\n # Whether to use multimask tokens for obj ptr; Only relevant when both\n # use_obj_ptrs_in_encoder=True and multimask_output_for_tracking=True\n use_multimask_token_for_obj_ptr: bool = False,\n # whether to use sigmoid to restrict ious prediction to [0-1]\n iou_prediction_use_sigmoid=False,\n # The memory bank's temporal stride during evaluation (i.e. the `r` parameter in XMem and Cutie; XMem and Cutie use r=5).\n # For r>1, the (self.num_maskmem - 1) non-conditioning memory frames consist of\n # (self.num_maskmem - 2) nearest frames from every r-th frames, plus the last frame.\n memory_temporal_stride_for_eval=1,\n # whether to apply non-overlapping constraints on the object masks in the memory encoder during evaluation (to avoid/alleviate superposing masks)\n non_overlap_masks_for_mem_enc=False,\n # whether to cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder=False,\n # the maximum number of object pointers from other frames in encoder cross attention (only relevant when `use_obj_ptrs_in_encoder=True`)\n max_obj_ptrs_in_encoder=16,\n # whether to add temporal positional encoding to the object pointers in the encoder (only relevant when `use_obj_ptrs_in_encoder=True`)\n add_tpos_enc_to_obj_ptrs=True,\n # whether to add an extra linear projection layer for the temporal positional encoding in the object pointers to avoid potential interference\n # with spatial positional encoding (only relevant when both `use_obj_ptrs_in_encoder=True` and `add_tpos_enc_to_obj_ptrs=True`)\n proj_tpos_enc_in_obj_ptrs=False,\n # whether to use signed distance (instead of unsigned absolute distance) in the temporal positional encoding in the object pointers\n # (only relevant when both `use_obj_ptrs_in_encoder=True` and `add_tpos_enc_to_obj_ptrs=True`)\n use_signed_tpos_enc_to_obj_ptrs=False,\n # whether to only attend to object pointers in the past (before the current frame) in the encoder during evaluation\n # (only relevant when `use_obj_ptrs_in_encoder=True`; this might avoid pointer information too far in the future to distract the initial tracking)\n only_obj_ptrs_in_the_past_for_eval=False,\n # Whether to predict if there is an object in the frame\n pred_obj_scores: bool = False,\n # Whether to use an MLP to predict object scores\n pred_obj_scores_mlp: bool = False,\n # Only relevant if pred_obj_scores=True and use_obj_ptrs_in_encoder=True;\n # Whether to have a fixed no obj pointer when there is no object present\n # or to use it as an additive embedding with obj_ptr produced by decoder\n fixed_no_obj_ptr: bool = False,\n # Soft no object, i.e. mix in no_obj_ptr softly,\n # hope to make recovery easier if there is a mistake and mitigate accumulation of errors\n soft_no_obj_ptr: bool = False,\n use_mlp_for_obj_ptr_proj: bool = False,\n # add no obj embedding to spatial frames\n no_obj_embed_spatial: bool = False,\n # extra arguments used to construct the SAM mask decoder; if not None, it should be a dict of kwargs to be passed into `MaskDecoder` class.\n sam_mask_decoder_extra_args=None,\n compile_image_encoder: bool = False,\n ):\n super().__init__()\n\n # Part 1: the image backbone\n self.image_encoder = image_encoder\n # Use level 0, 1, 2 for high-res setting, or just level 2 for the default setting\n self.use_high_res_features_in_sam = use_high_res_features_in_sam\n self.num_feature_levels = 3 if use_high_res_features_in_sam else 1\n self.use_obj_ptrs_in_encoder = use_obj_ptrs_in_encoder\n self.max_obj_ptrs_in_encoder = max_obj_ptrs_in_encoder\n if use_obj_ptrs_in_encoder:\n # A conv layer to downsample the mask prompt to stride 4 (the same stride as\n # low-res SAM mask logits) and to change its scales from 0~1 to SAM logit scale,\n # so that it can be fed into the SAM mask decoder to generate a pointer.\n self.mask_downsample = torch.nn.Conv2d(1, 1, kernel_size=4, stride=4)\n self.add_tpos_enc_to_obj_ptrs = add_tpos_enc_to_obj_ptrs\n if proj_tpos_enc_in_obj_ptrs:\n assert add_tpos_enc_to_obj_ptrs # these options need to be used together\n self.proj_tpos_enc_in_obj_ptrs = proj_tpos_enc_in_obj_ptrs\n self.use_signed_tpos_enc_to_obj_ptrs = use_signed_tpos_enc_to_obj_ptrs\n self.only_obj_ptrs_in_the_past_for_eval = only_obj_ptrs_in_the_past_for_eval\n\n # Part 2: memory attention to condition current frame's visual features\n # with memories (and obj ptrs) from past frames\n self.memory_attention = memory_attention\n self.hidden_dim = image_encoder.neck.d_model\n\n # Part 3: memory encoder for the previous frame's outputs\n self.memory_encoder = memory_encoder\n self.mem_dim = self.hidden_dim\n if hasattr(self.memory_encoder, \"out_proj\") and hasattr(\n self.memory_encoder.out_proj, \"weight\"\n ):\n # if there is compression of memories along channel dim\n self.mem_dim = self.memory_encoder.out_proj.weight.shape[0]\n self.num_maskmem = num_maskmem # Number of memories accessible\n # Temporal encoding of the memories\n self.maskmem_tpos_enc = torch.nn.Parameter(\n torch.zeros(num_maskmem, 1, 1, self.mem_dim)\n )\n trunc_normal_(self.maskmem_tpos_enc, std=0.02)\n # a single token to indicate no memory embedding from previous frames\n self.no_mem_embed = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))\n self.no_mem_pos_enc = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))\n trunc_normal_(self.no_mem_embed, std=0.02)\n trunc_normal_(self.no_mem_pos_enc, std=0.02)\n self.directly_add_no_mem_embed = directly_add_no_mem_embed\n # Apply sigmoid to the output raw mask logits (to turn them from\n # range (-inf, +inf) to range (0, 1)) before feeding them into the memory encoder\n self.sigmoid_scale_for_mem_enc = sigmoid_scale_for_mem_enc\n self.sigmoid_bias_for_mem_enc = sigmoid_bias_for_mem_enc\n self.binarize_mask_from_pts_for_mem_enc = binarize_mask_from_pts_for_mem_enc\n self.non_overlap_masks_for_mem_enc = non_overlap_masks_for_mem_enc\n self.memory_temporal_stride_for_eval = memory_temporal_stride_for_eval\n # On frames with mask input, whether to directly output the input mask without\n # using a SAM prompt encoder + mask decoder\n self.use_mask_input_as_output_without_sam = use_mask_input_as_output_without_sam\n self.multimask_output_in_sam = multimask_output_in_sam\n self.multimask_min_pt_num = multimask_min_pt_num\n self.multimask_max_pt_num = multimask_max_pt_num\n self.multimask_output_for_tracking = multimask_output_for_tracking\n self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr\n self.iou_prediction_use_sigmoid = iou_prediction_use_sigmoid\n\n # Part 4: SAM-style prompt encoder (for both mask and point inputs)\n # and SAM-style mask decoder for the final mask output\n self.image_size = image_size\n self.backbone_stride = backbone_stride\n self.sam_mask_decoder_extra_args = sam_mask_decoder_extra_args\n self.pred_obj_scores = pred_obj_scores\n self.pred_obj_scores_mlp = pred_obj_scores_mlp\n self.fixed_no_obj_ptr = fixed_no_obj_ptr\n self.soft_no_obj_ptr = soft_no_obj_ptr\n if self.fixed_no_obj_ptr:\n assert self.pred_obj_scores\n assert self.use_obj_ptrs_in_encoder\n if self.pred_obj_scores and self.use_obj_ptrs_in_encoder:\n self.no_obj_ptr = torch.nn.Parameter(torch.zeros(1, self.hidden_dim))\n trunc_normal_(self.no_obj_ptr, std=0.02)\n self.use_mlp_for_obj_ptr_proj = use_mlp_for_obj_ptr_proj\n self.no_obj_embed_spatial = None\n if no_obj_embed_spatial:\n self.no_obj_embed_spatial = torch.nn.Parameter(torch.zeros(1, self.mem_dim))\n trunc_normal_(self.no_obj_embed_spatial, std=0.02)\n\n self._build_sam_heads()\n self.max_cond_frames_in_attn = max_cond_frames_in_attn\n\n # Model compilation\n if compile_image_encoder:\n # Compile the forward function (not the full module) to allow loading checkpoints.\n print(\n \"Image encoder compilation is enabled. First forward pass will be slow.\"\n )\n self.image_encoder.forward = torch.compile(\n self.image_encoder.forward,\n mode=\"max-autotune\",\n fullgraph=True,\n dynamic=False,\n )\n\n @property\n def device(self):\n return next(self.parameters()).device\n\n def forward(self, *args, **kwargs):\n raise NotImplementedError(\n \"Please use the corresponding methods in SAM2VideoPredictor for inference or SAM2Train for training/fine-tuning\"\n \"See notebooks/video_predictor_example.ipynb for an inference example.\"\n )\n\n def _build_sam_heads(self):\n \"\"\"Build SAM-style prompt encoder and mask decoder.\"\"\"\n self.sam_prompt_embed_dim = self.hidden_dim\n self.sam_image_embedding_size = self.image_size // self.backbone_stride\n\n # build PromptEncoder and MaskDecoder from SAM\n # (their hyperparameters like `mask_in_chans=16` are from SAM code)\n self.sam_prompt_encoder = PromptEncoder(\n embed_dim=self.sam_prompt_embed_dim,\n image_embedding_size=(\n self.sam_image_embedding_size,\n self.sam_image_embedding_size,\n ),\n input_image_size=(self.image_size, self.image_size),\n mask_in_chans=16,\n )\n self.sam_mask_decoder = MaskDecoder(\n num_multimask_outputs=3,\n transformer=TwoWayTransformer(\n depth=2,\n embedding_dim=self.sam_prompt_embed_dim,\n mlp_dim=2048,\n num_heads=8,\n ),\n transformer_dim=self.sam_prompt_embed_dim,\n iou_head_depth=3,\n iou_head_hidden_dim=256,\n use_high_res_features=self.use_high_res_features_in_sam,\n iou_prediction_use_sigmoid=self.iou_prediction_use_sigmoid,\n pred_obj_scores=self.pred_obj_scores,\n pred_obj_scores_mlp=self.pred_obj_scores_mlp,\n use_multimask_token_for_obj_ptr=self.use_multimask_token_for_obj_ptr,\n **(self.sam_mask_decoder_extra_args or {}),\n )\n if self.use_obj_ptrs_in_encoder:\n # a linear projection on SAM output tokens to turn them into object pointers\n self.obj_ptr_proj = torch.nn.Linear(self.hidden_dim, self.hidden_dim)\n if self.use_mlp_for_obj_ptr_proj:\n self.obj_ptr_proj = MLP(\n self.hidden_dim, self.hidden_dim, self.hidden_dim, 3\n )\n else:\n self.obj_ptr_proj = torch.nn.Identity()\n if self.proj_tpos_enc_in_obj_ptrs:\n # a linear projection on temporal positional encoding in object pointers to\n # avoid potential interference with spatial positional encoding\n self.obj_ptr_tpos_proj = torch.nn.Linear(self.hidden_dim, self.mem_dim)\n else:\n self.obj_ptr_tpos_proj = torch.nn.Identity()\n\n def _forward_sam_heads(\n self,\n backbone_features,\n point_inputs=None,\n mask_inputs=None,\n high_res_features=None,\n multimask_output=False,\n ):\n \"\"\"\n Forward SAM prompt encoders and mask heads.\n\n Inputs:\n - backbone_features: image features of [B, C, H, W] shape\n - point_inputs: a dictionary with \"point_coords\" and \"point_labels\", where\n 1) \"point_coords\" has [B, P, 2] shape and float32 dtype and contains the\n absolute pixel-unit coordinate in (x, y) format of the P input points\n 2) \"point_labels\" has shape [B, P] and int32 dtype, where 1 means\n positive clicks, 0 means negative clicks, and -1 means padding\n - mask_inputs: a mask of [B, 1, H*16, W*16] shape, float or bool, with the\n same spatial size as the image.\n - high_res_features: either 1) None or 2) or a list of length 2 containing\n two feature maps of [B, C, 4*H, 4*W] and [B, C, 2*H, 2*W] shapes respectively,\n which will be used as high-resolution feature maps for SAM decoder.\n - multimask_output: if it's True, we output 3 candidate masks and their 3\n corresponding IoU estimates, and if it's False, we output only 1 mask and\n its corresponding IoU estimate.\n\n Outputs:\n - low_res_multimasks: [B, M, H*4, W*4] shape (where M = 3 if\n `multimask_output=True` and M = 1 if `multimask_output=False`), the SAM\n output mask logits (before sigmoid) for the low-resolution masks, with 4x\n the resolution (1/4 stride) of the input backbone_features.\n - high_res_multimasks: [B, M, H*16, W*16] shape (where M = 3\n if `multimask_output=True` and M = 1 if `multimask_output=False`),\n upsampled from the low-resolution masks, with shape size as the image\n (stride is 1 pixel).\n - ious, [B, M] shape, where (where M = 3 if `multimask_output=True` and M = 1\n if `multimask_output=False`), the estimated IoU of each output mask.\n - low_res_masks: [B, 1, H*4, W*4] shape, the best mask in `low_res_multimasks`.\n If `multimask_output=True`, it's the mask with the highest IoU estimate.\n If `multimask_output=False`, it's the same as `low_res_multimasks`.\n - high_res_masks: [B, 1, H*16, W*16] shape, the best mask in `high_res_multimasks`.\n If `multimask_output=True`, it's the mask with the highest IoU estimate.\n If `multimask_output=False`, it's the same as `high_res_multimasks`.\n - obj_ptr: [B, C] shape, the object pointer vector for the output mask, extracted\n based on the output token from the SAM mask decoder.\n \"\"\"\n B = backbone_features.size(0)\n device = backbone_features.device\n assert backbone_features.size(1) == self.sam_prompt_embed_dim\n assert backbone_features.size(2) == self.sam_image_embedding_size\n assert backbone_features.size(3) == self.sam_image_embedding_size\n\n # a) Handle point prompts\n if point_inputs is not None:\n sam_point_coords = point_inputs[\"point_coords\"]\n sam_point_labels = point_inputs[\"point_labels\"]\n assert sam_point_coords.size(0) == B and sam_point_labels.size(0) == B\n else:\n # If no points are provide, pad with an empty point (with label -1)\n sam_point_coords = torch.zeros(B, 1, 2, device=device)\n sam_point_labels = -torch.ones(B, 1, dtype=torch.int32, device=device)\n\n # b) Handle mask prompts\n if mask_inputs is not None:\n # If mask_inputs is provided, downsize it into low-res mask input if needed\n # and feed it as a dense mask prompt into the SAM mask encoder\n assert len(mask_inputs.shape) == 4 and mask_inputs.shape[:2] == (B, 1)\n if mask_inputs.shape[-2:] != self.sam_prompt_encoder.mask_input_size:\n sam_mask_prompt = F.interpolate(\n mask_inputs.float(),\n size=self.sam_prompt_encoder.mask_input_size,\n align_corners=False,\n mode=\"bilinear\",\n antialias=True, # use antialias for downsampling\n )\n else:\n sam_mask_prompt = mask_inputs\n else:\n # Otherwise, simply feed None (and SAM's prompt encoder will add\n # a learned `no_mask_embed` to indicate no mask input in this case).\n sam_mask_prompt = None\n\n sparse_embeddings, dense_embeddings = self.sam_prompt_encoder(\n points=(sam_point_coords, sam_point_labels),\n boxes=None,\n masks=sam_mask_prompt,\n )\n (\n low_res_multimasks,\n ious,\n sam_output_tokens,\n object_score_logits,\n ) = self.sam_mask_decoder(\n image_embeddings=backbone_features,\n image_pe=self.sam_prompt_encoder.get_dense_pe(),\n sparse_prompt_embeddings=sparse_embeddings,\n dense_prompt_embeddings=dense_embeddings,\n multimask_output=multimask_output,\n repeat_image=False, # the image is already batched\n high_res_features=high_res_features,\n )\n if self.pred_obj_scores:\n is_obj_appearing = object_score_logits > 0\n\n # Mask used for spatial memories is always a *hard* choice between obj and no obj,\n # consistent with the actual mask prediction\n low_res_multimasks = torch.where(\n is_obj_appearing[:, None, None],\n low_res_multimasks,\n NO_OBJ_SCORE,\n )\n\n # convert masks from possibly bfloat16 (or float16) to float32\n # (older PyTorch versions before 2.1 don't support `interpolate` on bf16)\n low_res_multimasks = low_res_multimasks.float()\n high_res_multimasks = F.interpolate(\n low_res_multimasks,\n size=(self.image_size, self.image_size),\n mode=\"bilinear\",\n align_corners=False,\n )\n\n sam_output_token = sam_output_tokens[:, 0]\n if multimask_output:\n # take the best mask prediction (with the highest IoU estimation)\n best_iou_inds = torch.argmax(ious, dim=-1)\n batch_inds = torch.arange(B, device=device)\n low_res_masks = low_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1)\n high_res_masks = high_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1)\n if sam_output_tokens.size(1) > 1:\n sam_output_token = sam_output_tokens[batch_inds, best_iou_inds]\n else:\n low_res_masks, high_res_masks = low_res_multimasks, high_res_multimasks\n\n # Extract object pointer from the SAM output token (with occlusion handling)\n obj_ptr = self.obj_ptr_proj(sam_output_token)\n if self.pred_obj_scores:\n # Allow *soft* no obj ptr, unlike for masks\n if self.soft_no_obj_ptr:\n lambda_is_obj_appearing = object_score_logits.sigmoid()\n else:\n lambda_is_obj_appearing = is_obj_appearing.float()\n\n if self.fixed_no_obj_ptr:\n obj_ptr = lambda_is_obj_appearing * obj_ptr\n obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr\n\n return (\n low_res_multimasks,\n high_res_multimasks,\n ious,\n low_res_masks,\n high_res_masks,\n obj_ptr,\n object_score_logits,\n )\n\n def _use_mask_as_output(self, backbone_features, high_res_features, mask_inputs):\n \"\"\"\n Directly turn binary `mask_inputs` into a output mask logits without using SAM.\n (same input and output shapes as in _forward_sam_heads above).\n \"\"\"\n # Use -10/+10 as logits for neg/pos pixels (very close to 0/1 in prob after sigmoid).\n out_scale, out_bias = 20.0, -10.0 # sigmoid(-10.0)=4.5398e-05\n mask_inputs_float = mask_inputs.float()\n high_res_masks = mask_inputs_float * out_scale + out_bias\n low_res_masks = F.interpolate(\n high_res_masks,\n size=(high_res_masks.size(-2) // 4, high_res_masks.size(-1) // 4),\n align_corners=False,\n mode=\"bilinear\",\n antialias=True, # use antialias for downsampling\n )\n # a dummy IoU prediction of all 1's under mask input\n ious = mask_inputs.new_ones(mask_inputs.size(0), 1).float()\n if not self.use_obj_ptrs_in_encoder:\n # all zeros as a dummy object pointer (of shape [B, C])\n obj_ptr = torch.zeros(\n mask_inputs.size(0), self.hidden_dim, device=mask_inputs.device\n )\n else:\n # produce an object pointer using the SAM decoder from the mask input\n _, _, _, _, _, obj_ptr, _ = self._forward_sam_heads(\n backbone_features=backbone_features,\n mask_inputs=self.mask_downsample(mask_inputs_float),\n high_res_features=high_res_features,\n )\n # In this method, we are treating mask_input as output, e.g. using it directly to create spatial mem;\n # Below, we follow the same design axiom to use mask_input to decide if obj appears or not instead of relying\n # on the object_scores from the SAM decoder.\n is_obj_appearing = torch.any(mask_inputs.flatten(1).float() > 0.0, dim=1)\n is_obj_appearing = is_obj_appearing[..., None]\n lambda_is_obj_appearing = is_obj_appearing.float()\n object_score_logits = out_scale * lambda_is_obj_appearing + out_bias\n if self.pred_obj_scores:\n if self.fixed_no_obj_ptr:\n obj_ptr = lambda_is_obj_appearing * obj_ptr\n obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr\n\n return (\n low_res_masks,\n high_res_masks,\n ious,\n low_res_masks,\n high_res_masks,\n obj_ptr,\n object_score_logits,\n )\n\n def forward_image(self, img_batch: torch.Tensor):\n \"\"\"Get the image feature on the input batch.\"\"\"\n backbone_out = self.image_encoder(img_batch)\n if self.use_high_res_features_in_sam:\n # precompute projected level 0 and level 1 features in SAM decoder\n # to avoid running it again on every SAM click\n backbone_out[\"backbone_fpn\"][0] = self.sam_mask_decoder.conv_s0(\n backbone_out[\"backbone_fpn\"][0]\n )\n backbone_out[\"backbone_fpn\"][1] = self.sam_mask_decoder.conv_s1(\n backbone_out[\"backbone_fpn\"][1]\n )\n return backbone_out\n\n def _prepare_backbone_features(self, backbone_out):\n \"\"\"Prepare and flatten visual features.\"\"\"\n backbone_out = backbone_out.copy()\n assert len(backbone_out[\"backbone_fpn\"]) == len(backbone_out[\"vision_pos_enc\"])\n assert len(backbone_out[\"backbone_fpn\"]) >= self.num_feature_levels\n\n feature_maps = backbone_out[\"backbone_fpn\"][-self.num_feature_levels :]\n vision_pos_embeds = backbone_out[\"vision_pos_enc\"][-self.num_feature_levels :]\n\n feat_sizes = [(x.shape[-2], x.shape[-1]) for x in vision_pos_embeds]\n # flatten NxCxHxW to HWxNxC\n vision_feats = [x.flatten(2).permute(2, 0, 1) for x in feature_maps]\n vision_pos_embeds = [x.flatten(2).permute(2, 0, 1) for x in vision_pos_embeds]\n\n return backbone_out, vision_feats, vision_pos_embeds, feat_sizes\n\n def _prepare_memory_conditioned_features(\n self,\n frame_idx,\n is_init_cond_frame,\n current_vision_feats,\n current_vision_pos_embeds,\n feat_sizes,\n output_dict,\n num_frames,\n track_in_reverse=False, # tracking in reverse time order (for demo usage)\n ):\n \"\"\"Fuse the current frame's visual feature map with previous memory.\"\"\"\n B = current_vision_feats[-1].size(1) # batch size on this frame\n C = self.hidden_dim\n H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size\n device = current_vision_feats[-1].device\n # The case of `self.num_maskmem == 0` below is primarily used for reproducing SAM on images.\n # In this case, we skip the fusion with any memory.\n if self.num_maskmem == 0: # Disable memory and skip fusion\n pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W)\n return pix_feat\n\n num_obj_ptr_tokens = 0\n tpos_sign_mul = -1 if track_in_reverse else 1\n # Step 1: condition the visual features of the current frame on previous memories\n if not is_init_cond_frame:\n # Retrieve the memories encoded with the maskmem backbone\n to_cat_memory, to_cat_memory_pos_embed = [], []\n # Add conditioning frames's output first (all cond frames have t_pos=0 for\n # when getting temporal positional embedding below)\n assert len(output_dict[\"cond_frame_outputs\"]) > 0\n # Select a maximum number of temporally closest cond frames for cross attention\n cond_outputs = output_dict[\"cond_frame_outputs\"]\n selected_cond_outputs, unselected_cond_outputs = select_closest_cond_frames(\n frame_idx, cond_outputs, self.max_cond_frames_in_attn\n )\n t_pos_and_prevs = [(0, out) for out in selected_cond_outputs.values()]\n # Add last (self.num_maskmem - 1) frames before current frame for non-conditioning memory\n # the earliest one has t_pos=1 and the latest one has t_pos=self.num_maskmem-1\n # We also allow taking the memory frame non-consecutively (with stride>1), in which case\n # we take (self.num_maskmem - 2) frames among every stride-th frames plus the last frame.\n stride = 1 if self.training else self.memory_temporal_stride_for_eval\n for t_pos in range(1, self.num_maskmem):\n t_rel = self.num_maskmem - t_pos # how many frames before current frame\n if t_rel == 1:\n # for t_rel == 1, we take the last frame (regardless of r)\n if not track_in_reverse:\n # the frame immediately before this frame (i.e. frame_idx - 1)\n prev_frame_idx = frame_idx - t_rel\n else:\n # the frame immediately after this frame (i.e. frame_idx + 1)\n prev_frame_idx = frame_idx + t_rel\n else:\n # for t_rel >= 2, we take the memory frame from every r-th frames\n if not track_in_reverse:\n # first find the nearest frame among every r-th frames before this frame\n # for r=1, this would be (frame_idx - 2)\n prev_frame_idx = ((frame_idx - 2) // stride) * stride\n # then seek further among every r-th frames\n prev_frame_idx = prev_frame_idx - (t_rel - 2) * stride\n else:\n # first find the nearest frame among every r-th frames after this frame\n # for r=1, this would be (frame_idx + 2)\n prev_frame_idx = -(-(frame_idx + 2) // stride) * stride\n # then seek further among every r-th frames\n prev_frame_idx = prev_frame_idx + (t_rel - 2) * stride\n out = output_dict[\"non_cond_frame_outputs\"].get(prev_frame_idx, None)\n if out is None:\n # If an unselected conditioning frame is among the last (self.num_maskmem - 1)\n # frames, we still attend to it as if it's a non-conditioning frame.\n out = unselected_cond_outputs.get(prev_frame_idx, None)\n t_pos_and_prevs.append((t_pos, out))\n\n for t_pos, prev in t_pos_and_prevs:\n if prev is None:\n continue # skip padding frames\n # \"maskmem_features\" might have been offloaded to CPU in demo use cases,\n # so we load it back to GPU (it's a no-op if it's already on GPU).\n feats = prev[\"maskmem_features\"].to(device, non_blocking=True)\n to_cat_memory.append(feats.flatten(2).permute(2, 0, 1))\n # Spatial positional encoding (it might have been offloaded to CPU in eval)\n maskmem_enc = prev[\"maskmem_pos_enc\"][-1].to(device)\n maskmem_enc = maskmem_enc.flatten(2).permute(2, 0, 1)\n # Temporal positional encoding\n maskmem_enc = (\n maskmem_enc + self.maskmem_tpos_enc[self.num_maskmem - t_pos - 1]\n )\n to_cat_memory_pos_embed.append(maskmem_enc)\n\n # Construct the list of past object pointers\n if self.use_obj_ptrs_in_encoder:\n max_obj_ptrs_in_encoder = min(num_frames, self.max_obj_ptrs_in_encoder)\n # First add those object pointers from selected conditioning frames\n # (optionally, only include object pointers in the past during evaluation)\n if not self.training and self.only_obj_ptrs_in_the_past_for_eval:\n ptr_cond_outputs = {\n t: out\n for t, out in selected_cond_outputs.items()\n if (t >= frame_idx if track_in_reverse else t <= frame_idx)\n }\n else:\n ptr_cond_outputs = selected_cond_outputs\n pos_and_ptrs = [\n # Temporal pos encoding contains how far away each pointer is from current frame\n (\n (\n (frame_idx - t) * tpos_sign_mul\n if self.use_signed_tpos_enc_to_obj_ptrs\n else abs(frame_idx - t)\n ),\n out[\"obj_ptr\"],\n )\n for t, out in ptr_cond_outputs.items()\n ]\n # Add up to (max_obj_ptrs_in_encoder - 1) non-conditioning frames before current frame\n for t_diff in range(1, max_obj_ptrs_in_encoder):\n t = frame_idx + t_diff if track_in_reverse else frame_idx - t_diff\n if t < 0 or (num_frames is not None and t >= num_frames):\n break\n out = output_dict[\"non_cond_frame_outputs\"].get(\n t, unselected_cond_outputs.get(t, None)\n )\n if out is not None:\n pos_and_ptrs.append((t_diff, out[\"obj_ptr\"]))\n # If we have at least one object pointer, add them to the across attention\n if len(pos_and_ptrs) > 0:\n pos_list, ptrs_list = zip(*pos_and_ptrs)\n # stack object pointers along dim=0 into [ptr_seq_len, B, C] shape\n obj_ptrs = torch.stack(ptrs_list, dim=0)\n # a temporal positional embedding based on how far each object pointer is from\n # the current frame (sine embedding normalized by the max pointer num).\n if self.add_tpos_enc_to_obj_ptrs:\n t_diff_max = max_obj_ptrs_in_encoder - 1\n tpos_dim = C if self.proj_tpos_enc_in_obj_ptrs else self.mem_dim\n obj_pos = torch.tensor(pos_list).to(\n device=device, non_blocking=True\n )\n obj_pos = get_1d_sine_pe(obj_pos / t_diff_max, dim=tpos_dim)\n obj_pos = self.obj_ptr_tpos_proj(obj_pos)\n obj_pos = obj_pos.unsqueeze(1).expand(-1, B, self.mem_dim)\n else:\n obj_pos = obj_ptrs.new_zeros(len(pos_list), B, self.mem_dim)\n if self.mem_dim < C:\n # split a pointer into (C // self.mem_dim) tokens for self.mem_dim < C\n obj_ptrs = obj_ptrs.reshape(\n -1, B, C // self.mem_dim, self.mem_dim\n )\n obj_ptrs = obj_ptrs.permute(0, 2, 1, 3).flatten(0, 1)\n obj_pos = obj_pos.repeat_interleave(C // self.mem_dim, dim=0)\n to_cat_memory.append(obj_ptrs)\n to_cat_memory_pos_embed.append(obj_pos)\n num_obj_ptr_tokens = obj_ptrs.shape[0]\n else:\n num_obj_ptr_tokens = 0\n else:\n # for initial conditioning frames, encode them without using any previous memory\n if self.directly_add_no_mem_embed:\n # directly add no-mem embedding (instead of using the transformer encoder)\n pix_feat_with_mem = current_vision_feats[-1] + self.no_mem_embed\n pix_feat_with_mem = pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W)\n return pix_feat_with_mem\n\n # Use a dummy token on the first frame (to avoid empty memory input to tranformer encoder)\n to_cat_memory = [self.no_mem_embed.expand(1, B, self.mem_dim)]\n to_cat_memory_pos_embed = [self.no_mem_pos_enc.expand(1, B, self.mem_dim)]\n\n # Step 2: Concatenate the memories and forward through the transformer encoder\n memory = torch.cat(to_cat_memory, dim=0)\n memory_pos_embed = torch.cat(to_cat_memory_pos_embed, dim=0)\n\n pix_feat_with_mem = self.memory_attention(\n curr=current_vision_feats,\n curr_pos=current_vision_pos_embeds,\n memory=memory,\n memory_pos=memory_pos_embed,\n num_obj_ptr_tokens=num_obj_ptr_tokens,\n )\n # reshape the output (HW)BC => BCHW\n pix_feat_with_mem = pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W)\n return pix_feat_with_mem\n\n def _encode_new_memory(\n self,\n current_vision_feats,\n feat_sizes,\n pred_masks_high_res,\n object_score_logits,\n is_mask_from_pts,\n ):\n \"\"\"Encode the current image and its prediction into a memory feature.\"\"\"\n B = current_vision_feats[-1].size(1) # batch size on this frame\n C = self.hidden_dim\n H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size\n # top-level feature, (HW)BC => BCHW\n pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W)\n if self.non_overlap_masks_for_mem_enc and not self.training:\n # optionally, apply non-overlapping constraints to the masks (it's applied\n # in the batch dimension and should only be used during eval, where all\n # the objects come from the same video under batch size 1).\n pred_masks_high_res = self._apply_non_overlapping_constraints(\n pred_masks_high_res\n )\n # scale the raw mask logits with a temperature before applying sigmoid\n binarize = self.binarize_mask_from_pts_for_mem_enc and is_mask_from_pts\n if binarize and not self.training:\n mask_for_mem = (pred_masks_high_res > 0).float()\n else:\n # apply sigmoid on the raw mask logits to turn them into range (0, 1)\n mask_for_mem = torch.sigmoid(pred_masks_high_res)\n # apply scale and bias terms to the sigmoid probabilities\n if self.sigmoid_scale_for_mem_enc != 1.0:\n mask_for_mem = mask_for_mem * self.sigmoid_scale_for_mem_enc\n if self.sigmoid_bias_for_mem_enc != 0.0:\n mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc\n maskmem_out = self.memory_encoder(\n pix_feat, mask_for_mem, skip_mask_sigmoid=True # sigmoid already applied\n )\n maskmem_features = maskmem_out[\"vision_features\"]\n maskmem_pos_enc = maskmem_out[\"vision_pos_enc\"]\n # add a no-object embedding to the spatial memory to indicate that the frame\n # is predicted to be occluded (i.e. no object is appearing in the frame)\n if self.no_obj_embed_spatial is not None:\n is_obj_appearing = (object_score_logits > 0).float()\n maskmem_features += (\n 1 - is_obj_appearing[..., None, None]\n ) * self.no_obj_embed_spatial[..., None, None].expand(\n *maskmem_features.shape\n )\n\n return maskmem_features, maskmem_pos_enc\n\n def _track_step(\n self,\n frame_idx,\n is_init_cond_frame,\n current_vision_feats,\n current_vision_pos_embeds,\n feat_sizes,\n point_inputs,\n mask_inputs,\n output_dict,\n num_frames,\n track_in_reverse,\n prev_sam_mask_logits,\n ):\n current_out = {\"point_inputs\": point_inputs, \"mask_inputs\": mask_inputs}\n # High-resolution feature maps for the SAM head, reshape (HW)BC => BCHW\n if len(current_vision_feats) > 1:\n high_res_features = [\n x.permute(1, 2, 0).view(x.size(1), x.size(2), *s)\n for x, s in zip(current_vision_feats[:-1], feat_sizes[:-1])\n ]\n else:\n high_res_features = None\n if mask_inputs is not None and self.use_mask_input_as_output_without_sam:\n # When use_mask_input_as_output_without_sam=True, we directly output the mask input\n # (see it as a GT mask) without using a SAM prompt encoder + mask decoder.\n pix_feat = current_vision_feats[-1].permute(1, 2, 0)\n pix_feat = pix_feat.view(-1, self.hidden_dim, *feat_sizes[-1])\n sam_outputs = self._use_mask_as_output(\n pix_feat, high_res_features, mask_inputs\n )\n else:\n # fused the visual feature with previous memory features in the memory bank\n pix_feat = self._prepare_memory_conditioned_features(\n frame_idx=frame_idx,\n is_init_cond_frame=is_init_cond_frame,\n current_vision_feats=current_vision_feats[-1:],\n current_vision_pos_embeds=current_vision_pos_embeds[-1:],\n feat_sizes=feat_sizes[-1:],\n output_dict=output_dict,\n num_frames=num_frames,\n track_in_reverse=track_in_reverse,\n )\n # apply SAM-style segmentation head\n # here we might feed previously predicted low-res SAM mask logits into the SAM mask decoder,\n # e.g. in demo where such logits come from earlier interaction instead of correction sampling\n # (in this case, any `mask_inputs` shouldn't reach here as they are sent to _use_mask_as_output instead)\n if prev_sam_mask_logits is not None:\n assert point_inputs is not None and mask_inputs is None\n mask_inputs = prev_sam_mask_logits\n multimask_output = self._use_multimask(is_init_cond_frame, point_inputs)\n sam_outputs = self._forward_sam_heads(\n backbone_features=pix_feat,\n point_inputs=point_inputs,\n mask_inputs=mask_inputs,\n high_res_features=high_res_features,\n multimask_output=multimask_output,\n )\n\n return current_out, sam_outputs, high_res_features, pix_feat\n\n def _encode_memory_in_output(\n self,\n current_vision_feats,\n feat_sizes,\n point_inputs,\n run_mem_encoder,\n high_res_masks,\n object_score_logits,\n current_out,\n ):\n if run_mem_encoder and self.num_maskmem > 0:\n high_res_masks_for_mem_enc = high_res_masks\n maskmem_features, maskmem_pos_enc = self._encode_new_memory(\n current_vision_feats=current_vision_feats,\n feat_sizes=feat_sizes,\n pred_masks_high_res=high_res_masks_for_mem_enc,\n object_score_logits=object_score_logits,\n is_mask_from_pts=(point_inputs is not None),\n )\n current_out[\"maskmem_features\"] = maskmem_features\n current_out[\"maskmem_pos_enc\"] = maskmem_pos_enc\n else:\n current_out[\"maskmem_features\"] = None\n current_out[\"maskmem_pos_enc\"] = None\n\n def track_step(\n self,\n frame_idx,\n is_init_cond_frame,\n current_vision_feats,\n current_vision_pos_embeds,\n feat_sizes,\n point_inputs,\n mask_inputs,\n output_dict,\n num_frames,\n track_in_reverse=False, # tracking in reverse time order (for demo usage)\n # Whether to run the memory encoder on the predicted masks. Sometimes we might want\n # to skip the memory encoder with `run_mem_encoder=False`. For example,\n # in demo we might call `track_step` multiple times for each user click,\n # and only encode the memory when the user finalizes their clicks. And in ablation\n # settings like SAM training on static images, we don't need the memory encoder.\n run_mem_encoder=True,\n # The previously predicted SAM mask logits (which can be fed together with new clicks in demo).\n prev_sam_mask_logits=None,\n ):\n current_out, sam_outputs, _, _ = self._track_step(\n frame_idx,\n is_init_cond_frame,\n current_vision_feats,\n current_vision_pos_embeds,\n feat_sizes,\n point_inputs,\n mask_inputs,\n output_dict,\n num_frames,\n track_in_reverse,\n prev_sam_mask_logits,\n )\n\n (\n _,\n _,\n _,\n low_res_masks,\n high_res_masks,\n obj_ptr,\n object_score_logits,\n ) = sam_outputs\n\n current_out[\"pred_masks\"] = low_res_masks\n current_out[\"pred_masks_high_res\"] = high_res_masks\n current_out[\"obj_ptr\"] = obj_ptr\n if not self.training:\n # Only add this in inference (to avoid unused param in activation checkpointing;\n # it's mainly used in the demo to encode spatial memories w/ consolidated masks)\n current_out[\"object_score_logits\"] = object_score_logits\n\n # Finally run the memory encoder on the predicted mask to encode\n # it into a new memory feature (that can be used in future frames)\n self._encode_memory_in_output(\n current_vision_feats,\n feat_sizes,\n point_inputs,\n run_mem_encoder,\n high_res_masks,\n object_score_logits,\n current_out,\n )\n\n return current_out\n\n def _use_multimask(self, is_init_cond_frame, point_inputs):\n \"\"\"Whether to use multimask output in the SAM head.\"\"\"\n num_pts = 0 if point_inputs is None else point_inputs[\"point_labels\"].size(1)\n multimask_output = (\n self.multimask_output_in_sam\n and (is_init_cond_frame or self.multimask_output_for_tracking)\n and (self.multimask_min_pt_num <= num_pts <= self.multimask_max_pt_num)\n )\n return multimask_output\n\n def _apply_non_overlapping_constraints(self, pred_masks):\n \"\"\"\n Apply non-overlapping constraints to the object scores in pred_masks. Here we\n keep only the highest scoring object at each spatial location in pred_masks.\n \"\"\"\n batch_size = pred_masks.size(0)\n if batch_size == 1:\n return pred_masks\n\n device = pred_masks.device\n # \"max_obj_inds\": object index of the object with the highest score at each location\n max_obj_inds = torch.argmax(pred_masks, dim=0, keepdim=True)\n # \"batch_obj_inds\": object index of each object slice (along dim 0) in `pred_masks`\n batch_obj_inds = torch.arange(batch_size, device=device)[:, None, None, None]\n keep = max_obj_inds == batch_obj_inds\n # suppress overlapping regions' scores below -10.0 so that the foreground regions\n # don't overlap (here sigmoid(-10.0)=4.5398e-05)\n pred_masks = torch.where(keep, pred_masks, torch.clamp(pred_masks, max=-10.0))\n return pred_masks\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/modeling/sam2_utils.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\n\nimport copy\nfrom typing import Tuple\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom sam2.utils.misc import mask_to_box\n\n\ndef select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num):\n \"\"\"\n Select up to `max_cond_frame_num` conditioning frames from `cond_frame_outputs`\n that are temporally closest to the current frame at `frame_idx`. Here, we take\n - a) the closest conditioning frame before `frame_idx` (if any);\n - b) the closest conditioning frame after `frame_idx` (if any);\n - c) any other temporally closest conditioning frames until reaching a total\n of `max_cond_frame_num` conditioning frames.\n\n Outputs:\n - selected_outputs: selected items (keys & values) from `cond_frame_outputs`.\n - unselected_outputs: items (keys & values) not selected in `cond_frame_outputs`.\n \"\"\"\n if max_cond_frame_num == -1 or len(cond_frame_outputs) <= max_cond_frame_num:\n selected_outputs = cond_frame_outputs\n unselected_outputs = {}\n else:\n assert max_cond_frame_num >= 2, \"we should allow using 2+ conditioning frames\"\n selected_outputs = {}\n\n # the closest conditioning frame before `frame_idx` (if any)\n idx_before = max((t for t in cond_frame_outputs if t < frame_idx), default=None)\n if idx_before is not None:\n selected_outputs[idx_before] = cond_frame_outputs[idx_before]\n\n # the closest conditioning frame after `frame_idx` (if any)\n idx_after = min((t for t in cond_frame_outputs if t >= frame_idx), default=None)\n if idx_after is not None:\n selected_outputs[idx_after] = cond_frame_outputs[idx_after]\n\n # add other temporally closest conditioning frames until reaching a total\n # of `max_cond_frame_num` conditioning frames.\n num_remain = max_cond_frame_num - len(selected_outputs)\n inds_remain = sorted(\n (t for t in cond_frame_outputs if t not in selected_outputs),\n key=lambda x: abs(x - frame_idx),\n )[:num_remain]\n selected_outputs.update((t, cond_frame_outputs[t]) for t in inds_remain)\n unselected_outputs = {\n t: v for t, v in cond_frame_outputs.items() if t not in selected_outputs\n }\n\n return selected_outputs, unselected_outputs\n\n\ndef get_1d_sine_pe(pos_inds, dim, temperature=10000):\n \"\"\"\n Get 1D sine positional embedding as in the original Transformer paper.\n \"\"\"\n pe_dim = dim // 2\n dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device)\n dim_t = temperature ** (2 * (dim_t // 2) / pe_dim)\n\n pos_embed = pos_inds.unsqueeze(-1) / dim_t\n pos_embed = torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1)\n return pos_embed\n\n\ndef get_activation_fn(activation):\n \"\"\"Return an activation function given a string\"\"\"\n if activation == \"relu\":\n return F.relu\n if activation == \"gelu\":\n return F.gelu\n if activation == \"glu\":\n return F.glu\n raise RuntimeError(f\"activation should be relu/gelu, not {activation}.\")\n\n\ndef get_clones(module, N):\n return nn.ModuleList([copy.deepcopy(module) for i in range(N)])\n\n\nclass DropPath(nn.Module):\n # adapted from https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py\n def __init__(self, drop_prob=0.0, scale_by_keep=True):\n super(DropPath, self).__init__()\n self.drop_prob = drop_prob\n self.scale_by_keep = scale_by_keep\n\n def forward(self, x):\n if self.drop_prob == 0.0 or not self.training:\n return x\n keep_prob = 1 - self.drop_prob\n shape = (x.shape[0],) + (1,) * (x.ndim - 1)\n random_tensor = x.new_empty(shape).bernoulli_(keep_prob)\n if keep_prob > 0.0 and self.scale_by_keep:\n random_tensor.div_(keep_prob)\n return x * random_tensor\n\n\n# Lightly adapted from\n# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa\nclass MLP(nn.Module):\n def __init__(\n self,\n input_dim: int,\n hidden_dim: int,\n output_dim: int,\n num_layers: int,\n activation: nn.Module = nn.ReLU,\n sigmoid_output: bool = False,\n ) -> None:\n super().__init__()\n self.num_layers = num_layers\n h = [hidden_dim] * (num_layers - 1)\n self.layers = nn.ModuleList(\n nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])\n )\n self.sigmoid_output = sigmoid_output\n self.act = activation()\n\n def forward(self, x):\n for i, layer in enumerate(self.layers):\n x = self.act(layer(x)) if i < self.num_layers - 1 else layer(x)\n if self.sigmoid_output:\n x = F.sigmoid(x)\n return x\n\n\n# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa\n# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa\nclass LayerNorm2d(nn.Module):\n def __init__(self, num_channels: int, eps: float = 1e-6) -> None:\n super().__init__()\n self.weight = nn.Parameter(torch.ones(num_channels))\n self.bias = nn.Parameter(torch.zeros(num_channels))\n self.eps = eps\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n u = x.mean(1, keepdim=True)\n s = (x - u).pow(2).mean(1, keepdim=True)\n x = (x - u) / torch.sqrt(s + self.eps)\n x = self.weight[:, None, None] * x + self.bias[:, None, None]\n return x\n\n\ndef sample_box_points(\n masks: torch.Tensor,\n noise: float = 0.1, # SAM default\n noise_bound: int = 20, # SAM default\n top_left_label: int = 2,\n bottom_right_label: int = 3,\n) -> Tuple[np.array, np.array]:\n \"\"\"\n Sample a noised version of the top left and bottom right corners of a given `bbox`\n\n Inputs:\n - masks: [B, 1, H,W] boxes, dtype=torch.Tensor\n - noise: noise as a fraction of box width and height, dtype=float\n - noise_bound: maximum amount of noise (in pure pixesl), dtype=int\n\n Returns:\n - box_coords: [B, num_pt, 2], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.float\n - box_labels: [B, num_pt], label 2 is reserverd for top left and 3 for bottom right corners, dtype=torch.int32\n \"\"\"\n device = masks.device\n box_coords = mask_to_box(masks)\n B, _, H, W = masks.shape\n box_labels = torch.tensor(\n [top_left_label, bottom_right_label], dtype=torch.int, device=device\n ).repeat(B)\n if noise > 0.0:\n if not isinstance(noise_bound, torch.Tensor):\n noise_bound = torch.tensor(noise_bound, device=device)\n bbox_w = box_coords[..., 2] - box_coords[..., 0]\n bbox_h = box_coords[..., 3] - box_coords[..., 1]\n max_dx = torch.min(bbox_w * noise, noise_bound)\n max_dy = torch.min(bbox_h * noise, noise_bound)\n box_noise = 2 * torch.rand(B, 1, 4, device=device) - 1\n box_noise = box_noise * torch.stack((max_dx, max_dy, max_dx, max_dy), dim=-1)\n\n box_coords = box_coords + box_noise\n img_bounds = (\n torch.tensor([W, H, W, H], device=device) - 1\n ) # uncentered pixel coords\n box_coords.clamp_(torch.zeros_like(img_bounds), img_bounds) # In place clamping\n\n box_coords = box_coords.reshape(-1, 2, 2) # always 2 points\n box_labels = box_labels.reshape(-1, 2)\n return box_coords, box_labels\n\n\ndef sample_random_points_from_errors(gt_masks, pred_masks, num_pt=1):\n \"\"\"\n Sample `num_pt` random points (along with their labels) independently from the error regions.\n\n Inputs:\n - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool\n - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None\n - num_pt: int, number of points to sample independently for each of the B error maps\n\n Outputs:\n - points: [B, num_pt, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point\n - labels: [B, num_pt], dtype=torch.int32, where 1 means positive clicks and 0 means\n negative clicks\n \"\"\"\n if pred_masks is None: # if pred_masks is not provided, treat it as empty\n pred_masks = torch.zeros_like(gt_masks)\n assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1\n assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape\n assert num_pt >= 0\n\n B, _, H_im, W_im = gt_masks.shape\n device = gt_masks.device\n\n # false positive region, a new point sampled in this region should have\n # negative label to correct the FP error\n fp_masks = ~gt_masks & pred_masks\n # false negative region, a new point sampled in this region should have\n # positive label to correct the FN error\n fn_masks = gt_masks & ~pred_masks\n # whether the prediction completely match the ground-truth on each mask\n all_correct = torch.all((gt_masks == pred_masks).flatten(2), dim=2)\n all_correct = all_correct[..., None, None]\n\n # channel 0 is FP map, while channel 1 is FN map\n pts_noise = torch.rand(B, num_pt, H_im, W_im, 2, device=device)\n # sample a negative new click from FP region or a positive new click\n # from FN region, depend on where the maximum falls,\n # and in case the predictions are all correct (no FP or FN), we just\n # sample a negative click from the background region\n pts_noise[..., 0] *= fp_masks | (all_correct & ~gt_masks)\n pts_noise[..., 1] *= fn_masks\n pts_idx = pts_noise.flatten(2).argmax(dim=2)\n labels = (pts_idx % 2).to(torch.int32)\n pts_idx = pts_idx // 2\n pts_x = pts_idx % W_im\n pts_y = pts_idx // W_im\n points = torch.stack([pts_x, pts_y], dim=2).to(torch.float)\n return points, labels\n\n\ndef sample_one_point_from_error_center(gt_masks, pred_masks, padding=True):\n \"\"\"\n Sample 1 random point (along with its label) from the center of each error region,\n that is, the point with the largest distance to the boundary of each error region.\n This is the RITM sampling method from https://github.com/saic-vul/ritm_interactive_segmentation/blob/master/isegm/inference/clicker.py\n\n Inputs:\n - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool\n - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None\n - padding: if True, pad with boundary of 1 px for distance transform\n\n Outputs:\n - points: [B, 1, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point\n - labels: [B, 1], dtype=torch.int32, where 1 means positive clicks and 0 means negative clicks\n \"\"\"\n import cv2\n\n if pred_masks is None:\n pred_masks = torch.zeros_like(gt_masks)\n assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1\n assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape\n\n B, _, _, W_im = gt_masks.shape\n device = gt_masks.device\n\n # false positive region, a new point sampled in this region should have\n # negative label to correct the FP error\n fp_masks = ~gt_masks & pred_masks\n # false negative region, a new point sampled in this region should have\n # positive label to correct the FN error\n fn_masks = gt_masks & ~pred_masks\n\n fp_masks = fp_masks.cpu().numpy()\n fn_masks = fn_masks.cpu().numpy()\n points = torch.zeros(B, 1, 2, dtype=torch.float)\n labels = torch.ones(B, 1, dtype=torch.int32)\n for b in range(B):\n fn_mask = fn_masks[b, 0]\n fp_mask = fp_masks[b, 0]\n if padding:\n fn_mask = np.pad(fn_mask, ((1, 1), (1, 1)), \"constant\")\n fp_mask = np.pad(fp_mask, ((1, 1), (1, 1)), \"constant\")\n # compute the distance of each point in FN/FP region to its boundary\n fn_mask_dt = cv2.distanceTransform(fn_mask.astype(np.uint8), cv2.DIST_L2, 0)\n fp_mask_dt = cv2.distanceTransform(fp_mask.astype(np.uint8), cv2.DIST_L2, 0)\n if padding:\n fn_mask_dt = fn_mask_dt[1:-1, 1:-1]\n fp_mask_dt = fp_mask_dt[1:-1, 1:-1]\n\n # take the point in FN/FP region with the largest distance to its boundary\n fn_mask_dt_flat = fn_mask_dt.reshape(-1)\n fp_mask_dt_flat = fp_mask_dt.reshape(-1)\n fn_argmax = np.argmax(fn_mask_dt_flat)\n fp_argmax = np.argmax(fp_mask_dt_flat)\n is_positive = fn_mask_dt_flat[fn_argmax] > fp_mask_dt_flat[fp_argmax]\n pt_idx = fn_argmax if is_positive else fp_argmax\n points[b, 0, 0] = pt_idx % W_im # x\n points[b, 0, 1] = pt_idx // W_im # y\n labels[b, 0] = int(is_positive)\n\n points = points.to(device)\n labels = labels.to(device)\n return points, labels\n\n\ndef get_next_point(gt_masks, pred_masks, method):\n if method == \"uniform\":\n return sample_random_points_from_errors(gt_masks, pred_masks)\n elif method == \"center\":\n return sample_one_point_from_error_center(gt_masks, pred_masks)\n else:\n raise ValueError(f\"unknown sampling method {method}\")\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/sam2_hiera_b+.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 112\n num_heads: 2\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [896, 448, 224, 112]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/sam2_hiera_l.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 144\n num_heads: 2\n stages: [2, 6, 36, 4]\n global_att_blocks: [23, 33, 43]\n window_pos_embed_bkg_spatial_size: [7, 7]\n window_spec: [8, 4, 16, 8]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [1152, 576, 288, 144]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/sam2_hiera_s.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 96\n num_heads: 1\n stages: [1, 2, 11, 2]\n global_att_blocks: [7, 10, 13]\n window_pos_embed_bkg_spatial_size: [7, 7]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [768, 384, 192, 96]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/sam2_hiera_t.yaml", "content": "# @package _global_\n\n# Model\nmodel:\n _target_: sam2.modeling.sam2_base.SAM2Base\n image_encoder:\n _target_: sam2.modeling.backbones.image_encoder.ImageEncoder\n scalp: 1\n trunk:\n _target_: sam2.modeling.backbones.hieradet.Hiera\n embed_dim: 96\n num_heads: 1\n stages: [1, 2, 7, 2]\n global_att_blocks: [5, 7, 9]\n window_pos_embed_bkg_spatial_size: [7, 7]\n neck:\n _target_: sam2.modeling.backbones.image_encoder.FpnNeck\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 256\n normalize: true\n scale: null\n temperature: 10000\n d_model: 256\n backbone_channel_list: [768, 384, 192, 96]\n fpn_top_down_levels: [2, 3] # output level 0 and 1 directly use the backbone features\n fpn_interp_model: nearest\n\n memory_attention:\n _target_: sam2.modeling.memory_attention.MemoryAttention\n d_model: 256\n pos_enc_at_input: true\n layer:\n _target_: sam2.modeling.memory_attention.MemoryAttentionLayer\n activation: relu\n dim_feedforward: 2048\n dropout: 0.1\n pos_enc_at_attn: false\n self_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n d_model: 256\n pos_enc_at_cross_attn_keys: true\n pos_enc_at_cross_attn_queries: false\n cross_attention:\n _target_: sam2.modeling.sam.transformer.RoPEAttention\n rope_theta: 10000.0\n feat_sizes: [64, 64]\n rope_k_repeat: True\n embedding_dim: 256\n num_heads: 1\n downsample_rate: 1\n dropout: 0.1\n kv_in_dim: 64\n num_layers: 4\n\n memory_encoder:\n _target_: sam2.modeling.memory_encoder.MemoryEncoder\n out_dim: 64\n position_encoding:\n _target_: sam2.modeling.position_encoding.PositionEmbeddingSine\n num_pos_feats: 64\n normalize: true\n scale: null\n temperature: 10000\n mask_downsampler:\n _target_: sam2.modeling.memory_encoder.MaskDownSampler\n kernel_size: 3\n stride: 2\n padding: 1\n fuser:\n _target_: sam2.modeling.memory_encoder.Fuser\n layer:\n _target_: sam2.modeling.memory_encoder.CXBlock\n dim: 256\n kernel_size: 7\n padding: 3\n layer_scale_init_value: 1e-6\n use_dwconv: True # depth-wise convs\n num_layers: 2\n\n num_maskmem: 7\n image_size: 1024\n # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask\n # SAM decoder\n sigmoid_scale_for_mem_enc: 20.0\n sigmoid_bias_for_mem_enc: -10.0\n use_mask_input_as_output_without_sam: true\n # Memory\n directly_add_no_mem_embed: true\n # use high-resolution feature map in the SAM mask decoder\n use_high_res_features_in_sam: true\n # output 3 masks on the first click on initial conditioning frames\n multimask_output_in_sam: true\n # SAM heads\n iou_prediction_use_sigmoid: True\n # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder\n use_obj_ptrs_in_encoder: true\n add_tpos_enc_to_obj_ptrs: false\n only_obj_ptrs_in_the_past_for_eval: true\n # object occlusion prediction\n pred_obj_scores: true\n pred_obj_scores_mlp: true\n fixed_no_obj_ptr: true\n # multimask tracking settings\n multimask_output_for_tracking: true\n use_multimask_token_for_obj_ptr: true\n multimask_min_pt_num: 0\n multimask_max_pt_num: 1\n use_mlp_for_obj_ptr_proj: true\n # Compilation flag\n # HieraT does not currently support compilation, should always be set to False\n compile_image_encoder: False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/sam2_image_predictor.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\n\nfrom typing import List, Optional, Tuple, Union\n\nimport numpy as np\nimport torch\nfrom PIL.Image import Image\n\nfrom sam2.modeling.sam2_base import SAM2Base\n\nfrom sam2.utils.transforms import SAM2Transforms\n\n\nclass SAM2ImagePredictor:\n def __init__(\n self,\n sam_model: SAM2Base,\n mask_threshold=0.0,\n max_hole_area=0.0,\n max_sprinkle_area=0.0,\n **kwargs,\n ) -> None:\n \"\"\"\n Uses SAM-2 to calculate the image embedding for an image, and then\n allow repeated, efficient mask prediction given prompts.\n\n Arguments:\n sam_model (Sam-2): The model to use for mask prediction.\n mask_threshold (float): The threshold to use when converting mask logits\n to binary masks. Masks are thresholded at 0 by default.\n max_hole_area (int): If max_hole_area > 0, we fill small holes in up to\n the maximum area of max_hole_area in low_res_masks.\n max_sprinkle_area (int): If max_sprinkle_area > 0, we remove small sprinkles up to\n the maximum area of max_sprinkle_area in low_res_masks.\n \"\"\"\n super().__init__()\n self.model = sam_model\n self._transforms = SAM2Transforms(\n resolution=self.model.image_size,\n mask_threshold=mask_threshold,\n max_hole_area=max_hole_area,\n max_sprinkle_area=max_sprinkle_area,\n )\n\n # Predictor state\n self._is_image_set = False\n self._features = None\n self._orig_hw = None\n # Whether the predictor is set for single image or a batch of images\n self._is_batch = False\n\n # Predictor config\n self.mask_threshold = mask_threshold\n\n # Spatial dim for backbone feature maps\n self._bb_feat_sizes = [\n (256, 256),\n (128, 128),\n (64, 64),\n ]\n\n @classmethod\n def from_pretrained(cls, model_id: str, **kwargs) -> \"SAM2ImagePredictor\":\n \"\"\"\n Load a pretrained model from the Hugging Face hub.\n\n Arguments:\n model_id (str): The Hugging Face repository ID.\n **kwargs: Additional arguments to pass to the model constructor.\n\n Returns:\n (SAM2ImagePredictor): The loaded model.\n \"\"\"\n from sam2.build_sam import build_sam2_hf\n\n sam_model = build_sam2_hf(model_id, **kwargs)\n return cls(sam_model, **kwargs)\n\n @torch.no_grad()\n def set_image(\n self,\n image: Union[np.ndarray, Image],\n ) -> None:\n \"\"\"\n Calculates the image embeddings for the provided image, allowing\n masks to be predicted with the 'predict' method.\n\n Arguments:\n image (np.ndarray or PIL Image): The input image to embed in RGB format. The image should be in HWC format if np.ndarray, or WHC format if PIL Image\n with pixel values in [0, 255].\n image_format (str): The color format of the image, in ['RGB', 'BGR'].\n \"\"\"\n self.reset_predictor()\n # Transform the image to the form expected by the model\n if isinstance(image, np.ndarray):\n logging.info(\"For numpy array image, we assume (HxWxC) format\")\n self._orig_hw = [image.shape[:2]]\n elif isinstance(image, Image):\n w, h = image.size\n self._orig_hw = [(h, w)]\n else:\n raise NotImplementedError(\"Image format not supported\")\n\n input_image = self._transforms(image)\n input_image = input_image[None, ...].to(self.device)\n\n assert (\n len(input_image.shape) == 4 and input_image.shape[1] == 3\n ), f\"input_image must be of size 1x3xHxW, got {input_image.shape}\"\n logging.info(\"Computing image embeddings for the provided image...\")\n backbone_out = self.model.forward_image(input_image)\n _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out)\n # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos\n if self.model.directly_add_no_mem_embed:\n vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed\n\n feats = [\n feat.permute(1, 2, 0).view(1, -1, *feat_size)\n for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])\n ][::-1]\n self._features = {\"image_embed\": feats[-1], \"high_res_feats\": feats[:-1]}\n self._is_image_set = True\n logging.info(\"Image embeddings computed.\")\n\n @torch.no_grad()\n def set_image_batch(\n self,\n image_list: List[Union[np.ndarray]],\n ) -> None:\n \"\"\"\n Calculates the image embeddings for the provided image batch, allowing\n masks to be predicted with the 'predict_batch' method.\n\n Arguments:\n image_list (List[np.ndarray]): The input images to embed in RGB format. The image should be in HWC format if np.ndarray\n with pixel values in [0, 255].\n \"\"\"\n self.reset_predictor()\n assert isinstance(image_list, list)\n self._orig_hw = []\n for image in image_list:\n assert isinstance(\n image, np.ndarray\n ), \"Images are expected to be an np.ndarray in RGB format, and of shape HWC\"\n self._orig_hw.append(image.shape[:2])\n # Transform the image to the form expected by the model\n img_batch = self._transforms.forward_batch(image_list)\n img_batch = img_batch.to(self.device)\n batch_size = img_batch.shape[0]\n assert (\n len(img_batch.shape) == 4 and img_batch.shape[1] == 3\n ), f\"img_batch must be of size Bx3xHxW, got {img_batch.shape}\"\n logging.info(\"Computing image embeddings for the provided images...\")\n backbone_out = self.model.forward_image(img_batch)\n _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out)\n # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos\n if self.model.directly_add_no_mem_embed:\n vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed\n\n feats = [\n feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)\n for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])\n ][::-1]\n self._features = {\"image_embed\": feats[-1], \"high_res_feats\": feats[:-1]}\n self._is_image_set = True\n self._is_batch = True\n logging.info(\"Image embeddings computed.\")\n\n def predict_batch(\n self,\n point_coords_batch: List[np.ndarray] = None,\n point_labels_batch: List[np.ndarray] = None,\n box_batch: List[np.ndarray] = None,\n mask_input_batch: List[np.ndarray] = None,\n multimask_output: bool = True,\n return_logits: bool = False,\n normalize_coords=True,\n ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]:\n \"\"\"This function is very similar to predict(...), however it is used for batched mode, when the model is expected to generate predictions on multiple images.\n It returns a tuple of lists of masks, ious, and low_res_masks_logits.\n \"\"\"\n assert self._is_batch, \"This function should only be used when in batched mode\"\n if not self._is_image_set:\n raise RuntimeError(\n \"An image must be set with .set_image_batch(...) before mask prediction.\"\n )\n num_images = len(self._features[\"image_embed\"])\n all_masks = []\n all_ious = []\n all_low_res_masks = []\n for img_idx in range(num_images):\n # Transform input prompts\n point_coords = (\n point_coords_batch[img_idx] if point_coords_batch is not None else None\n )\n point_labels = (\n point_labels_batch[img_idx] if point_labels_batch is not None else None\n )\n box = box_batch[img_idx] if box_batch is not None else None\n mask_input = (\n mask_input_batch[img_idx] if mask_input_batch is not None else None\n )\n mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(\n point_coords,\n point_labels,\n box,\n mask_input,\n normalize_coords,\n img_idx=img_idx,\n )\n masks, iou_predictions, low_res_masks = self._predict(\n unnorm_coords,\n labels,\n unnorm_box,\n mask_input,\n multimask_output,\n return_logits=return_logits,\n img_idx=img_idx,\n )\n masks_np = masks.squeeze(0).float().detach().cpu().numpy()\n iou_predictions_np = (\n iou_predictions.squeeze(0).float().detach().cpu().numpy()\n )\n low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()\n all_masks.append(masks_np)\n all_ious.append(iou_predictions_np)\n all_low_res_masks.append(low_res_masks_np)\n\n return all_masks, all_ious, all_low_res_masks\n\n def predict(\n self,\n point_coords: Optional[np.ndarray] = None,\n point_labels: Optional[np.ndarray] = None,\n box: Optional[np.ndarray] = None,\n mask_input: Optional[np.ndarray] = None,\n multimask_output: bool = True,\n return_logits: bool = False,\n normalize_coords=True,\n ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n \"\"\"\n Predict masks for the given input prompts, using the currently set image.\n\n Arguments:\n point_coords (np.ndarray or None): A Nx2 array of point prompts to the\n model. Each point is in (X,Y) in pixels.\n point_labels (np.ndarray or None): A length N array of labels for the\n point prompts. 1 indicates a foreground point and 0 indicates a\n background point.\n box (np.ndarray or None): A length 4 array given a box prompt to the\n model, in XYXY format.\n mask_input (np.ndarray): A low resolution mask input to the model, typically\n coming from a previous prediction iteration. Has form 1xHxW, where\n for SAM, H=W=256.\n multimask_output (bool): If true, the model will return three masks.\n For ambiguous input prompts (such as a single click), this will often\n produce better masks than a single prediction. If only a single\n mask is needed, the model's predicted quality score can be used\n to select the best mask. For non-ambiguous prompts, such as multiple\n input prompts, multimask_output=False can give better results.\n return_logits (bool): If true, returns un-thresholded masks logits\n instead of a binary mask.\n normalize_coords (bool): If true, the point coordinates will be normalized to the range [0,1] and point_coords is expected to be wrt. image dimensions.\n\n Returns:\n (np.ndarray): The output masks in CxHxW format, where C is the\n number of masks, and (H, W) is the original image size.\n (np.ndarray): An array of length C containing the model's\n predictions for the quality of each mask.\n (np.ndarray): An array of shape CxHxW, where C is the number\n of masks and H=W=256. These low resolution logits can be passed to\n a subsequent iteration as mask input.\n \"\"\"\n if not self._is_image_set:\n raise RuntimeError(\n \"An image must be set with .set_image(...) before mask prediction.\"\n )\n\n # Transform input prompts\n\n mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(\n point_coords, point_labels, box, mask_input, normalize_coords\n )\n\n masks, iou_predictions, low_res_masks = self._predict(\n unnorm_coords,\n labels,\n unnorm_box,\n mask_input,\n multimask_output,\n return_logits=return_logits,\n )\n\n masks_np = masks.squeeze(0).float().detach().cpu().numpy()\n iou_predictions_np = iou_predictions.squeeze(0).float().detach().cpu().numpy()\n low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()\n return masks_np, iou_predictions_np, low_res_masks_np\n\n def _prep_prompts(\n self, point_coords, point_labels, box, mask_logits, normalize_coords, img_idx=-1\n ):\n\n unnorm_coords, labels, unnorm_box, mask_input = None, None, None, None\n if point_coords is not None:\n assert (\n point_labels is not None\n ), \"point_labels must be supplied if point_coords is supplied.\"\n point_coords = torch.as_tensor(\n point_coords, dtype=torch.float, device=self.device\n )\n unnorm_coords = self._transforms.transform_coords(\n point_coords, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]\n )\n labels = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)\n if len(unnorm_coords.shape) == 2:\n unnorm_coords, labels = unnorm_coords[None, ...], labels[None, ...]\n if box is not None:\n box = torch.as_tensor(box, dtype=torch.float, device=self.device)\n unnorm_box = self._transforms.transform_boxes(\n box, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]\n ) # Bx2x2\n if mask_logits is not None:\n mask_input = torch.as_tensor(\n mask_logits, dtype=torch.float, device=self.device\n )\n if len(mask_input.shape) == 3:\n mask_input = mask_input[None, :, :, :]\n return mask_input, unnorm_coords, labels, unnorm_box\n\n @torch.no_grad()\n def _predict(\n self,\n point_coords: Optional[torch.Tensor],\n point_labels: Optional[torch.Tensor],\n boxes: Optional[torch.Tensor] = None,\n mask_input: Optional[torch.Tensor] = None,\n multimask_output: bool = True,\n return_logits: bool = False,\n img_idx: int = -1,\n ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n \"\"\"\n Predict masks for the given input prompts, using the currently set image.\n Input prompts are batched torch tensors and are expected to already be\n transformed to the input frame using SAM2Transforms.\n\n Arguments:\n point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the\n model. Each point is in (X,Y) in pixels.\n point_labels (torch.Tensor or None): A BxN array of labels for the\n point prompts. 1 indicates a foreground point and 0 indicates a\n background point.\n boxes (np.ndarray or None): A Bx4 array given a box prompt to the\n model, in XYXY format.\n mask_input (np.ndarray): A low resolution mask input to the model, typically\n coming from a previous prediction iteration. Has form Bx1xHxW, where\n for SAM, H=W=256. Masks returned by a previous iteration of the\n predict method do not need further transformation.\n multimask_output (bool): If true, the model will return three masks.\n For ambiguous input prompts (such as a single click), this will often\n produce better masks than a single prediction. If only a single\n mask is needed, the model's predicted quality score can be used\n to select the best mask. For non-ambiguous prompts, such as multiple\n input prompts, multimask_output=False can give better results.\n return_logits (bool): If true, returns un-thresholded masks logits\n instead of a binary mask.\n\n Returns:\n (torch.Tensor): The output masks in BxCxHxW format, where C is the\n number of masks, and (H, W) is the original image size.\n (torch.Tensor): An array of shape BxC containing the model's\n predictions for the quality of each mask.\n (torch.Tensor): An array of shape BxCxHxW, where C is the number\n of masks and H=W=256. These low res logits can be passed to\n a subsequent iteration as mask input.\n \"\"\"\n if not self._is_image_set:\n raise RuntimeError(\n \"An image must be set with .set_image(...) before mask prediction.\"\n )\n\n if point_coords is not None:\n concat_points = (point_coords, point_labels)\n else:\n concat_points = None\n\n # Embed prompts\n if boxes is not None:\n box_coords = boxes.reshape(-1, 2, 2)\n box_labels = torch.tensor([[2, 3]], dtype=torch.int, device=boxes.device)\n box_labels = box_labels.repeat(boxes.size(0), 1)\n # we merge \"boxes\" and \"points\" into a single \"concat_points\" input (where\n # boxes are added at the beginning) to sam_prompt_encoder\n if concat_points is not None:\n concat_coords = torch.cat([box_coords, concat_points[0]], dim=1)\n concat_labels = torch.cat([box_labels, concat_points[1]], dim=1)\n concat_points = (concat_coords, concat_labels)\n else:\n concat_points = (box_coords, box_labels)\n\n sparse_embeddings, dense_embeddings = self.model.sam_prompt_encoder(\n points=concat_points,\n boxes=None,\n masks=mask_input,\n )\n\n # Predict masks\n batched_mode = (\n concat_points is not None and concat_points[0].shape[0] > 1\n ) # multi object prediction\n high_res_features = [\n feat_level[img_idx].unsqueeze(0)\n for feat_level in self._features[\"high_res_feats\"]\n ]\n low_res_masks, iou_predictions, _, _ = self.model.sam_mask_decoder(\n image_embeddings=self._features[\"image_embed\"][img_idx].unsqueeze(0),\n image_pe=self.model.sam_prompt_encoder.get_dense_pe(),\n sparse_prompt_embeddings=sparse_embeddings,\n dense_prompt_embeddings=dense_embeddings,\n multimask_output=multimask_output,\n repeat_image=batched_mode,\n high_res_features=high_res_features,\n )\n\n # Upscale the masks to the original image resolution\n masks = self._transforms.postprocess_masks(\n low_res_masks, self._orig_hw[img_idx]\n )\n low_res_masks = torch.clamp(low_res_masks, -32.0, 32.0)\n if not return_logits:\n masks = masks > self.mask_threshold\n\n return masks, iou_predictions, low_res_masks\n\n def get_image_embedding(self) -> torch.Tensor:\n \"\"\"\n Returns the image embeddings for the currently set image, with\n shape 1xCxHxW, where C is the embedding dimension and (H,W) are\n the embedding spatial dimension of SAM (typically C=256, H=W=64).\n \"\"\"\n if not self._is_image_set:\n raise RuntimeError(\n \"An image must be set with .set_image(...) to generate an embedding.\"\n )\n assert (\n self._features is not None\n ), \"Features must exist if an image has been set.\"\n return self._features[\"image_embed\"]\n\n @property\n def device(self) -> torch.device:\n return self.model.device\n\n def reset_predictor(self) -> None:\n \"\"\"\n Resets the image embeddings and other state variables.\n \"\"\"\n self._is_image_set = False\n self._features = None\n self._orig_hw = None\n self._is_batch = False\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/sam2_video_predictor.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport warnings\nfrom collections import OrderedDict\n\nimport torch\nimport torch.nn.functional as F\n\nfrom tqdm import tqdm\n\nfrom sam2.modeling.sam2_base import NO_OBJ_SCORE, SAM2Base\nfrom sam2.utils.misc import concat_points, fill_holes_in_mask_scores, load_video_frames\n\n\nclass SAM2VideoPredictor(SAM2Base):\n \"\"\"The predictor class to handle user interactions and manage inference states.\"\"\"\n\n def __init__(\n self,\n fill_hole_area=0,\n # whether to apply non-overlapping constraints on the output object masks\n non_overlap_masks=False,\n # whether to clear non-conditioning memory of the surrounding frames (which may contain outdated information) after adding correction clicks;\n # note that this would only apply to *single-object tracking* unless `clear_non_cond_mem_for_multi_obj` is also set to True)\n clear_non_cond_mem_around_input=False,\n # if `add_all_frames_to_correct_as_cond` is True, we also append to the conditioning frame list any frame that receives a later correction click\n # if `add_all_frames_to_correct_as_cond` is False, we conditioning frame list to only use those initial conditioning frames\n add_all_frames_to_correct_as_cond=False,\n **kwargs,\n ):\n super().__init__(**kwargs)\n self.fill_hole_area = fill_hole_area\n self.non_overlap_masks = non_overlap_masks\n self.clear_non_cond_mem_around_input = clear_non_cond_mem_around_input\n self.add_all_frames_to_correct_as_cond = add_all_frames_to_correct_as_cond\n\n @torch.inference_mode()\n def init_state(\n self,\n video_path,\n offload_video_to_cpu=False,\n offload_state_to_cpu=False,\n async_loading_frames=False,\n ):\n \"\"\"Initialize an inference state.\"\"\"\n compute_device = self.device # device of the model\n images, video_height, video_width = load_video_frames(\n video_path=video_path,\n image_size=self.image_size,\n offload_video_to_cpu=offload_video_to_cpu,\n async_loading_frames=async_loading_frames,\n compute_device=compute_device,\n )\n inference_state = {}\n inference_state[\"images\"] = images\n inference_state[\"num_frames\"] = len(images)\n # whether to offload the video frames to CPU memory\n # turning on this option saves the GPU memory with only a very small overhead\n inference_state[\"offload_video_to_cpu\"] = offload_video_to_cpu\n # whether to offload the inference state to CPU memory\n # turning on this option saves the GPU memory at the cost of a lower tracking fps\n # (e.g. in a test case of 768x768 model, fps dropped from 27 to 24 when tracking one object\n # and from 24 to 21 when tracking two objects)\n inference_state[\"offload_state_to_cpu\"] = offload_state_to_cpu\n # the original video height and width, used for resizing final output scores\n inference_state[\"video_height\"] = video_height\n inference_state[\"video_width\"] = video_width\n inference_state[\"device\"] = compute_device\n if offload_state_to_cpu:\n inference_state[\"storage_device\"] = torch.device(\"cpu\")\n else:\n inference_state[\"storage_device\"] = compute_device\n # inputs on each frame\n inference_state[\"point_inputs_per_obj\"] = {}\n inference_state[\"mask_inputs_per_obj\"] = {}\n # visual features on a small number of recently visited frames for quick interactions\n inference_state[\"cached_features\"] = {}\n # values that don't change across frames (so we only need to hold one copy of them)\n inference_state[\"constants\"] = {}\n # mapping between client-side object id and model-side object index\n inference_state[\"obj_id_to_idx\"] = OrderedDict()\n inference_state[\"obj_idx_to_id\"] = OrderedDict()\n inference_state[\"obj_ids\"] = []\n # Slice (view) of each object tracking results, sharing the same memory with \"output_dict\"\n inference_state[\"output_dict_per_obj\"] = {}\n # A temporary storage to hold new outputs when user interact with a frame\n # to add clicks or mask (it's merged into \"output_dict\" before propagation starts)\n inference_state[\"temp_output_dict_per_obj\"] = {}\n # Frames that already holds consolidated outputs from click or mask inputs\n # (we directly use their consolidated outputs during tracking)\n # metadata for each tracking frame (e.g. which direction it's tracked)\n inference_state[\"frames_tracked_per_obj\"] = {}\n # Warm up the visual backbone and cache the image feature on frame 0\n self._get_image_feature(inference_state, frame_idx=0, batch_size=1)\n return inference_state\n\n @classmethod\n def from_pretrained(cls, model_id: str, **kwargs) -> \"SAM2VideoPredictor\":\n \"\"\"\n Load a pretrained model from the Hugging Face hub.\n\n Arguments:\n model_id (str): The Hugging Face repository ID.\n **kwargs: Additional arguments to pass to the model constructor.\n\n Returns:\n (SAM2VideoPredictor): The loaded model.\n \"\"\"\n from sam2.build_sam import build_sam2_video_predictor_hf\n\n sam_model = build_sam2_video_predictor_hf(model_id, **kwargs)\n return sam_model\n\n def _obj_id_to_idx(self, inference_state, obj_id):\n \"\"\"Map client-side object id to model-side object index.\"\"\"\n obj_idx = inference_state[\"obj_id_to_idx\"].get(obj_id, None)\n if obj_idx is not None:\n return obj_idx\n\n # We always allow adding new objects (including after tracking starts).\n allow_new_object = True\n if allow_new_object:\n # get the next object slot\n obj_idx = len(inference_state[\"obj_id_to_idx\"])\n inference_state[\"obj_id_to_idx\"][obj_id] = obj_idx\n inference_state[\"obj_idx_to_id\"][obj_idx] = obj_id\n inference_state[\"obj_ids\"] = list(inference_state[\"obj_id_to_idx\"])\n # set up input and output structures for this object\n inference_state[\"point_inputs_per_obj\"][obj_idx] = {}\n inference_state[\"mask_inputs_per_obj\"][obj_idx] = {}\n inference_state[\"output_dict_per_obj\"][obj_idx] = {\n \"cond_frame_outputs\": {}, # dict containing {frame_idx: }\n \"non_cond_frame_outputs\": {}, # dict containing {frame_idx: }\n }\n inference_state[\"temp_output_dict_per_obj\"][obj_idx] = {\n \"cond_frame_outputs\": {}, # dict containing {frame_idx: }\n \"non_cond_frame_outputs\": {}, # dict containing {frame_idx: }\n }\n inference_state[\"frames_tracked_per_obj\"][obj_idx] = {}\n return obj_idx\n else:\n raise RuntimeError(\n f\"Cannot add new object id {obj_id} after tracking starts. \"\n f\"All existing object ids: {inference_state['obj_ids']}. \"\n f\"Please call 'reset_state' to restart from scratch.\"\n )\n\n def _obj_idx_to_id(self, inference_state, obj_idx):\n \"\"\"Map model-side object index to client-side object id.\"\"\"\n return inference_state[\"obj_idx_to_id\"][obj_idx]\n\n def _get_obj_num(self, inference_state):\n \"\"\"Get the total number of unique object ids received so far in this session.\"\"\"\n return len(inference_state[\"obj_idx_to_id\"])\n\n @torch.inference_mode()\n def add_new_points_or_box(\n self,\n inference_state,\n frame_idx,\n obj_id,\n points=None,\n labels=None,\n clear_old_points=True,\n normalize_coords=True,\n box=None,\n ):\n \"\"\"Add new points to a frame.\"\"\"\n obj_idx = self._obj_id_to_idx(inference_state, obj_id)\n point_inputs_per_frame = inference_state[\"point_inputs_per_obj\"][obj_idx]\n mask_inputs_per_frame = inference_state[\"mask_inputs_per_obj\"][obj_idx]\n\n if (points is not None) != (labels is not None):\n raise ValueError(\"points and labels must be provided together\")\n if points is None and box is None:\n raise ValueError(\"at least one of points or box must be provided as input\")\n\n if points is None:\n points = torch.zeros(0, 2, dtype=torch.float32)\n elif not isinstance(points, torch.Tensor):\n points = torch.tensor(points, dtype=torch.float32)\n if labels is None:\n labels = torch.zeros(0, dtype=torch.int32)\n elif not isinstance(labels, torch.Tensor):\n labels = torch.tensor(labels, dtype=torch.int32)\n if points.dim() == 2:\n points = points.unsqueeze(0) # add batch dimension\n if labels.dim() == 1:\n labels = labels.unsqueeze(0) # add batch dimension\n\n # If `box` is provided, we add it as the first two points with labels 2 and 3\n # along with the user-provided points (consistent with how SAM 2 is trained).\n if box is not None:\n if not clear_old_points:\n raise ValueError(\n \"cannot add box without clearing old points, since \"\n \"box prompt must be provided before any point prompt \"\n \"(please use clear_old_points=True instead)\"\n )\n if not isinstance(box, torch.Tensor):\n box = torch.tensor(box, dtype=torch.float32, device=points.device)\n box_coords = box.reshape(1, 2, 2)\n box_labels = torch.tensor([2, 3], dtype=torch.int32, device=labels.device)\n box_labels = box_labels.reshape(1, 2)\n points = torch.cat([box_coords, points], dim=1)\n labels = torch.cat([box_labels, labels], dim=1)\n\n if normalize_coords:\n video_H = inference_state[\"video_height\"]\n video_W = inference_state[\"video_width\"]\n points = points / torch.tensor([video_W, video_H]).to(points.device)\n # scale the (normalized) coordinates by the model's internal image size\n points = points * self.image_size\n points = points.to(inference_state[\"device\"])\n labels = labels.to(inference_state[\"device\"])\n\n if not clear_old_points:\n point_inputs = point_inputs_per_frame.get(frame_idx, None)\n else:\n point_inputs = None\n point_inputs = concat_points(point_inputs, points, labels)\n\n point_inputs_per_frame[frame_idx] = point_inputs\n mask_inputs_per_frame.pop(frame_idx, None)\n # If this frame hasn't been tracked before, we treat it as an initial conditioning\n # frame, meaning that the inputs points are to generate segments on this frame without\n # using any memory from other frames, like in SAM. Otherwise (if it has been tracked),\n # the input points will be used to correct the already tracked masks.\n obj_frames_tracked = inference_state[\"frames_tracked_per_obj\"][obj_idx]\n is_init_cond_frame = frame_idx not in obj_frames_tracked\n # whether to track in reverse time order\n if is_init_cond_frame:\n reverse = False\n else:\n reverse = obj_frames_tracked[frame_idx][\"reverse\"]\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n obj_temp_output_dict = inference_state[\"temp_output_dict_per_obj\"][obj_idx]\n # Add a frame to conditioning output if it's an initial conditioning frame or\n # if the model sees all frames receiving clicks/mask as conditioning frames.\n is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond\n storage_key = \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n\n # Get any previously predicted mask logits on this object and feed it along with\n # the new clicks into the SAM mask decoder.\n prev_sam_mask_logits = None\n # lookup temporary output dict first, which contains the most recent output\n # (if not found, then lookup conditioning and non-conditioning frame output)\n prev_out = obj_temp_output_dict[storage_key].get(frame_idx)\n if prev_out is None:\n prev_out = obj_output_dict[\"cond_frame_outputs\"].get(frame_idx)\n if prev_out is None:\n prev_out = obj_output_dict[\"non_cond_frame_outputs\"].get(frame_idx)\n\n if prev_out is not None and prev_out[\"pred_masks\"] is not None:\n device = inference_state[\"device\"]\n prev_sam_mask_logits = prev_out[\"pred_masks\"].to(device, non_blocking=True)\n # Clamp the scale of prev_sam_mask_logits to avoid rare numerical issues.\n prev_sam_mask_logits = torch.clamp(prev_sam_mask_logits, -32.0, 32.0)\n current_out, _ = self._run_single_frame_inference(\n inference_state=inference_state,\n output_dict=obj_output_dict, # run on the slice of a single object\n frame_idx=frame_idx,\n batch_size=1, # run on the slice of a single object\n is_init_cond_frame=is_init_cond_frame,\n point_inputs=point_inputs,\n mask_inputs=None,\n reverse=reverse,\n # Skip the memory encoder when adding clicks or mask. We execute the memory encoder\n # at the beginning of `propagate_in_video` (after user finalize their clicks). This\n # allows us to enforce non-overlapping constraints on all objects before encoding\n # them into memory.\n run_mem_encoder=False,\n prev_sam_mask_logits=prev_sam_mask_logits,\n )\n # Add the output to the output dict (to be used as future memory)\n obj_temp_output_dict[storage_key][frame_idx] = current_out\n\n # Resize the output mask to the original video resolution\n obj_ids = inference_state[\"obj_ids\"]\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n return frame_idx, obj_ids, video_res_masks\n\n def add_new_points(self, *args, **kwargs):\n \"\"\"Deprecated method. Please use `add_new_points_or_box` instead.\"\"\"\n return self.add_new_points_or_box(*args, **kwargs)\n\n @torch.inference_mode()\n def add_new_mask(\n self,\n inference_state,\n frame_idx,\n obj_id,\n mask,\n ):\n \"\"\"Add new mask to a frame.\"\"\"\n obj_idx = self._obj_id_to_idx(inference_state, obj_id)\n point_inputs_per_frame = inference_state[\"point_inputs_per_obj\"][obj_idx]\n mask_inputs_per_frame = inference_state[\"mask_inputs_per_obj\"][obj_idx]\n\n if not isinstance(mask, torch.Tensor):\n mask = torch.tensor(mask, dtype=torch.bool)\n assert mask.dim() == 2\n mask_H, mask_W = mask.shape\n mask_inputs_orig = mask[None, None] # add batch and channel dimension\n mask_inputs_orig = mask_inputs_orig.float().to(inference_state[\"device\"])\n\n # resize the mask if it doesn't match the model's image size\n if mask_H != self.image_size or mask_W != self.image_size:\n mask_inputs = torch.nn.functional.interpolate(\n mask_inputs_orig,\n size=(self.image_size, self.image_size),\n align_corners=False,\n mode=\"bilinear\",\n antialias=True, # use antialias for downsampling\n )\n mask_inputs = (mask_inputs >= 0.5).float()\n else:\n mask_inputs = mask_inputs_orig\n\n mask_inputs_per_frame[frame_idx] = mask_inputs\n point_inputs_per_frame.pop(frame_idx, None)\n # If this frame hasn't been tracked before, we treat it as an initial conditioning\n # frame, meaning that the inputs points are to generate segments on this frame without\n # using any memory from other frames, like in SAM. Otherwise (if it has been tracked),\n # the input points will be used to correct the already tracked masks.\n obj_frames_tracked = inference_state[\"frames_tracked_per_obj\"][obj_idx]\n is_init_cond_frame = frame_idx not in obj_frames_tracked\n # whether to track in reverse time order\n if is_init_cond_frame:\n reverse = False\n else:\n reverse = obj_frames_tracked[frame_idx][\"reverse\"]\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n obj_temp_output_dict = inference_state[\"temp_output_dict_per_obj\"][obj_idx]\n # Add a frame to conditioning output if it's an initial conditioning frame or\n # if the model sees all frames receiving clicks/mask as conditioning frames.\n is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond\n storage_key = \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n\n current_out, _ = self._run_single_frame_inference(\n inference_state=inference_state,\n output_dict=obj_output_dict, # run on the slice of a single object\n frame_idx=frame_idx,\n batch_size=1, # run on the slice of a single object\n is_init_cond_frame=is_init_cond_frame,\n point_inputs=None,\n mask_inputs=mask_inputs,\n reverse=reverse,\n # Skip the memory encoder when adding clicks or mask. We execute the memory encoder\n # at the beginning of `propagate_in_video` (after user finalize their clicks). This\n # allows us to enforce non-overlapping constraints on all objects before encoding\n # them into memory.\n run_mem_encoder=False,\n )\n # Add the output to the output dict (to be used as future memory)\n obj_temp_output_dict[storage_key][frame_idx] = current_out\n\n # Resize the output mask to the original video resolution\n obj_ids = inference_state[\"obj_ids\"]\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n return frame_idx, obj_ids, video_res_masks\n\n def _get_orig_video_res_output(self, inference_state, any_res_masks):\n \"\"\"\n Resize the object scores to the original video resolution (video_res_masks)\n and apply non-overlapping constraints for final output.\n \"\"\"\n device = inference_state[\"device\"]\n video_H = inference_state[\"video_height\"]\n video_W = inference_state[\"video_width\"]\n any_res_masks = any_res_masks.to(device, non_blocking=True)\n if any_res_masks.shape[-2:] == (video_H, video_W):\n video_res_masks = any_res_masks\n else:\n video_res_masks = torch.nn.functional.interpolate(\n any_res_masks,\n size=(video_H, video_W),\n mode=\"bilinear\",\n align_corners=False,\n )\n if self.non_overlap_masks:\n video_res_masks = self._apply_non_overlapping_constraints(video_res_masks)\n return any_res_masks, video_res_masks\n\n def _consolidate_temp_output_across_obj(\n self,\n inference_state,\n frame_idx,\n is_cond,\n consolidate_at_video_res=False,\n ):\n \"\"\"\n Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on\n a frame into a single output for all objects, including\n 1) fill any missing objects either from `output_dict_per_obj` (if they exist in\n `output_dict_per_obj` for this frame) or leave them as placeholder values\n (if they don't exist in `output_dict_per_obj` for this frame);\n 2) if specified, rerun memory encoder after apply non-overlapping constraints\n on the object scores.\n \"\"\"\n batch_size = self._get_obj_num(inference_state)\n storage_key = \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n # Optionally, we allow consolidating the temporary outputs at the original\n # video resolution (to provide a better editing experience for mask prompts).\n if consolidate_at_video_res:\n consolidated_H = inference_state[\"video_height\"]\n consolidated_W = inference_state[\"video_width\"]\n consolidated_mask_key = \"pred_masks_video_res\"\n else:\n consolidated_H = consolidated_W = self.image_size // 4\n consolidated_mask_key = \"pred_masks\"\n\n # Initialize `consolidated_out`. Its \"maskmem_features\" and \"maskmem_pos_enc\"\n # will be added when rerunning the memory encoder after applying non-overlapping\n # constraints to object scores. Its \"pred_masks\" are prefilled with a large\n # negative value (NO_OBJ_SCORE) to represent missing objects.\n consolidated_out = {\n consolidated_mask_key: torch.full(\n size=(batch_size, 1, consolidated_H, consolidated_W),\n fill_value=NO_OBJ_SCORE,\n dtype=torch.float32,\n device=inference_state[\"storage_device\"],\n ),\n }\n for obj_idx in range(batch_size):\n obj_temp_output_dict = inference_state[\"temp_output_dict_per_obj\"][obj_idx]\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n out = obj_temp_output_dict[storage_key].get(frame_idx, None)\n # If the object doesn't appear in \"temp_output_dict_per_obj\" on this frame,\n # we fall back and look up its previous output in \"output_dict_per_obj\".\n # We look up both \"cond_frame_outputs\" and \"non_cond_frame_outputs\" in\n # \"output_dict_per_obj\" to find a previous output for this object.\n if out is None:\n out = obj_output_dict[\"cond_frame_outputs\"].get(frame_idx, None)\n if out is None:\n out = obj_output_dict[\"non_cond_frame_outputs\"].get(frame_idx, None)\n # If the object doesn't appear in \"output_dict_per_obj\" either, we skip it\n # and leave its mask scores to the default scores (i.e. the NO_OBJ_SCORE\n # placeholder above) and set its object pointer to be a dummy pointer.\n if out is None:\n continue\n # Add the temporary object output mask to consolidated output mask\n obj_mask = out[\"pred_masks\"]\n consolidated_pred_masks = consolidated_out[consolidated_mask_key]\n if obj_mask.shape[-2:] == consolidated_pred_masks.shape[-2:]:\n consolidated_pred_masks[obj_idx : obj_idx + 1] = obj_mask\n else:\n # Resize first if temporary object mask has a different resolution\n resized_obj_mask = torch.nn.functional.interpolate(\n obj_mask,\n size=consolidated_pred_masks.shape[-2:],\n mode=\"bilinear\",\n align_corners=False,\n )\n consolidated_pred_masks[obj_idx : obj_idx + 1] = resized_obj_mask\n\n return consolidated_out\n\n @torch.inference_mode()\n def propagate_in_video_preflight(self, inference_state):\n \"\"\"Prepare inference_state and consolidate temporary outputs before tracking.\"\"\"\n # Check and make sure that every object has received input points or masks.\n batch_size = self._get_obj_num(inference_state)\n if batch_size == 0:\n raise RuntimeError(\n \"No input points or masks are provided for any object; please add inputs first.\"\n )\n\n # Consolidate per-object temporary outputs in \"temp_output_dict_per_obj\" and\n # add them into \"output_dict\".\n for obj_idx in range(batch_size):\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n obj_temp_output_dict = inference_state[\"temp_output_dict_per_obj\"][obj_idx]\n for is_cond in [False, True]:\n # Separately consolidate conditioning and non-conditioning temp outputs\n storage_key = (\n \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n )\n # Find all the frames that contain temporary outputs for any objects\n # (these should be the frames that have just received clicks for mask inputs\n # via `add_new_points_or_box` or `add_new_mask`)\n for frame_idx, out in obj_temp_output_dict[storage_key].items():\n # Run memory encoder on the temporary outputs (if the memory feature is missing)\n if out[\"maskmem_features\"] is None:\n high_res_masks = torch.nn.functional.interpolate(\n out[\"pred_masks\"].to(inference_state[\"device\"]),\n size=(self.image_size, self.image_size),\n mode=\"bilinear\",\n align_corners=False,\n )\n maskmem_features, maskmem_pos_enc = self._run_memory_encoder(\n inference_state=inference_state,\n frame_idx=frame_idx,\n batch_size=1, # run on the slice of a single object\n high_res_masks=high_res_masks,\n object_score_logits=out[\"object_score_logits\"],\n # these frames are what the user interacted with\n is_mask_from_pts=True,\n )\n out[\"maskmem_features\"] = maskmem_features\n out[\"maskmem_pos_enc\"] = maskmem_pos_enc\n\n obj_output_dict[storage_key][frame_idx] = out\n if self.clear_non_cond_mem_around_input:\n # clear non-conditioning memory of the surrounding frames\n self._clear_obj_non_cond_mem_around_input(\n inference_state, frame_idx, obj_idx\n )\n\n # clear temporary outputs in `temp_output_dict_per_obj`\n obj_temp_output_dict[storage_key].clear()\n\n # check and make sure that every object has received input points or masks\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n if len(obj_output_dict[\"cond_frame_outputs\"]) == 0:\n obj_id = self._obj_idx_to_id(inference_state, obj_idx)\n raise RuntimeError(\n f\"No input points or masks are provided for object id {obj_id}; please add inputs first.\"\n )\n # edge case: if an output is added to \"cond_frame_outputs\", we remove any prior\n # output on the same frame in \"non_cond_frame_outputs\"\n for frame_idx in obj_output_dict[\"cond_frame_outputs\"]:\n obj_output_dict[\"non_cond_frame_outputs\"].pop(frame_idx, None)\n\n @torch.inference_mode()\n def propagate_in_video(\n self,\n inference_state,\n start_frame_idx=None,\n max_frame_num_to_track=None,\n reverse=False,\n ):\n \"\"\"Propagate the input points across frames to track in the entire video.\"\"\"\n self.propagate_in_video_preflight(inference_state)\n\n obj_ids = inference_state[\"obj_ids\"]\n num_frames = inference_state[\"num_frames\"]\n batch_size = self._get_obj_num(inference_state)\n\n # set start index, end index, and processing order\n if start_frame_idx is None:\n # default: start from the earliest frame with input points\n start_frame_idx = min(\n t\n for obj_output_dict in inference_state[\"output_dict_per_obj\"].values()\n for t in obj_output_dict[\"cond_frame_outputs\"]\n )\n if max_frame_num_to_track is None:\n # default: track all the frames in the video\n max_frame_num_to_track = num_frames\n if reverse:\n end_frame_idx = max(start_frame_idx - max_frame_num_to_track, 0)\n if start_frame_idx > 0:\n processing_order = range(start_frame_idx, end_frame_idx - 1, -1)\n else:\n processing_order = [] # skip reverse tracking if starting from frame 0\n else:\n end_frame_idx = min(\n start_frame_idx + max_frame_num_to_track, num_frames - 1\n )\n processing_order = range(start_frame_idx, end_frame_idx + 1)\n\n for frame_idx in tqdm(processing_order, desc=\"propagate in video\"):\n pred_masks_per_obj = [None] * batch_size\n for obj_idx in range(batch_size):\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n # We skip those frames already in consolidated outputs (these are frames\n # that received input clicks or mask). Note that we cannot directly run\n # batched forward on them via `_run_single_frame_inference` because the\n # number of clicks on each object might be different.\n if frame_idx in obj_output_dict[\"cond_frame_outputs\"]:\n storage_key = \"cond_frame_outputs\"\n current_out = obj_output_dict[storage_key][frame_idx]\n device = inference_state[\"device\"]\n pred_masks = current_out[\"pred_masks\"].to(device, non_blocking=True)\n if self.clear_non_cond_mem_around_input:\n # clear non-conditioning memory of the surrounding frames\n self._clear_obj_non_cond_mem_around_input(\n inference_state, frame_idx, obj_idx\n )\n else:\n storage_key = \"non_cond_frame_outputs\"\n current_out, pred_masks = self._run_single_frame_inference(\n inference_state=inference_state,\n output_dict=obj_output_dict,\n frame_idx=frame_idx,\n batch_size=1, # run on the slice of a single object\n is_init_cond_frame=False,\n point_inputs=None,\n mask_inputs=None,\n reverse=reverse,\n run_mem_encoder=True,\n )\n obj_output_dict[storage_key][frame_idx] = current_out\n\n inference_state[\"frames_tracked_per_obj\"][obj_idx][frame_idx] = {\n \"reverse\": reverse\n }\n pred_masks_per_obj[obj_idx] = pred_masks\n\n # Resize the output mask to the original video resolution (we directly use\n # the mask scores on GPU for output to avoid any CPU conversion in between)\n if len(pred_masks_per_obj) > 1:\n all_pred_masks = torch.cat(pred_masks_per_obj, dim=0)\n else:\n all_pred_masks = pred_masks_per_obj[0]\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, all_pred_masks\n )\n yield frame_idx, obj_ids, video_res_masks\n\n @torch.inference_mode()\n def clear_all_prompts_in_frame(\n self, inference_state, frame_idx, obj_id, need_output=True\n ):\n \"\"\"Remove all input points or mask in a specific frame for a given object.\"\"\"\n obj_idx = self._obj_id_to_idx(inference_state, obj_id)\n\n # Clear the conditioning information on the given frame\n inference_state[\"point_inputs_per_obj\"][obj_idx].pop(frame_idx, None)\n inference_state[\"mask_inputs_per_obj\"][obj_idx].pop(frame_idx, None)\n\n temp_output_dict_per_obj = inference_state[\"temp_output_dict_per_obj\"]\n temp_output_dict_per_obj[obj_idx][\"cond_frame_outputs\"].pop(frame_idx, None)\n temp_output_dict_per_obj[obj_idx][\"non_cond_frame_outputs\"].pop(frame_idx, None)\n\n # Remove the frame's conditioning output (possibly downgrading it to non-conditioning)\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n out = obj_output_dict[\"cond_frame_outputs\"].pop(frame_idx, None)\n if out is not None:\n # The frame is not a conditioning frame anymore since it's not receiving inputs,\n # so we \"downgrade\" its output (if exists) to a non-conditioning frame output.\n obj_output_dict[\"non_cond_frame_outputs\"][frame_idx] = out\n inference_state[\"frames_tracked_per_obj\"][obj_idx].pop(frame_idx, None)\n\n if not need_output:\n return\n # Finally, output updated masks per object (after removing the inputs above)\n obj_ids = inference_state[\"obj_ids\"]\n is_cond = any(\n frame_idx in obj_temp_output_dict[\"cond_frame_outputs\"]\n for obj_temp_output_dict in temp_output_dict_per_obj.values()\n )\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n return frame_idx, obj_ids, video_res_masks\n\n @torch.inference_mode()\n def reset_state(self, inference_state):\n \"\"\"Remove all input points or mask in all frames throughout the video.\"\"\"\n self._reset_tracking_results(inference_state)\n # Remove all object ids\n inference_state[\"obj_id_to_idx\"].clear()\n inference_state[\"obj_idx_to_id\"].clear()\n inference_state[\"obj_ids\"].clear()\n inference_state[\"point_inputs_per_obj\"].clear()\n inference_state[\"mask_inputs_per_obj\"].clear()\n inference_state[\"output_dict_per_obj\"].clear()\n inference_state[\"temp_output_dict_per_obj\"].clear()\n inference_state[\"frames_tracked_per_obj\"].clear()\n\n def _reset_tracking_results(self, inference_state):\n \"\"\"Reset all tracking inputs and results across the videos.\"\"\"\n for v in inference_state[\"point_inputs_per_obj\"].values():\n v.clear()\n for v in inference_state[\"mask_inputs_per_obj\"].values():\n v.clear()\n for v in inference_state[\"output_dict_per_obj\"].values():\n v[\"cond_frame_outputs\"].clear()\n v[\"non_cond_frame_outputs\"].clear()\n for v in inference_state[\"temp_output_dict_per_obj\"].values():\n v[\"cond_frame_outputs\"].clear()\n v[\"non_cond_frame_outputs\"].clear()\n for v in inference_state[\"frames_tracked_per_obj\"].values():\n v.clear()\n\n def _get_image_feature(self, inference_state, frame_idx, batch_size):\n \"\"\"Compute the image features on a given frame.\"\"\"\n # Look up in the cache first\n image, backbone_out = inference_state[\"cached_features\"].get(\n frame_idx, (None, None)\n )\n if backbone_out is None:\n # Cache miss -- we will run inference on a single image\n device = inference_state[\"device\"]\n image = inference_state[\"images\"][frame_idx].to(device).float().unsqueeze(0)\n backbone_out = self.forward_image(image)\n # Cache the most recent frame's feature (for repeated interactions with\n # a frame; we can use an LRU cache for more frames in the future).\n inference_state[\"cached_features\"] = {frame_idx: (image, backbone_out)}\n\n # expand the features to have the same dimension as the number of objects\n expanded_image = image.expand(batch_size, -1, -1, -1)\n expanded_backbone_out = {\n \"backbone_fpn\": backbone_out[\"backbone_fpn\"].copy(),\n \"vision_pos_enc\": backbone_out[\"vision_pos_enc\"].copy(),\n }\n for i, feat in enumerate(expanded_backbone_out[\"backbone_fpn\"]):\n expanded_backbone_out[\"backbone_fpn\"][i] = feat.expand(\n batch_size, -1, -1, -1\n )\n for i, pos in enumerate(expanded_backbone_out[\"vision_pos_enc\"]):\n pos = pos.expand(batch_size, -1, -1, -1)\n expanded_backbone_out[\"vision_pos_enc\"][i] = pos\n\n features = self._prepare_backbone_features(expanded_backbone_out)\n features = (expanded_image,) + features\n return features\n\n def _run_single_frame_inference(\n self,\n inference_state,\n output_dict,\n frame_idx,\n batch_size,\n is_init_cond_frame,\n point_inputs,\n mask_inputs,\n reverse,\n run_mem_encoder,\n prev_sam_mask_logits=None,\n ):\n \"\"\"Run tracking on a single frame based on current inputs and previous memory.\"\"\"\n # Retrieve correct image features\n (\n _,\n _,\n current_vision_feats,\n current_vision_pos_embeds,\n feat_sizes,\n ) = self._get_image_feature(inference_state, frame_idx, batch_size)\n\n # point and mask should not appear as input simultaneously on the same frame\n assert point_inputs is None or mask_inputs is None\n current_out = self.track_step(\n frame_idx=frame_idx,\n is_init_cond_frame=is_init_cond_frame,\n current_vision_feats=current_vision_feats,\n current_vision_pos_embeds=current_vision_pos_embeds,\n feat_sizes=feat_sizes,\n point_inputs=point_inputs,\n mask_inputs=mask_inputs,\n output_dict=output_dict,\n num_frames=inference_state[\"num_frames\"],\n track_in_reverse=reverse,\n run_mem_encoder=run_mem_encoder,\n prev_sam_mask_logits=prev_sam_mask_logits,\n )\n\n # optionally offload the output to CPU memory to save GPU space\n storage_device = inference_state[\"storage_device\"]\n maskmem_features = current_out[\"maskmem_features\"]\n if maskmem_features is not None:\n maskmem_features = maskmem_features.to(torch.bfloat16)\n maskmem_features = maskmem_features.to(storage_device, non_blocking=True)\n pred_masks_gpu = current_out[\"pred_masks\"]\n # potentially fill holes in the predicted masks\n if self.fill_hole_area > 0:\n pred_masks_gpu = fill_holes_in_mask_scores(\n pred_masks_gpu, self.fill_hole_area\n )\n pred_masks = pred_masks_gpu.to(storage_device, non_blocking=True)\n # \"maskmem_pos_enc\" is the same across frames, so we only need to store one copy of it\n maskmem_pos_enc = self._get_maskmem_pos_enc(inference_state, current_out)\n # object pointer is a small tensor, so we always keep it on GPU memory for fast access\n obj_ptr = current_out[\"obj_ptr\"]\n object_score_logits = current_out[\"object_score_logits\"]\n # make a compact version of this frame's output to reduce the state size\n compact_current_out = {\n \"maskmem_features\": maskmem_features,\n \"maskmem_pos_enc\": maskmem_pos_enc,\n \"pred_masks\": pred_masks,\n \"obj_ptr\": obj_ptr,\n \"object_score_logits\": object_score_logits,\n }\n return compact_current_out, pred_masks_gpu\n\n def _run_memory_encoder(\n self,\n inference_state,\n frame_idx,\n batch_size,\n high_res_masks,\n object_score_logits,\n is_mask_from_pts,\n ):\n \"\"\"\n Run the memory encoder on `high_res_masks`. This is usually after applying\n non-overlapping constraints to object scores. Since their scores changed, their\n memory also need to be computed again with the memory encoder.\n \"\"\"\n # Retrieve correct image features\n _, _, current_vision_feats, _, feat_sizes = self._get_image_feature(\n inference_state, frame_idx, batch_size\n )\n maskmem_features, maskmem_pos_enc = self._encode_new_memory(\n current_vision_feats=current_vision_feats,\n feat_sizes=feat_sizes,\n pred_masks_high_res=high_res_masks,\n object_score_logits=object_score_logits,\n is_mask_from_pts=is_mask_from_pts,\n )\n\n # optionally offload the output to CPU memory to save GPU space\n storage_device = inference_state[\"storage_device\"]\n maskmem_features = maskmem_features.to(torch.bfloat16)\n maskmem_features = maskmem_features.to(storage_device, non_blocking=True)\n # \"maskmem_pos_enc\" is the same across frames, so we only need to store one copy of it\n maskmem_pos_enc = self._get_maskmem_pos_enc(\n inference_state, {\"maskmem_pos_enc\": maskmem_pos_enc}\n )\n return maskmem_features, maskmem_pos_enc\n\n def _get_maskmem_pos_enc(self, inference_state, current_out):\n \"\"\"\n `maskmem_pos_enc` is the same across frames and objects, so we cache it as\n a constant in the inference session to reduce session storage size.\n \"\"\"\n model_constants = inference_state[\"constants\"]\n # \"out_maskmem_pos_enc\" should be either a list of tensors or None\n out_maskmem_pos_enc = current_out[\"maskmem_pos_enc\"]\n if out_maskmem_pos_enc is not None:\n if \"maskmem_pos_enc\" not in model_constants:\n assert isinstance(out_maskmem_pos_enc, list)\n # only take the slice for one object, since it's same across objects\n maskmem_pos_enc = [x[0:1].clone() for x in out_maskmem_pos_enc]\n model_constants[\"maskmem_pos_enc\"] = maskmem_pos_enc\n else:\n maskmem_pos_enc = model_constants[\"maskmem_pos_enc\"]\n # expand the cached maskmem_pos_enc to the actual batch size\n batch_size = out_maskmem_pos_enc[0].size(0)\n expanded_maskmem_pos_enc = [\n x.expand(batch_size, -1, -1, -1) for x in maskmem_pos_enc\n ]\n else:\n expanded_maskmem_pos_enc = None\n return expanded_maskmem_pos_enc\n\n @torch.inference_mode()\n def remove_object(self, inference_state, obj_id, strict=False, need_output=True):\n \"\"\"\n Remove an object id from the tracking state. If strict is True, we check whether\n the object id actually exists and raise an error if it doesn't exist.\n \"\"\"\n old_obj_idx_to_rm = inference_state[\"obj_id_to_idx\"].get(obj_id, None)\n updated_frames = []\n # Check whether this object_id to remove actually exists and possibly raise an error.\n if old_obj_idx_to_rm is None:\n if not strict:\n return inference_state[\"obj_ids\"], updated_frames\n raise RuntimeError(\n f\"Cannot remove object id {obj_id} as it doesn't exist. \"\n f\"All existing object ids: {inference_state['obj_ids']}.\"\n )\n\n # If this is the only remaining object id, we simply reset the state.\n if len(inference_state[\"obj_id_to_idx\"]) == 1:\n self.reset_state(inference_state)\n return inference_state[\"obj_ids\"], updated_frames\n\n # There are still remaining objects after removing this object id. In this case,\n # we need to delete the object storage from inference state tensors.\n # Step 0: clear the input on those frames where this object id has point or mask input\n # (note that this step is required as it might downgrade conditioning frames to\n # non-conditioning ones)\n obj_input_frames_inds = set()\n obj_input_frames_inds.update(\n inference_state[\"point_inputs_per_obj\"][old_obj_idx_to_rm]\n )\n obj_input_frames_inds.update(\n inference_state[\"mask_inputs_per_obj\"][old_obj_idx_to_rm]\n )\n for frame_idx in obj_input_frames_inds:\n self.clear_all_prompts_in_frame(\n inference_state, frame_idx, obj_id, need_output=False\n )\n\n # Step 1: Update the object id mapping (note that it must be done after Step 0,\n # since Step 0 still requires the old object id mappings in inference_state)\n old_obj_ids = inference_state[\"obj_ids\"]\n old_obj_inds = list(range(len(old_obj_ids)))\n remain_old_obj_inds = old_obj_inds.copy()\n remain_old_obj_inds.remove(old_obj_idx_to_rm)\n new_obj_ids = [old_obj_ids[old_idx] for old_idx in remain_old_obj_inds]\n new_obj_inds = list(range(len(new_obj_ids)))\n # build new mappings\n old_idx_to_new_idx = dict(zip(remain_old_obj_inds, new_obj_inds))\n inference_state[\"obj_id_to_idx\"] = dict(zip(new_obj_ids, new_obj_inds))\n inference_state[\"obj_idx_to_id\"] = dict(zip(new_obj_inds, new_obj_ids))\n inference_state[\"obj_ids\"] = new_obj_ids\n\n # Step 2: For per-object tensor storage, we shift their obj_idx in the dict keys.\n def _map_keys(container):\n new_kvs = []\n for k in old_obj_inds:\n v = container.pop(k)\n if k in old_idx_to_new_idx:\n new_kvs.append((old_idx_to_new_idx[k], v))\n container.update(new_kvs)\n\n _map_keys(inference_state[\"point_inputs_per_obj\"])\n _map_keys(inference_state[\"mask_inputs_per_obj\"])\n _map_keys(inference_state[\"output_dict_per_obj\"])\n _map_keys(inference_state[\"temp_output_dict_per_obj\"])\n _map_keys(inference_state[\"frames_tracked_per_obj\"])\n\n # Step 3: Further collect the outputs on those frames in `obj_input_frames_inds`, which\n # could show an updated mask for objects previously occluded by the object being removed\n if need_output:\n temp_output_dict_per_obj = inference_state[\"temp_output_dict_per_obj\"]\n for frame_idx in obj_input_frames_inds:\n is_cond = any(\n frame_idx in obj_temp_output_dict[\"cond_frame_outputs\"]\n for obj_temp_output_dict in temp_output_dict_per_obj.values()\n )\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n updated_frames.append((frame_idx, video_res_masks))\n\n return inference_state[\"obj_ids\"], updated_frames\n\n def _clear_non_cond_mem_around_input(self, inference_state, frame_idx):\n \"\"\"\n Remove the non-conditioning memory around the input frame. When users provide\n correction clicks, the surrounding frames' non-conditioning memories can still\n contain outdated object appearance information and could confuse the model.\n\n This method clears those non-conditioning memories surrounding the interacted\n frame to avoid giving the model both old and new information about the object.\n \"\"\"\n r = self.memory_temporal_stride_for_eval\n frame_idx_begin = frame_idx - r * self.num_maskmem\n frame_idx_end = frame_idx + r * self.num_maskmem\n batch_size = self._get_obj_num(inference_state)\n for obj_idx in range(batch_size):\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n non_cond_frame_outputs = obj_output_dict[\"non_cond_frame_outputs\"]\n for t in range(frame_idx_begin, frame_idx_end + 1):\n non_cond_frame_outputs.pop(t, None)\n\n\nclass SAM2VideoPredictorVOS(SAM2VideoPredictor):\n \"\"\"Optimized for the VOS setting\"\"\"\n\n def __init__(self, *args, **kwargs):\n super().__init__(*args, **kwargs)\n self._compile_all_components()\n\n def _compile_all_components(self):\n print(\"Compiling all components for VOS setting. First time may be very slow.\")\n self.memory_encoder.forward = torch.compile(\n self.memory_encoder.forward,\n mode=\"max-autotune\",\n fullgraph=True,\n dynamic=False,\n )\n\n self.memory_attention.forward = torch.compile(\n self.memory_attention.forward,\n mode=\"max-autotune\",\n fullgraph=True,\n dynamic=True, # Num. of memories varies\n )\n\n self.sam_prompt_encoder.forward = torch.compile(\n self.sam_prompt_encoder.forward,\n mode=\"max-autotune\",\n fullgraph=True,\n dynamic=False, # Accuracy regression on True\n )\n\n self.sam_mask_decoder.forward = torch.compile(\n self.sam_mask_decoder.forward,\n mode=\"max-autotune\",\n fullgraph=True,\n dynamic=False, # Accuracy regression on True\n )\n\n def forward_image(self, img_batch: torch.Tensor):\n \"\"\"\n Identical to the corresponding method in the parent (SAM2VideoPredictor), but\n cloning the backbone features and pos encoding to enable compilation.\n \"\"\"\n backbone_out = self.image_encoder(img_batch)\n if self.use_high_res_features_in_sam:\n # precompute projected level 0 and level 1 features in SAM decoder\n # to avoid running it again on every SAM click\n backbone_out[\"backbone_fpn\"][0] = self.sam_mask_decoder.conv_s0(\n backbone_out[\"backbone_fpn\"][0]\n )\n backbone_out[\"backbone_fpn\"][1] = self.sam_mask_decoder.conv_s1(\n backbone_out[\"backbone_fpn\"][1]\n )\n # Clone to help torch.compile\n for i in range(len(backbone_out[\"backbone_fpn\"])):\n backbone_out[\"backbone_fpn\"][i] = backbone_out[\"backbone_fpn\"][i].clone()\n backbone_out[\"vision_pos_enc\"][i] = backbone_out[\"vision_pos_enc\"][\n i\n ].clone()\n return backbone_out\n\n def _forward_sam_heads(\n self,\n backbone_features,\n point_inputs=None,\n mask_inputs=None,\n high_res_features=None,\n multimask_output=False,\n ):\n \"\"\"\n Identical to the corresponding method in the parent (SAM2VideoPredictor), but\n cloning the outputs of prompt_encoder and mask_decoder to enable compilation.\n \"\"\"\n B = backbone_features.size(0)\n device = backbone_features.device\n assert backbone_features.size(1) == self.sam_prompt_embed_dim\n assert backbone_features.size(2) == self.sam_image_embedding_size\n assert backbone_features.size(3) == self.sam_image_embedding_size\n\n # a) Handle point prompts\n if point_inputs is not None:\n sam_point_coords = point_inputs[\"point_coords\"]\n sam_point_labels = point_inputs[\"point_labels\"]\n assert sam_point_coords.size(0) == B and sam_point_labels.size(0) == B\n else:\n # If no points are provide, pad with an empty point (with label -1)\n sam_point_coords = torch.zeros(B, 1, 2, device=device)\n sam_point_labels = -torch.ones(B, 1, dtype=torch.int32, device=device)\n\n # b) Handle mask prompts\n if mask_inputs is not None:\n # If mask_inputs is provided, downsize it into low-res mask input if needed\n # and feed it as a dense mask prompt into the SAM mask encoder\n assert len(mask_inputs.shape) == 4 and mask_inputs.shape[:2] == (B, 1)\n if mask_inputs.shape[-2:] != self.sam_prompt_encoder.mask_input_size:\n sam_mask_prompt = F.interpolate(\n mask_inputs.float(),\n size=self.sam_prompt_encoder.mask_input_size,\n align_corners=False,\n mode=\"bilinear\",\n antialias=True, # use antialias for downsampling\n )\n else:\n sam_mask_prompt = mask_inputs\n else:\n # Otherwise, simply feed None (and SAM's prompt encoder will add\n # a learned `no_mask_embed` to indicate no mask input in this case).\n sam_mask_prompt = None\n\n sparse_embeddings, dense_embeddings = self.sam_prompt_encoder(\n points=(sam_point_coords, sam_point_labels),\n boxes=None,\n masks=sam_mask_prompt,\n )\n # Clone image_pe and the outputs of sam_prompt_encoder\n # to enable compilation\n sparse_embeddings = sparse_embeddings.clone()\n dense_embeddings = dense_embeddings.clone()\n image_pe = self.sam_prompt_encoder.get_dense_pe().clone()\n (\n low_res_multimasks,\n ious,\n sam_output_tokens,\n object_score_logits,\n ) = self.sam_mask_decoder(\n image_embeddings=backbone_features,\n image_pe=image_pe,\n sparse_prompt_embeddings=sparse_embeddings,\n dense_prompt_embeddings=dense_embeddings,\n multimask_output=multimask_output,\n repeat_image=False, # the image is already batched\n high_res_features=high_res_features,\n )\n # Clone the output of sam_mask_decoder\n # to enable compilation\n low_res_multimasks = low_res_multimasks.clone()\n ious = ious.clone()\n sam_output_tokens = sam_output_tokens.clone()\n object_score_logits = object_score_logits.clone()\n\n if self.pred_obj_scores:\n is_obj_appearing = object_score_logits > 0\n\n # Mask used for spatial memories is always a *hard* choice between obj and no obj,\n # consistent with the actual mask prediction\n low_res_multimasks = torch.where(\n is_obj_appearing[:, None, None],\n low_res_multimasks,\n NO_OBJ_SCORE,\n )\n\n # convert masks from possibly bfloat16 (or float16) to float32\n # (older PyTorch versions before 2.1 don't support `interpolate` on bf16)\n low_res_multimasks = low_res_multimasks.float()\n high_res_multimasks = F.interpolate(\n low_res_multimasks,\n size=(self.image_size, self.image_size),\n mode=\"bilinear\",\n align_corners=False,\n )\n\n sam_output_token = sam_output_tokens[:, 0]\n if multimask_output:\n # take the best mask prediction (with the highest IoU estimation)\n best_iou_inds = torch.argmax(ious, dim=-1)\n batch_inds = torch.arange(B, device=device)\n low_res_masks = low_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1)\n high_res_masks = high_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1)\n if sam_output_tokens.size(1) > 1:\n sam_output_token = sam_output_tokens[batch_inds, best_iou_inds]\n else:\n low_res_masks, high_res_masks = low_res_multimasks, high_res_multimasks\n\n # Extract object pointer from the SAM output token (with occlusion handling)\n obj_ptr = self.obj_ptr_proj(sam_output_token)\n if self.pred_obj_scores:\n # Allow *soft* no obj ptr, unlike for masks\n if self.soft_no_obj_ptr:\n lambda_is_obj_appearing = object_score_logits.sigmoid()\n else:\n lambda_is_obj_appearing = is_obj_appearing.float()\n\n if self.fixed_no_obj_ptr:\n obj_ptr = lambda_is_obj_appearing * obj_ptr\n obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr\n\n return (\n low_res_multimasks,\n high_res_multimasks,\n ious,\n low_res_masks,\n high_res_masks,\n obj_ptr,\n object_score_logits,\n )\n\n def _encode_new_memory(\n self,\n current_vision_feats,\n feat_sizes,\n pred_masks_high_res,\n object_score_logits,\n is_mask_from_pts,\n ):\n \"\"\"\n Identical to the corresponding method in the parent (SAM2VideoPredictor), but\n cloning the memories and their pos enc to enable compilation.\n \"\"\"\n B = current_vision_feats[-1].size(1) # batch size on this frame\n C = self.hidden_dim\n H, W = feat_sizes[-1] # top-level (lowest-resolution) feature size\n # top-level feature, (HW)BC => BCHW\n pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W)\n if self.non_overlap_masks_for_mem_enc and not self.training:\n # optionally, apply non-overlapping constraints to the masks (it's applied\n # in the batch dimension and should only be used during eval, where all\n # the objects come from the same video under batch size 1).\n pred_masks_high_res = self._apply_non_overlapping_constraints(\n pred_masks_high_res\n )\n # scale the raw mask logits with a temperature before applying sigmoid\n binarize = self.binarize_mask_from_pts_for_mem_enc and is_mask_from_pts\n if binarize and not self.training:\n mask_for_mem = (pred_masks_high_res > 0).float()\n else:\n # apply sigmoid on the raw mask logits to turn them into range (0, 1)\n mask_for_mem = torch.sigmoid(pred_masks_high_res)\n # apply scale and bias terms to the sigmoid probabilities\n if self.sigmoid_scale_for_mem_enc != 1.0:\n mask_for_mem = mask_for_mem * self.sigmoid_scale_for_mem_enc\n if self.sigmoid_bias_for_mem_enc != 0.0:\n mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc\n maskmem_out = self.memory_encoder(\n pix_feat, mask_for_mem, skip_mask_sigmoid=True # sigmoid already applied\n )\n # Clone the feats and pos_enc to enable compilation\n maskmem_features = maskmem_out[\"vision_features\"].clone()\n maskmem_pos_enc = [m.clone() for m in maskmem_out[\"vision_pos_enc\"]]\n # add a no-object embedding to the spatial memory to indicate that the frame\n # is predicted to be occluded (i.e. no object is appearing in the frame)\n if self.no_obj_embed_spatial is not None:\n is_obj_appearing = (object_score_logits > 0).float()\n maskmem_features += (\n 1 - is_obj_appearing[..., None, None]\n ) * self.no_obj_embed_spatial[..., None, None].expand(\n *maskmem_features.shape\n )\n\n return maskmem_features, maskmem_pos_enc\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/sam2_video_predictor_legacy.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport warnings\nfrom collections import OrderedDict\n\nimport torch\n\nfrom tqdm import tqdm\n\nfrom sam2.modeling.sam2_base import NO_OBJ_SCORE, SAM2Base\nfrom sam2.utils.misc import concat_points, fill_holes_in_mask_scores, load_video_frames\n\n\nclass SAM2VideoPredictor(SAM2Base):\n \"\"\"The predictor class to handle user interactions and manage inference states.\"\"\"\n\n def __init__(\n self,\n fill_hole_area=0,\n # whether to apply non-overlapping constraints on the output object masks\n non_overlap_masks=False,\n # whether to clear non-conditioning memory of the surrounding frames (which may contain outdated information) after adding correction clicks;\n # note that this would only apply to *single-object tracking* unless `clear_non_cond_mem_for_multi_obj` is also set to True)\n clear_non_cond_mem_around_input=False,\n # whether to also clear non-conditioning memory of the surrounding frames (only effective when `clear_non_cond_mem_around_input` is True).\n clear_non_cond_mem_for_multi_obj=False,\n # if `add_all_frames_to_correct_as_cond` is True, we also append to the conditioning frame list any frame that receives a later correction click\n # if `add_all_frames_to_correct_as_cond` is False, we conditioning frame list to only use those initial conditioning frames\n add_all_frames_to_correct_as_cond=False,\n **kwargs,\n ):\n super().__init__(**kwargs)\n self.fill_hole_area = fill_hole_area\n self.non_overlap_masks = non_overlap_masks\n self.clear_non_cond_mem_around_input = clear_non_cond_mem_around_input\n self.clear_non_cond_mem_for_multi_obj = clear_non_cond_mem_for_multi_obj\n self.add_all_frames_to_correct_as_cond = add_all_frames_to_correct_as_cond\n\n @torch.inference_mode()\n def init_state(\n self,\n video_path,\n offload_video_to_cpu=False,\n offload_state_to_cpu=False,\n async_loading_frames=False,\n ):\n \"\"\"Initialize an inference state.\"\"\"\n compute_device = self.device # device of the model\n images, video_height, video_width = load_video_frames(\n video_path=video_path,\n image_size=self.image_size,\n offload_video_to_cpu=offload_video_to_cpu,\n async_loading_frames=async_loading_frames,\n compute_device=compute_device,\n )\n inference_state = {}\n inference_state[\"images\"] = images\n inference_state[\"num_frames\"] = len(images)\n # whether to offload the video frames to CPU memory\n # turning on this option saves the GPU memory with only a very small overhead\n inference_state[\"offload_video_to_cpu\"] = offload_video_to_cpu\n # whether to offload the inference state to CPU memory\n # turning on this option saves the GPU memory at the cost of a lower tracking fps\n # (e.g. in a test case of 768x768 model, fps dropped from 27 to 24 when tracking one object\n # and from 24 to 21 when tracking two objects)\n inference_state[\"offload_state_to_cpu\"] = offload_state_to_cpu\n # the original video height and width, used for resizing final output scores\n inference_state[\"video_height\"] = video_height\n inference_state[\"video_width\"] = video_width\n inference_state[\"device\"] = compute_device\n if offload_state_to_cpu:\n inference_state[\"storage_device\"] = torch.device(\"cpu\")\n else:\n inference_state[\"storage_device\"] = compute_device\n # inputs on each frame\n inference_state[\"point_inputs_per_obj\"] = {}\n inference_state[\"mask_inputs_per_obj\"] = {}\n # visual features on a small number of recently visited frames for quick interactions\n inference_state[\"cached_features\"] = {}\n # values that don't change across frames (so we only need to hold one copy of them)\n inference_state[\"constants\"] = {}\n # mapping between client-side object id and model-side object index\n inference_state[\"obj_id_to_idx\"] = OrderedDict()\n inference_state[\"obj_idx_to_id\"] = OrderedDict()\n inference_state[\"obj_ids\"] = []\n # A storage to hold the model's tracking results and states on each frame\n inference_state[\"output_dict\"] = {\n \"cond_frame_outputs\": {}, # dict containing {frame_idx: }\n \"non_cond_frame_outputs\": {}, # dict containing {frame_idx: }\n }\n # Slice (view) of each object tracking results, sharing the same memory with \"output_dict\"\n inference_state[\"output_dict_per_obj\"] = {}\n # A temporary storage to hold new outputs when user interact with a frame\n # to add clicks or mask (it's merged into \"output_dict\" before propagation starts)\n inference_state[\"temp_output_dict_per_obj\"] = {}\n # Frames that already holds consolidated outputs from click or mask inputs\n # (we directly use their consolidated outputs during tracking)\n inference_state[\"consolidated_frame_inds\"] = {\n \"cond_frame_outputs\": set(), # set containing frame indices\n \"non_cond_frame_outputs\": set(), # set containing frame indices\n }\n # metadata for each tracking frame (e.g. which direction it's tracked)\n inference_state[\"tracking_has_started\"] = False\n inference_state[\"frames_already_tracked\"] = {}\n # Warm up the visual backbone and cache the image feature on frame 0\n self._get_image_feature(inference_state, frame_idx=0, batch_size=1)\n return inference_state\n\n @classmethod\n def from_pretrained(cls, model_id: str, **kwargs) -> \"SAM2VideoPredictor\":\n \"\"\"\n Load a pretrained model from the Hugging Face hub.\n\n Arguments:\n model_id (str): The Hugging Face repository ID.\n **kwargs: Additional arguments to pass to the model constructor.\n\n Returns:\n (SAM2VideoPredictor): The loaded model.\n \"\"\"\n from sam2.build_sam import build_sam2_video_predictor_hf\n\n sam_model = build_sam2_video_predictor_hf(model_id, **kwargs)\n return sam_model\n\n def _obj_id_to_idx(self, inference_state, obj_id):\n \"\"\"Map client-side object id to model-side object index.\"\"\"\n obj_idx = inference_state[\"obj_id_to_idx\"].get(obj_id, None)\n if obj_idx is not None:\n return obj_idx\n\n # This is a new object id not sent to the server before. We only allow adding\n # new objects *before* the tracking starts.\n allow_new_object = not inference_state[\"tracking_has_started\"]\n if allow_new_object:\n # get the next object slot\n obj_idx = len(inference_state[\"obj_id_to_idx\"])\n inference_state[\"obj_id_to_idx\"][obj_id] = obj_idx\n inference_state[\"obj_idx_to_id\"][obj_idx] = obj_id\n inference_state[\"obj_ids\"] = list(inference_state[\"obj_id_to_idx\"])\n # set up input and output structures for this object\n inference_state[\"point_inputs_per_obj\"][obj_idx] = {}\n inference_state[\"mask_inputs_per_obj\"][obj_idx] = {}\n inference_state[\"output_dict_per_obj\"][obj_idx] = {\n \"cond_frame_outputs\": {}, # dict containing {frame_idx: }\n \"non_cond_frame_outputs\": {}, # dict containing {frame_idx: }\n }\n inference_state[\"temp_output_dict_per_obj\"][obj_idx] = {\n \"cond_frame_outputs\": {}, # dict containing {frame_idx: }\n \"non_cond_frame_outputs\": {}, # dict containing {frame_idx: }\n }\n return obj_idx\n else:\n raise RuntimeError(\n f\"Cannot add new object id {obj_id} after tracking starts. \"\n f\"All existing object ids: {inference_state['obj_ids']}. \"\n f\"Please call 'reset_state' to restart from scratch.\"\n )\n\n def _obj_idx_to_id(self, inference_state, obj_idx):\n \"\"\"Map model-side object index to client-side object id.\"\"\"\n return inference_state[\"obj_idx_to_id\"][obj_idx]\n\n def _get_obj_num(self, inference_state):\n \"\"\"Get the total number of unique object ids received so far in this session.\"\"\"\n return len(inference_state[\"obj_idx_to_id\"])\n\n @torch.inference_mode()\n def add_new_points_or_box(\n self,\n inference_state,\n frame_idx,\n obj_id,\n points=None,\n labels=None,\n clear_old_points=True,\n normalize_coords=True,\n box=None,\n ):\n \"\"\"Add new points to a frame.\"\"\"\n obj_idx = self._obj_id_to_idx(inference_state, obj_id)\n point_inputs_per_frame = inference_state[\"point_inputs_per_obj\"][obj_idx]\n mask_inputs_per_frame = inference_state[\"mask_inputs_per_obj\"][obj_idx]\n\n if (points is not None) != (labels is not None):\n raise ValueError(\"points and labels must be provided together\")\n if points is None and box is None:\n raise ValueError(\"at least one of points or box must be provided as input\")\n\n if points is None:\n points = torch.zeros(0, 2, dtype=torch.float32)\n elif not isinstance(points, torch.Tensor):\n points = torch.tensor(points, dtype=torch.float32)\n if labels is None:\n labels = torch.zeros(0, dtype=torch.int32)\n elif not isinstance(labels, torch.Tensor):\n labels = torch.tensor(labels, dtype=torch.int32)\n if points.dim() == 2:\n points = points.unsqueeze(0) # add batch dimension\n if labels.dim() == 1:\n labels = labels.unsqueeze(0) # add batch dimension\n\n # If `box` is provided, we add it as the first two points with labels 2 and 3\n # along with the user-provided points (consistent with how SAM 2 is trained).\n if box is not None:\n if not clear_old_points:\n raise ValueError(\n \"cannot add box without clearing old points, since \"\n \"box prompt must be provided before any point prompt \"\n \"(please use clear_old_points=True instead)\"\n )\n if inference_state[\"tracking_has_started\"]:\n warnings.warn(\n \"You are adding a box after tracking starts. SAM 2 may not always be \"\n \"able to incorporate a box prompt for *refinement*. If you intend to \"\n \"use box prompt as an *initial* input before tracking, please call \"\n \"'reset_state' on the inference state to restart from scratch.\",\n category=UserWarning,\n stacklevel=2,\n )\n if not isinstance(box, torch.Tensor):\n box = torch.tensor(box, dtype=torch.float32, device=points.device)\n box_coords = box.reshape(1, 2, 2)\n box_labels = torch.tensor([2, 3], dtype=torch.int32, device=labels.device)\n box_labels = box_labels.reshape(1, 2)\n points = torch.cat([box_coords, points], dim=1)\n labels = torch.cat([box_labels, labels], dim=1)\n\n if normalize_coords:\n video_H = inference_state[\"video_height\"]\n video_W = inference_state[\"video_width\"]\n points = points / torch.tensor([video_W, video_H]).to(points.device)\n # scale the (normalized) coordinates by the model's internal image size\n points = points * self.image_size\n points = points.to(inference_state[\"device\"])\n labels = labels.to(inference_state[\"device\"])\n\n if not clear_old_points:\n point_inputs = point_inputs_per_frame.get(frame_idx, None)\n else:\n point_inputs = None\n point_inputs = concat_points(point_inputs, points, labels)\n\n point_inputs_per_frame[frame_idx] = point_inputs\n mask_inputs_per_frame.pop(frame_idx, None)\n # If this frame hasn't been tracked before, we treat it as an initial conditioning\n # frame, meaning that the inputs points are to generate segments on this frame without\n # using any memory from other frames, like in SAM. Otherwise (if it has been tracked),\n # the input points will be used to correct the already tracked masks.\n is_init_cond_frame = frame_idx not in inference_state[\"frames_already_tracked\"]\n # whether to track in reverse time order\n if is_init_cond_frame:\n reverse = False\n else:\n reverse = inference_state[\"frames_already_tracked\"][frame_idx][\"reverse\"]\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n obj_temp_output_dict = inference_state[\"temp_output_dict_per_obj\"][obj_idx]\n # Add a frame to conditioning output if it's an initial conditioning frame or\n # if the model sees all frames receiving clicks/mask as conditioning frames.\n is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond\n storage_key = \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n\n # Get any previously predicted mask logits on this object and feed it along with\n # the new clicks into the SAM mask decoder.\n prev_sam_mask_logits = None\n # lookup temporary output dict first, which contains the most recent output\n # (if not found, then lookup conditioning and non-conditioning frame output)\n prev_out = obj_temp_output_dict[storage_key].get(frame_idx)\n if prev_out is None:\n prev_out = obj_output_dict[\"cond_frame_outputs\"].get(frame_idx)\n if prev_out is None:\n prev_out = obj_output_dict[\"non_cond_frame_outputs\"].get(frame_idx)\n\n if prev_out is not None and prev_out[\"pred_masks\"] is not None:\n device = inference_state[\"device\"]\n prev_sam_mask_logits = prev_out[\"pred_masks\"].to(device, non_blocking=True)\n # Clamp the scale of prev_sam_mask_logits to avoid rare numerical issues.\n prev_sam_mask_logits = torch.clamp(prev_sam_mask_logits, -32.0, 32.0)\n current_out, _ = self._run_single_frame_inference(\n inference_state=inference_state,\n output_dict=obj_output_dict, # run on the slice of a single object\n frame_idx=frame_idx,\n batch_size=1, # run on the slice of a single object\n is_init_cond_frame=is_init_cond_frame,\n point_inputs=point_inputs,\n mask_inputs=None,\n reverse=reverse,\n # Skip the memory encoder when adding clicks or mask. We execute the memory encoder\n # at the beginning of `propagate_in_video` (after user finalize their clicks). This\n # allows us to enforce non-overlapping constraints on all objects before encoding\n # them into memory.\n run_mem_encoder=False,\n prev_sam_mask_logits=prev_sam_mask_logits,\n )\n # Add the output to the output dict (to be used as future memory)\n obj_temp_output_dict[storage_key][frame_idx] = current_out\n\n # Resize the output mask to the original video resolution\n obj_ids = inference_state[\"obj_ids\"]\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n run_mem_encoder=False,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n return frame_idx, obj_ids, video_res_masks\n\n def add_new_points(self, *args, **kwargs):\n \"\"\"Deprecated method. Please use `add_new_points_or_box` instead.\"\"\"\n return self.add_new_points_or_box(*args, **kwargs)\n\n @torch.inference_mode()\n def add_new_mask(\n self,\n inference_state,\n frame_idx,\n obj_id,\n mask,\n ):\n \"\"\"Add new mask to a frame.\"\"\"\n obj_idx = self._obj_id_to_idx(inference_state, obj_id)\n point_inputs_per_frame = inference_state[\"point_inputs_per_obj\"][obj_idx]\n mask_inputs_per_frame = inference_state[\"mask_inputs_per_obj\"][obj_idx]\n\n if not isinstance(mask, torch.Tensor):\n mask = torch.tensor(mask, dtype=torch.bool)\n assert mask.dim() == 2\n mask_H, mask_W = mask.shape\n mask_inputs_orig = mask[None, None] # add batch and channel dimension\n mask_inputs_orig = mask_inputs_orig.float().to(inference_state[\"device\"])\n\n # resize the mask if it doesn't match the model's image size\n if mask_H != self.image_size or mask_W != self.image_size:\n mask_inputs = torch.nn.functional.interpolate(\n mask_inputs_orig,\n size=(self.image_size, self.image_size),\n align_corners=False,\n mode=\"bilinear\",\n antialias=True, # use antialias for downsampling\n )\n mask_inputs = (mask_inputs >= 0.5).float()\n else:\n mask_inputs = mask_inputs_orig\n\n mask_inputs_per_frame[frame_idx] = mask_inputs\n point_inputs_per_frame.pop(frame_idx, None)\n # If this frame hasn't been tracked before, we treat it as an initial conditioning\n # frame, meaning that the inputs points are to generate segments on this frame without\n # using any memory from other frames, like in SAM. Otherwise (if it has been tracked),\n # the input points will be used to correct the already tracked masks.\n is_init_cond_frame = frame_idx not in inference_state[\"frames_already_tracked\"]\n # whether to track in reverse time order\n if is_init_cond_frame:\n reverse = False\n else:\n reverse = inference_state[\"frames_already_tracked\"][frame_idx][\"reverse\"]\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n obj_temp_output_dict = inference_state[\"temp_output_dict_per_obj\"][obj_idx]\n # Add a frame to conditioning output if it's an initial conditioning frame or\n # if the model sees all frames receiving clicks/mask as conditioning frames.\n is_cond = is_init_cond_frame or self.add_all_frames_to_correct_as_cond\n storage_key = \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n\n current_out, _ = self._run_single_frame_inference(\n inference_state=inference_state,\n output_dict=obj_output_dict, # run on the slice of a single object\n frame_idx=frame_idx,\n batch_size=1, # run on the slice of a single object\n is_init_cond_frame=is_init_cond_frame,\n point_inputs=None,\n mask_inputs=mask_inputs,\n reverse=reverse,\n # Skip the memory encoder when adding clicks or mask. We execute the memory encoder\n # at the beginning of `propagate_in_video` (after user finalize their clicks). This\n # allows us to enforce non-overlapping constraints on all objects before encoding\n # them into memory.\n run_mem_encoder=False,\n )\n # Add the output to the output dict (to be used as future memory)\n obj_temp_output_dict[storage_key][frame_idx] = current_out\n\n # Resize the output mask to the original video resolution\n obj_ids = inference_state[\"obj_ids\"]\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n run_mem_encoder=False,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n return frame_idx, obj_ids, video_res_masks\n\n def _get_orig_video_res_output(self, inference_state, any_res_masks):\n \"\"\"\n Resize the object scores to the original video resolution (video_res_masks)\n and apply non-overlapping constraints for final output.\n \"\"\"\n device = inference_state[\"device\"]\n video_H = inference_state[\"video_height\"]\n video_W = inference_state[\"video_width\"]\n any_res_masks = any_res_masks.to(device, non_blocking=True)\n if any_res_masks.shape[-2:] == (video_H, video_W):\n video_res_masks = any_res_masks\n else:\n video_res_masks = torch.nn.functional.interpolate(\n any_res_masks,\n size=(video_H, video_W),\n mode=\"bilinear\",\n align_corners=False,\n )\n if self.non_overlap_masks:\n video_res_masks = self._apply_non_overlapping_constraints(video_res_masks)\n return any_res_masks, video_res_masks\n\n def _consolidate_temp_output_across_obj(\n self,\n inference_state,\n frame_idx,\n is_cond,\n run_mem_encoder,\n consolidate_at_video_res=False,\n ):\n \"\"\"\n Consolidate the per-object temporary outputs in `temp_output_dict_per_obj` on\n a frame into a single output for all objects, including\n 1) fill any missing objects either from `output_dict_per_obj` (if they exist in\n `output_dict_per_obj` for this frame) or leave them as placeholder values\n (if they don't exist in `output_dict_per_obj` for this frame);\n 2) if specified, rerun memory encoder after apply non-overlapping constraints\n on the object scores.\n \"\"\"\n batch_size = self._get_obj_num(inference_state)\n storage_key = \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n # Optionally, we allow consolidating the temporary outputs at the original\n # video resolution (to provide a better editing experience for mask prompts).\n if consolidate_at_video_res:\n assert not run_mem_encoder, \"memory encoder cannot run at video resolution\"\n consolidated_H = inference_state[\"video_height\"]\n consolidated_W = inference_state[\"video_width\"]\n consolidated_mask_key = \"pred_masks_video_res\"\n else:\n consolidated_H = consolidated_W = self.image_size // 4\n consolidated_mask_key = \"pred_masks\"\n\n # Initialize `consolidated_out`. Its \"maskmem_features\" and \"maskmem_pos_enc\"\n # will be added when rerunning the memory encoder after applying non-overlapping\n # constraints to object scores. Its \"pred_masks\" are prefilled with a large\n # negative value (NO_OBJ_SCORE) to represent missing objects.\n consolidated_out = {\n \"maskmem_features\": None,\n \"maskmem_pos_enc\": None,\n consolidated_mask_key: torch.full(\n size=(batch_size, 1, consolidated_H, consolidated_W),\n fill_value=NO_OBJ_SCORE,\n dtype=torch.float32,\n device=inference_state[\"storage_device\"],\n ),\n \"obj_ptr\": torch.full(\n size=(batch_size, self.hidden_dim),\n fill_value=NO_OBJ_SCORE,\n dtype=torch.float32,\n device=inference_state[\"device\"],\n ),\n \"object_score_logits\": torch.full(\n size=(batch_size, 1),\n # default to 10.0 for object_score_logits, i.e. assuming the object is\n # present as sigmoid(10)=1, same as in `predict_masks` of `MaskDecoder`\n fill_value=10.0,\n dtype=torch.float32,\n device=inference_state[\"device\"],\n ),\n }\n empty_mask_ptr = None\n for obj_idx in range(batch_size):\n obj_temp_output_dict = inference_state[\"temp_output_dict_per_obj\"][obj_idx]\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx]\n out = obj_temp_output_dict[storage_key].get(frame_idx, None)\n # If the object doesn't appear in \"temp_output_dict_per_obj\" on this frame,\n # we fall back and look up its previous output in \"output_dict_per_obj\".\n # We look up both \"cond_frame_outputs\" and \"non_cond_frame_outputs\" in\n # \"output_dict_per_obj\" to find a previous output for this object.\n if out is None:\n out = obj_output_dict[\"cond_frame_outputs\"].get(frame_idx, None)\n if out is None:\n out = obj_output_dict[\"non_cond_frame_outputs\"].get(frame_idx, None)\n # If the object doesn't appear in \"output_dict_per_obj\" either, we skip it\n # and leave its mask scores to the default scores (i.e. the NO_OBJ_SCORE\n # placeholder above) and set its object pointer to be a dummy pointer.\n if out is None:\n # Fill in dummy object pointers for those objects without any inputs or\n # tracking outcomes on this frame (only do it under `run_mem_encoder=True`,\n # i.e. when we need to build the memory for tracking).\n if run_mem_encoder:\n if empty_mask_ptr is None:\n empty_mask_ptr = self._get_empty_mask_ptr(\n inference_state, frame_idx\n )\n # fill object pointer with a dummy pointer (based on an empty mask)\n consolidated_out[\"obj_ptr\"][obj_idx : obj_idx + 1] = empty_mask_ptr\n continue\n # Add the temporary object output mask to consolidated output mask\n obj_mask = out[\"pred_masks\"]\n consolidated_pred_masks = consolidated_out[consolidated_mask_key]\n if obj_mask.shape[-2:] == consolidated_pred_masks.shape[-2:]:\n consolidated_pred_masks[obj_idx : obj_idx + 1] = obj_mask\n else:\n # Resize first if temporary object mask has a different resolution\n resized_obj_mask = torch.nn.functional.interpolate(\n obj_mask,\n size=consolidated_pred_masks.shape[-2:],\n mode=\"bilinear\",\n align_corners=False,\n )\n consolidated_pred_masks[obj_idx : obj_idx + 1] = resized_obj_mask\n consolidated_out[\"obj_ptr\"][obj_idx : obj_idx + 1] = out[\"obj_ptr\"]\n consolidated_out[\"object_score_logits\"][obj_idx : obj_idx + 1] = out[\n \"object_score_logits\"\n ]\n\n # Optionally, apply non-overlapping constraints on the consolidated scores\n # and rerun the memory encoder\n if run_mem_encoder:\n device = inference_state[\"device\"]\n high_res_masks = torch.nn.functional.interpolate(\n consolidated_out[\"pred_masks\"].to(device, non_blocking=True),\n size=(self.image_size, self.image_size),\n mode=\"bilinear\",\n align_corners=False,\n )\n if self.non_overlap_masks_for_mem_enc:\n high_res_masks = self._apply_non_overlapping_constraints(high_res_masks)\n maskmem_features, maskmem_pos_enc = self._run_memory_encoder(\n inference_state=inference_state,\n frame_idx=frame_idx,\n batch_size=batch_size,\n high_res_masks=high_res_masks,\n object_score_logits=consolidated_out[\"object_score_logits\"],\n is_mask_from_pts=True, # these frames are what the user interacted with\n )\n consolidated_out[\"maskmem_features\"] = maskmem_features\n consolidated_out[\"maskmem_pos_enc\"] = maskmem_pos_enc\n\n return consolidated_out\n\n def _get_empty_mask_ptr(self, inference_state, frame_idx):\n \"\"\"Get a dummy object pointer based on an empty mask on the current frame.\"\"\"\n # A dummy (empty) mask with a single object\n batch_size = 1\n mask_inputs = torch.zeros(\n (batch_size, 1, self.image_size, self.image_size),\n dtype=torch.float32,\n device=inference_state[\"device\"],\n )\n\n # Retrieve correct image features\n (\n _,\n _,\n current_vision_feats,\n current_vision_pos_embeds,\n feat_sizes,\n ) = self._get_image_feature(inference_state, frame_idx, batch_size)\n\n # Feed the empty mask and image feature above to get a dummy object pointer\n current_out = self.track_step(\n frame_idx=frame_idx,\n is_init_cond_frame=True,\n current_vision_feats=current_vision_feats,\n current_vision_pos_embeds=current_vision_pos_embeds,\n feat_sizes=feat_sizes,\n point_inputs=None,\n mask_inputs=mask_inputs,\n output_dict={},\n num_frames=inference_state[\"num_frames\"],\n track_in_reverse=False,\n run_mem_encoder=False,\n prev_sam_mask_logits=None,\n )\n return current_out[\"obj_ptr\"]\n\n @torch.inference_mode()\n def propagate_in_video_preflight(self, inference_state):\n \"\"\"Prepare inference_state and consolidate temporary outputs before tracking.\"\"\"\n # Tracking has started and we don't allow adding new objects until session is reset.\n inference_state[\"tracking_has_started\"] = True\n batch_size = self._get_obj_num(inference_state)\n\n # Consolidate per-object temporary outputs in \"temp_output_dict_per_obj\" and\n # add them into \"output_dict\".\n temp_output_dict_per_obj = inference_state[\"temp_output_dict_per_obj\"]\n output_dict = inference_state[\"output_dict\"]\n # \"consolidated_frame_inds\" contains indices of those frames where consolidated\n # temporary outputs have been added (either in this call or any previous calls\n # to `propagate_in_video_preflight`).\n consolidated_frame_inds = inference_state[\"consolidated_frame_inds\"]\n for is_cond in [False, True]:\n # Separately consolidate conditioning and non-conditioning temp outputs\n storage_key = \"cond_frame_outputs\" if is_cond else \"non_cond_frame_outputs\"\n # Find all the frames that contain temporary outputs for any objects\n # (these should be the frames that have just received clicks for mask inputs\n # via `add_new_points_or_box` or `add_new_mask`)\n temp_frame_inds = set()\n for obj_temp_output_dict in temp_output_dict_per_obj.values():\n temp_frame_inds.update(obj_temp_output_dict[storage_key].keys())\n consolidated_frame_inds[storage_key].update(temp_frame_inds)\n # consolidate the temporary output across all objects on this frame\n for frame_idx in temp_frame_inds:\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state, frame_idx, is_cond=is_cond, run_mem_encoder=True\n )\n # merge them into \"output_dict\" and also create per-object slices\n output_dict[storage_key][frame_idx] = consolidated_out\n self._add_output_per_object(\n inference_state, frame_idx, consolidated_out, storage_key\n )\n clear_non_cond_mem = self.clear_non_cond_mem_around_input and (\n self.clear_non_cond_mem_for_multi_obj or batch_size <= 1\n )\n if clear_non_cond_mem:\n # clear non-conditioning memory of the surrounding frames\n self._clear_non_cond_mem_around_input(inference_state, frame_idx)\n\n # clear temporary outputs in `temp_output_dict_per_obj`\n for obj_temp_output_dict in temp_output_dict_per_obj.values():\n obj_temp_output_dict[storage_key].clear()\n\n # edge case: if an output is added to \"cond_frame_outputs\", we remove any prior\n # output on the same frame in \"non_cond_frame_outputs\"\n for frame_idx in output_dict[\"cond_frame_outputs\"]:\n output_dict[\"non_cond_frame_outputs\"].pop(frame_idx, None)\n for obj_output_dict in inference_state[\"output_dict_per_obj\"].values():\n for frame_idx in obj_output_dict[\"cond_frame_outputs\"]:\n obj_output_dict[\"non_cond_frame_outputs\"].pop(frame_idx, None)\n for frame_idx in consolidated_frame_inds[\"cond_frame_outputs\"]:\n assert frame_idx in output_dict[\"cond_frame_outputs\"]\n consolidated_frame_inds[\"non_cond_frame_outputs\"].discard(frame_idx)\n\n # Make sure that the frame indices in \"consolidated_frame_inds\" are exactly those frames\n # with either points or mask inputs (which should be true under a correct workflow).\n all_consolidated_frame_inds = (\n consolidated_frame_inds[\"cond_frame_outputs\"]\n | consolidated_frame_inds[\"non_cond_frame_outputs\"]\n )\n input_frames_inds = set()\n for point_inputs_per_frame in inference_state[\"point_inputs_per_obj\"].values():\n input_frames_inds.update(point_inputs_per_frame.keys())\n for mask_inputs_per_frame in inference_state[\"mask_inputs_per_obj\"].values():\n input_frames_inds.update(mask_inputs_per_frame.keys())\n assert all_consolidated_frame_inds == input_frames_inds\n\n @torch.inference_mode()\n def propagate_in_video(\n self,\n inference_state,\n start_frame_idx=None,\n max_frame_num_to_track=None,\n reverse=False,\n ):\n \"\"\"Propagate the input points across frames to track in the entire video.\"\"\"\n self.propagate_in_video_preflight(inference_state)\n\n output_dict = inference_state[\"output_dict\"]\n consolidated_frame_inds = inference_state[\"consolidated_frame_inds\"]\n obj_ids = inference_state[\"obj_ids\"]\n num_frames = inference_state[\"num_frames\"]\n batch_size = self._get_obj_num(inference_state)\n if len(output_dict[\"cond_frame_outputs\"]) == 0:\n raise RuntimeError(\"No points are provided; please add points first\")\n clear_non_cond_mem = self.clear_non_cond_mem_around_input and (\n self.clear_non_cond_mem_for_multi_obj or batch_size <= 1\n )\n\n # set start index, end index, and processing order\n if start_frame_idx is None:\n # default: start from the earliest frame with input points\n start_frame_idx = min(output_dict[\"cond_frame_outputs\"])\n if max_frame_num_to_track is None:\n # default: track all the frames in the video\n max_frame_num_to_track = num_frames\n if reverse:\n end_frame_idx = max(start_frame_idx - max_frame_num_to_track, 0)\n if start_frame_idx > 0:\n processing_order = range(start_frame_idx, end_frame_idx - 1, -1)\n else:\n processing_order = [] # skip reverse tracking if starting from frame 0\n else:\n end_frame_idx = min(\n start_frame_idx + max_frame_num_to_track, num_frames - 1\n )\n processing_order = range(start_frame_idx, end_frame_idx + 1)\n\n for frame_idx in tqdm(processing_order, desc=\"propagate in video\"):\n # We skip those frames already in consolidated outputs (these are frames\n # that received input clicks or mask). Note that we cannot directly run\n # batched forward on them via `_run_single_frame_inference` because the\n # number of clicks on each object might be different.\n if frame_idx in consolidated_frame_inds[\"cond_frame_outputs\"]:\n storage_key = \"cond_frame_outputs\"\n current_out = output_dict[storage_key][frame_idx]\n pred_masks = current_out[\"pred_masks\"]\n if clear_non_cond_mem:\n # clear non-conditioning memory of the surrounding frames\n self._clear_non_cond_mem_around_input(inference_state, frame_idx)\n elif frame_idx in consolidated_frame_inds[\"non_cond_frame_outputs\"]:\n storage_key = \"non_cond_frame_outputs\"\n current_out = output_dict[storage_key][frame_idx]\n pred_masks = current_out[\"pred_masks\"]\n else:\n storage_key = \"non_cond_frame_outputs\"\n current_out, pred_masks = self._run_single_frame_inference(\n inference_state=inference_state,\n output_dict=output_dict,\n frame_idx=frame_idx,\n batch_size=batch_size,\n is_init_cond_frame=False,\n point_inputs=None,\n mask_inputs=None,\n reverse=reverse,\n run_mem_encoder=True,\n )\n output_dict[storage_key][frame_idx] = current_out\n # Create slices of per-object outputs for subsequent interaction with each\n # individual object after tracking.\n self._add_output_per_object(\n inference_state, frame_idx, current_out, storage_key\n )\n inference_state[\"frames_already_tracked\"][frame_idx] = {\"reverse\": reverse}\n\n # Resize the output mask to the original video resolution (we directly use\n # the mask scores on GPU for output to avoid any CPU conversion in between)\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, pred_masks\n )\n yield frame_idx, obj_ids, video_res_masks\n\n def _add_output_per_object(\n self, inference_state, frame_idx, current_out, storage_key\n ):\n \"\"\"\n Split a multi-object output into per-object output slices and add them into\n `output_dict_per_obj`. The resulting slices share the same tensor storage.\n \"\"\"\n maskmem_features = current_out[\"maskmem_features\"]\n assert maskmem_features is None or isinstance(maskmem_features, torch.Tensor)\n\n maskmem_pos_enc = current_out[\"maskmem_pos_enc\"]\n assert maskmem_pos_enc is None or isinstance(maskmem_pos_enc, list)\n\n output_dict_per_obj = inference_state[\"output_dict_per_obj\"]\n for obj_idx, obj_output_dict in output_dict_per_obj.items():\n obj_slice = slice(obj_idx, obj_idx + 1)\n obj_out = {\n \"maskmem_features\": None,\n \"maskmem_pos_enc\": None,\n \"pred_masks\": current_out[\"pred_masks\"][obj_slice],\n \"obj_ptr\": current_out[\"obj_ptr\"][obj_slice],\n \"object_score_logits\": current_out[\"object_score_logits\"][obj_slice],\n }\n if maskmem_features is not None:\n obj_out[\"maskmem_features\"] = maskmem_features[obj_slice]\n if maskmem_pos_enc is not None:\n obj_out[\"maskmem_pos_enc\"] = [x[obj_slice] for x in maskmem_pos_enc]\n obj_output_dict[storage_key][frame_idx] = obj_out\n\n @torch.inference_mode()\n def clear_all_prompts_in_frame(\n self, inference_state, frame_idx, obj_id, need_output=True\n ):\n \"\"\"Remove all input points or mask in a specific frame for a given object.\"\"\"\n obj_idx = self._obj_id_to_idx(inference_state, obj_id)\n\n # Clear the conditioning information on the given frame\n inference_state[\"point_inputs_per_obj\"][obj_idx].pop(frame_idx, None)\n inference_state[\"mask_inputs_per_obj\"][obj_idx].pop(frame_idx, None)\n\n temp_output_dict_per_obj = inference_state[\"temp_output_dict_per_obj\"]\n temp_output_dict_per_obj[obj_idx][\"cond_frame_outputs\"].pop(frame_idx, None)\n temp_output_dict_per_obj[obj_idx][\"non_cond_frame_outputs\"].pop(frame_idx, None)\n\n # Check and see if there are still any inputs left on this frame\n batch_size = self._get_obj_num(inference_state)\n frame_has_input = False\n for obj_idx2 in range(batch_size):\n if frame_idx in inference_state[\"point_inputs_per_obj\"][obj_idx2]:\n frame_has_input = True\n break\n if frame_idx in inference_state[\"mask_inputs_per_obj\"][obj_idx2]:\n frame_has_input = True\n break\n\n # If this frame has no remaining inputs for any objects, we further clear its\n # conditioning frame status\n if not frame_has_input:\n output_dict = inference_state[\"output_dict\"]\n consolidated_frame_inds = inference_state[\"consolidated_frame_inds\"]\n consolidated_frame_inds[\"cond_frame_outputs\"].discard(frame_idx)\n consolidated_frame_inds[\"non_cond_frame_outputs\"].discard(frame_idx)\n # Remove the frame's conditioning output (possibly downgrading it to non-conditioning)\n out = output_dict[\"cond_frame_outputs\"].pop(frame_idx, None)\n if out is not None:\n # The frame is not a conditioning frame anymore since it's not receiving inputs,\n # so we \"downgrade\" its output (if exists) to a non-conditioning frame output.\n output_dict[\"non_cond_frame_outputs\"][frame_idx] = out\n inference_state[\"frames_already_tracked\"].pop(frame_idx, None)\n # Similarly, do it for the sliced output on each object.\n for obj_idx2 in range(batch_size):\n obj_output_dict = inference_state[\"output_dict_per_obj\"][obj_idx2]\n obj_out = obj_output_dict[\"cond_frame_outputs\"].pop(frame_idx, None)\n if obj_out is not None:\n obj_output_dict[\"non_cond_frame_outputs\"][frame_idx] = obj_out\n\n # If all the conditioning frames have been removed, we also clear the tracking outputs\n if len(output_dict[\"cond_frame_outputs\"]) == 0:\n self._reset_tracking_results(inference_state)\n\n if not need_output:\n return\n # Finally, output updated masks per object (after removing the inputs above)\n obj_ids = inference_state[\"obj_ids\"]\n is_cond = any(\n frame_idx in obj_temp_output_dict[\"cond_frame_outputs\"]\n for obj_temp_output_dict in temp_output_dict_per_obj.values()\n )\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n run_mem_encoder=False,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n return frame_idx, obj_ids, video_res_masks\n\n @torch.inference_mode()\n def reset_state(self, inference_state):\n \"\"\"Remove all input points or mask in all frames throughout the video.\"\"\"\n self._reset_tracking_results(inference_state)\n # Remove all object ids\n inference_state[\"obj_id_to_idx\"].clear()\n inference_state[\"obj_idx_to_id\"].clear()\n inference_state[\"obj_ids\"].clear()\n inference_state[\"point_inputs_per_obj\"].clear()\n inference_state[\"mask_inputs_per_obj\"].clear()\n inference_state[\"output_dict_per_obj\"].clear()\n inference_state[\"temp_output_dict_per_obj\"].clear()\n\n def _reset_tracking_results(self, inference_state):\n \"\"\"Reset all tracking inputs and results across the videos.\"\"\"\n for v in inference_state[\"point_inputs_per_obj\"].values():\n v.clear()\n for v in inference_state[\"mask_inputs_per_obj\"].values():\n v.clear()\n for v in inference_state[\"output_dict_per_obj\"].values():\n v[\"cond_frame_outputs\"].clear()\n v[\"non_cond_frame_outputs\"].clear()\n for v in inference_state[\"temp_output_dict_per_obj\"].values():\n v[\"cond_frame_outputs\"].clear()\n v[\"non_cond_frame_outputs\"].clear()\n inference_state[\"output_dict\"][\"cond_frame_outputs\"].clear()\n inference_state[\"output_dict\"][\"non_cond_frame_outputs\"].clear()\n inference_state[\"consolidated_frame_inds\"][\"cond_frame_outputs\"].clear()\n inference_state[\"consolidated_frame_inds\"][\"non_cond_frame_outputs\"].clear()\n inference_state[\"tracking_has_started\"] = False\n inference_state[\"frames_already_tracked\"].clear()\n\n def _get_image_feature(self, inference_state, frame_idx, batch_size):\n \"\"\"Compute the image features on a given frame.\"\"\"\n # Look up in the cache first\n image, backbone_out = inference_state[\"cached_features\"].get(\n frame_idx, (None, None)\n )\n if backbone_out is None:\n # Cache miss -- we will run inference on a single image\n device = inference_state[\"device\"]\n image = inference_state[\"images\"][frame_idx].to(device).float().unsqueeze(0)\n backbone_out = self.forward_image(image)\n # Cache the most recent frame's feature (for repeated interactions with\n # a frame; we can use an LRU cache for more frames in the future).\n inference_state[\"cached_features\"] = {frame_idx: (image, backbone_out)}\n\n # expand the features to have the same dimension as the number of objects\n expanded_image = image.expand(batch_size, -1, -1, -1)\n expanded_backbone_out = {\n \"backbone_fpn\": backbone_out[\"backbone_fpn\"].copy(),\n \"vision_pos_enc\": backbone_out[\"vision_pos_enc\"].copy(),\n }\n for i, feat in enumerate(expanded_backbone_out[\"backbone_fpn\"]):\n expanded_backbone_out[\"backbone_fpn\"][i] = feat.expand(\n batch_size, -1, -1, -1\n )\n for i, pos in enumerate(expanded_backbone_out[\"vision_pos_enc\"]):\n pos = pos.expand(batch_size, -1, -1, -1)\n expanded_backbone_out[\"vision_pos_enc\"][i] = pos\n\n features = self._prepare_backbone_features(expanded_backbone_out)\n features = (expanded_image,) + features\n return features\n\n def _run_single_frame_inference(\n self,\n inference_state,\n output_dict,\n frame_idx,\n batch_size,\n is_init_cond_frame,\n point_inputs,\n mask_inputs,\n reverse,\n run_mem_encoder,\n prev_sam_mask_logits=None,\n ):\n \"\"\"Run tracking on a single frame based on current inputs and previous memory.\"\"\"\n # Retrieve correct image features\n (\n _,\n _,\n current_vision_feats,\n current_vision_pos_embeds,\n feat_sizes,\n ) = self._get_image_feature(inference_state, frame_idx, batch_size)\n\n # point and mask should not appear as input simultaneously on the same frame\n assert point_inputs is None or mask_inputs is None\n current_out = self.track_step(\n frame_idx=frame_idx,\n is_init_cond_frame=is_init_cond_frame,\n current_vision_feats=current_vision_feats,\n current_vision_pos_embeds=current_vision_pos_embeds,\n feat_sizes=feat_sizes,\n point_inputs=point_inputs,\n mask_inputs=mask_inputs,\n output_dict=output_dict,\n num_frames=inference_state[\"num_frames\"],\n track_in_reverse=reverse,\n run_mem_encoder=run_mem_encoder,\n prev_sam_mask_logits=prev_sam_mask_logits,\n )\n\n # optionally offload the output to CPU memory to save GPU space\n storage_device = inference_state[\"storage_device\"]\n maskmem_features = current_out[\"maskmem_features\"]\n if maskmem_features is not None:\n maskmem_features = maskmem_features.to(torch.bfloat16)\n maskmem_features = maskmem_features.to(storage_device, non_blocking=True)\n pred_masks_gpu = current_out[\"pred_masks\"]\n # potentially fill holes in the predicted masks\n if self.fill_hole_area > 0:\n pred_masks_gpu = fill_holes_in_mask_scores(\n pred_masks_gpu, self.fill_hole_area\n )\n pred_masks = pred_masks_gpu.to(storage_device, non_blocking=True)\n # \"maskmem_pos_enc\" is the same across frames, so we only need to store one copy of it\n maskmem_pos_enc = self._get_maskmem_pos_enc(inference_state, current_out)\n # object pointer is a small tensor, so we always keep it on GPU memory for fast access\n obj_ptr = current_out[\"obj_ptr\"]\n object_score_logits = current_out[\"object_score_logits\"]\n # make a compact version of this frame's output to reduce the state size\n compact_current_out = {\n \"maskmem_features\": maskmem_features,\n \"maskmem_pos_enc\": maskmem_pos_enc,\n \"pred_masks\": pred_masks,\n \"obj_ptr\": obj_ptr,\n \"object_score_logits\": object_score_logits,\n }\n return compact_current_out, pred_masks_gpu\n\n def _run_memory_encoder(\n self,\n inference_state,\n frame_idx,\n batch_size,\n high_res_masks,\n object_score_logits,\n is_mask_from_pts,\n ):\n \"\"\"\n Run the memory encoder on `high_res_masks`. This is usually after applying\n non-overlapping constraints to object scores. Since their scores changed, their\n memory also need to be computed again with the memory encoder.\n \"\"\"\n # Retrieve correct image features\n _, _, current_vision_feats, _, feat_sizes = self._get_image_feature(\n inference_state, frame_idx, batch_size\n )\n maskmem_features, maskmem_pos_enc = self._encode_new_memory(\n current_vision_feats=current_vision_feats,\n feat_sizes=feat_sizes,\n pred_masks_high_res=high_res_masks,\n object_score_logits=object_score_logits,\n is_mask_from_pts=is_mask_from_pts,\n )\n\n # optionally offload the output to CPU memory to save GPU space\n storage_device = inference_state[\"storage_device\"]\n maskmem_features = maskmem_features.to(torch.bfloat16)\n maskmem_features = maskmem_features.to(storage_device, non_blocking=True)\n # \"maskmem_pos_enc\" is the same across frames, so we only need to store one copy of it\n maskmem_pos_enc = self._get_maskmem_pos_enc(\n inference_state, {\"maskmem_pos_enc\": maskmem_pos_enc}\n )\n return maskmem_features, maskmem_pos_enc\n\n def _get_maskmem_pos_enc(self, inference_state, current_out):\n \"\"\"\n `maskmem_pos_enc` is the same across frames and objects, so we cache it as\n a constant in the inference session to reduce session storage size.\n \"\"\"\n model_constants = inference_state[\"constants\"]\n # \"out_maskmem_pos_enc\" should be either a list of tensors or None\n out_maskmem_pos_enc = current_out[\"maskmem_pos_enc\"]\n if out_maskmem_pos_enc is not None:\n if \"maskmem_pos_enc\" not in model_constants:\n assert isinstance(out_maskmem_pos_enc, list)\n # only take the slice for one object, since it's same across objects\n maskmem_pos_enc = [x[0:1].clone() for x in out_maskmem_pos_enc]\n model_constants[\"maskmem_pos_enc\"] = maskmem_pos_enc\n else:\n maskmem_pos_enc = model_constants[\"maskmem_pos_enc\"]\n # expand the cached maskmem_pos_enc to the actual batch size\n batch_size = out_maskmem_pos_enc[0].size(0)\n expanded_maskmem_pos_enc = [\n x.expand(batch_size, -1, -1, -1) for x in maskmem_pos_enc\n ]\n else:\n expanded_maskmem_pos_enc = None\n return expanded_maskmem_pos_enc\n\n @torch.inference_mode()\n def remove_object(self, inference_state, obj_id, strict=False, need_output=True):\n \"\"\"\n Remove an object id from the tracking state. If strict is True, we check whether\n the object id actually exists and raise an error if it doesn't exist.\n \"\"\"\n old_obj_idx_to_rm = inference_state[\"obj_id_to_idx\"].get(obj_id, None)\n updated_frames = []\n # Check whether this object_id to remove actually exists and possibly raise an error.\n if old_obj_idx_to_rm is None:\n if not strict:\n return inference_state[\"obj_ids\"], updated_frames\n raise RuntimeError(\n f\"Cannot remove object id {obj_id} as it doesn't exist. \"\n f\"All existing object ids: {inference_state['obj_ids']}.\"\n )\n\n # If this is the only remaining object id, we simply reset the state.\n if len(inference_state[\"obj_id_to_idx\"]) == 1:\n self.reset_state(inference_state)\n return inference_state[\"obj_ids\"], updated_frames\n\n # There are still remaining objects after removing this object id. In this case,\n # we need to delete the object storage from inference state tensors.\n # Step 0: clear the input on those frames where this object id has point or mask input\n # (note that this step is required as it might downgrade conditioning frames to\n # non-conditioning ones)\n obj_input_frames_inds = set()\n obj_input_frames_inds.update(\n inference_state[\"point_inputs_per_obj\"][old_obj_idx_to_rm]\n )\n obj_input_frames_inds.update(\n inference_state[\"mask_inputs_per_obj\"][old_obj_idx_to_rm]\n )\n for frame_idx in obj_input_frames_inds:\n self.clear_all_prompts_in_frame(\n inference_state, frame_idx, obj_id, need_output=False\n )\n\n # Step 1: Update the object id mapping (note that it must be done after Step 0,\n # since Step 0 still requires the old object id mappings in inference_state)\n old_obj_ids = inference_state[\"obj_ids\"]\n old_obj_inds = list(range(len(old_obj_ids)))\n remain_old_obj_inds = old_obj_inds.copy()\n remain_old_obj_inds.remove(old_obj_idx_to_rm)\n new_obj_ids = [old_obj_ids[old_idx] for old_idx in remain_old_obj_inds]\n new_obj_inds = list(range(len(new_obj_ids)))\n # build new mappings\n old_idx_to_new_idx = dict(zip(remain_old_obj_inds, new_obj_inds))\n inference_state[\"obj_id_to_idx\"] = dict(zip(new_obj_ids, new_obj_inds))\n inference_state[\"obj_idx_to_id\"] = dict(zip(new_obj_inds, new_obj_ids))\n inference_state[\"obj_ids\"] = new_obj_ids\n\n # Step 2: For per-object tensor storage, we shift their obj_idx in the dict keys.\n # (note that \"consolidated_frame_inds\" doesn't need to be updated in this step as\n # it's already handled in Step 0)\n def _map_keys(container):\n new_kvs = []\n for k in old_obj_inds:\n v = container.pop(k)\n if k in old_idx_to_new_idx:\n new_kvs.append((old_idx_to_new_idx[k], v))\n container.update(new_kvs)\n\n _map_keys(inference_state[\"point_inputs_per_obj\"])\n _map_keys(inference_state[\"mask_inputs_per_obj\"])\n _map_keys(inference_state[\"output_dict_per_obj\"])\n _map_keys(inference_state[\"temp_output_dict_per_obj\"])\n\n # Step 3: For packed tensor storage, we index the remaining ids and rebuild the per-object slices.\n def _slice_state(output_dict, storage_key):\n for frame_idx, out in output_dict[storage_key].items():\n out[\"maskmem_features\"] = out[\"maskmem_features\"][remain_old_obj_inds]\n out[\"maskmem_pos_enc\"] = [\n x[remain_old_obj_inds] for x in out[\"maskmem_pos_enc\"]\n ]\n # \"maskmem_pos_enc\" is the same across frames, so we only need to store one copy of it\n out[\"maskmem_pos_enc\"] = self._get_maskmem_pos_enc(inference_state, out)\n out[\"pred_masks\"] = out[\"pred_masks\"][remain_old_obj_inds]\n out[\"obj_ptr\"] = out[\"obj_ptr\"][remain_old_obj_inds]\n out[\"object_score_logits\"] = out[\"object_score_logits\"][\n remain_old_obj_inds\n ]\n # also update the per-object slices\n self._add_output_per_object(\n inference_state, frame_idx, out, storage_key\n )\n\n _slice_state(inference_state[\"output_dict\"], \"cond_frame_outputs\")\n _slice_state(inference_state[\"output_dict\"], \"non_cond_frame_outputs\")\n\n # Step 4: Further collect the outputs on those frames in `obj_input_frames_inds`, which\n # could show an updated mask for objects previously occluded by the object being removed\n if need_output:\n temp_output_dict_per_obj = inference_state[\"temp_output_dict_per_obj\"]\n for frame_idx in obj_input_frames_inds:\n is_cond = any(\n frame_idx in obj_temp_output_dict[\"cond_frame_outputs\"]\n for obj_temp_output_dict in temp_output_dict_per_obj.values()\n )\n consolidated_out = self._consolidate_temp_output_across_obj(\n inference_state,\n frame_idx,\n is_cond=is_cond,\n run_mem_encoder=False,\n consolidate_at_video_res=True,\n )\n _, video_res_masks = self._get_orig_video_res_output(\n inference_state, consolidated_out[\"pred_masks_video_res\"]\n )\n updated_frames.append((frame_idx, video_res_masks))\n\n return inference_state[\"obj_ids\"], updated_frames\n\n def _clear_non_cond_mem_around_input(self, inference_state, frame_idx):\n \"\"\"\n Remove the non-conditioning memory around the input frame. When users provide\n correction clicks, the surrounding frames' non-conditioning memories can still\n contain outdated object appearance information and could confuse the model.\n\n This method clears those non-conditioning memories surrounding the interacted\n frame to avoid giving the model both old and new information about the object.\n \"\"\"\n r = self.memory_temporal_stride_for_eval\n frame_idx_begin = frame_idx - r * self.num_maskmem\n frame_idx_end = frame_idx + r * self.num_maskmem\n output_dict = inference_state[\"output_dict\"]\n non_cond_frame_outputs = output_dict[\"non_cond_frame_outputs\"]\n for t in range(frame_idx_begin, frame_idx_end + 1):\n non_cond_frame_outputs.pop(t, None)\n for obj_output_dict in inference_state[\"output_dict_per_obj\"].values():\n obj_output_dict[\"non_cond_frame_outputs\"].pop(t, None)\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/utils/__init__.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/utils/amg.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport math\nfrom copy import deepcopy\nfrom itertools import product\nfrom typing import Any, Dict, Generator, ItemsView, List, Tuple\n\nimport numpy as np\nimport torch\n\n# Very lightly adapted from https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/amg.py\n\n\nclass MaskData:\n \"\"\"\n A structure for storing masks and their related data in batched format.\n Implements basic filtering and concatenation.\n \"\"\"\n\n def __init__(self, **kwargs) -> None:\n for v in kwargs.values():\n assert isinstance(\n v, (list, np.ndarray, torch.Tensor)\n ), \"MaskData only supports list, numpy arrays, and torch tensors.\"\n self._stats = dict(**kwargs)\n\n def __setitem__(self, key: str, item: Any) -> None:\n assert isinstance(\n item, (list, np.ndarray, torch.Tensor)\n ), \"MaskData only supports list, numpy arrays, and torch tensors.\"\n self._stats[key] = item\n\n def __delitem__(self, key: str) -> None:\n del self._stats[key]\n\n def __getitem__(self, key: str) -> Any:\n return self._stats[key]\n\n def items(self) -> ItemsView[str, Any]:\n return self._stats.items()\n\n def filter(self, keep: torch.Tensor) -> None:\n for k, v in self._stats.items():\n if v is None:\n self._stats[k] = None\n elif isinstance(v, torch.Tensor):\n self._stats[k] = v[torch.as_tensor(keep, device=v.device)]\n elif isinstance(v, np.ndarray):\n self._stats[k] = v[keep.detach().cpu().numpy()]\n elif isinstance(v, list) and keep.dtype == torch.bool:\n self._stats[k] = [a for i, a in enumerate(v) if keep[i]]\n elif isinstance(v, list):\n self._stats[k] = [v[i] for i in keep]\n else:\n raise TypeError(f\"MaskData key {k} has an unsupported type {type(v)}.\")\n\n def cat(self, new_stats: \"MaskData\") -> None:\n for k, v in new_stats.items():\n if k not in self._stats or self._stats[k] is None:\n self._stats[k] = deepcopy(v)\n elif isinstance(v, torch.Tensor):\n self._stats[k] = torch.cat([self._stats[k], v], dim=0)\n elif isinstance(v, np.ndarray):\n self._stats[k] = np.concatenate([self._stats[k], v], axis=0)\n elif isinstance(v, list):\n self._stats[k] = self._stats[k] + deepcopy(v)\n else:\n raise TypeError(f\"MaskData key {k} has an unsupported type {type(v)}.\")\n\n def to_numpy(self) -> None:\n for k, v in self._stats.items():\n if isinstance(v, torch.Tensor):\n self._stats[k] = v.float().detach().cpu().numpy()\n\n\ndef is_box_near_crop_edge(\n boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0\n) -> torch.Tensor:\n \"\"\"Filter masks at the edge of a crop, but not at the edge of the original image.\"\"\"\n crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device)\n orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device)\n boxes = uncrop_boxes_xyxy(boxes, crop_box).float()\n near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0)\n near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0)\n near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge)\n return torch.any(near_crop_edge, dim=1)\n\n\ndef box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor:\n box_xywh = deepcopy(box_xyxy)\n box_xywh[2] = box_xywh[2] - box_xywh[0]\n box_xywh[3] = box_xywh[3] - box_xywh[1]\n return box_xywh\n\n\ndef batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:\n assert len(args) > 0 and all(\n len(a) == len(args[0]) for a in args\n ), \"Batched iteration must have inputs of all the same size.\"\n n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)\n for b in range(n_batches):\n yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]\n\n\ndef mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]:\n \"\"\"\n Encodes masks to an uncompressed RLE, in the format expected by\n pycoco tools.\n \"\"\"\n # Put in fortran order and flatten h,w\n b, h, w = tensor.shape\n tensor = tensor.permute(0, 2, 1).flatten(1)\n\n # Compute change indices\n diff = tensor[:, 1:] ^ tensor[:, :-1]\n change_indices = diff.nonzero()\n\n # Encode run length\n out = []\n for i in range(b):\n cur_idxs = change_indices[change_indices[:, 0] == i, 1]\n cur_idxs = torch.cat(\n [\n torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device),\n cur_idxs + 1,\n torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device),\n ]\n )\n btw_idxs = cur_idxs[1:] - cur_idxs[:-1]\n counts = [] if tensor[i, 0] == 0 else [0]\n counts.extend(btw_idxs.detach().cpu().tolist())\n out.append({\"size\": [h, w], \"counts\": counts})\n return out\n\n\ndef rle_to_mask(rle: Dict[str, Any]) -> np.ndarray:\n \"\"\"Compute a binary mask from an uncompressed RLE.\"\"\"\n h, w = rle[\"size\"]\n mask = np.empty(h * w, dtype=bool)\n idx = 0\n parity = False\n for count in rle[\"counts\"]:\n mask[idx : idx + count] = parity\n idx += count\n parity ^= True\n mask = mask.reshape(w, h)\n return mask.transpose() # Put in C order\n\n\ndef area_from_rle(rle: Dict[str, Any]) -> int:\n return sum(rle[\"counts\"][1::2])\n\n\ndef calculate_stability_score(\n masks: torch.Tensor, mask_threshold: float, threshold_offset: float\n) -> torch.Tensor:\n \"\"\"\n Computes the stability score for a batch of masks. The stability\n score is the IoU between the binary masks obtained by thresholding\n the predicted mask logits at high and low values.\n \"\"\"\n # One mask is always contained inside the other.\n # Save memory by preventing unnecessary cast to torch.int64\n intersections = (\n (masks > (mask_threshold + threshold_offset))\n .sum(-1, dtype=torch.int16)\n .sum(-1, dtype=torch.int32)\n )\n unions = (\n (masks > (mask_threshold - threshold_offset))\n .sum(-1, dtype=torch.int16)\n .sum(-1, dtype=torch.int32)\n )\n return intersections / unions\n\n\ndef build_point_grid(n_per_side: int) -> np.ndarray:\n \"\"\"Generates a 2D grid of points evenly spaced in [0,1]x[0,1].\"\"\"\n offset = 1 / (2 * n_per_side)\n points_one_side = np.linspace(offset, 1 - offset, n_per_side)\n points_x = np.tile(points_one_side[None, :], (n_per_side, 1))\n points_y = np.tile(points_one_side[:, None], (1, n_per_side))\n points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2)\n return points\n\n\ndef build_all_layer_point_grids(\n n_per_side: int, n_layers: int, scale_per_layer: int\n) -> List[np.ndarray]:\n \"\"\"Generates point grids for all crop layers.\"\"\"\n points_by_layer = []\n for i in range(n_layers + 1):\n n_points = int(n_per_side / (scale_per_layer**i))\n points_by_layer.append(build_point_grid(n_points))\n return points_by_layer\n\n\ndef generate_crop_boxes(\n im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float\n) -> Tuple[List[List[int]], List[int]]:\n \"\"\"\n Generates a list of crop boxes of different sizes. Each layer\n has (2**i)**2 boxes for the ith layer.\n \"\"\"\n crop_boxes, layer_idxs = [], []\n im_h, im_w = im_size\n short_side = min(im_h, im_w)\n\n # Original image\n crop_boxes.append([0, 0, im_w, im_h])\n layer_idxs.append(0)\n\n def crop_len(orig_len, n_crops, overlap):\n return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops))\n\n for i_layer in range(n_layers):\n n_crops_per_side = 2 ** (i_layer + 1)\n overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side))\n\n crop_w = crop_len(im_w, n_crops_per_side, overlap)\n crop_h = crop_len(im_h, n_crops_per_side, overlap)\n\n crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)]\n crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)]\n\n # Crops in XYWH format\n for x0, y0 in product(crop_box_x0, crop_box_y0):\n box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)]\n crop_boxes.append(box)\n layer_idxs.append(i_layer + 1)\n\n return crop_boxes, layer_idxs\n\n\ndef uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor:\n x0, y0, _, _ = crop_box\n offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device)\n # Check if boxes has a channel dimension\n if len(boxes.shape) == 3:\n offset = offset.unsqueeze(1)\n return boxes + offset\n\n\ndef uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor:\n x0, y0, _, _ = crop_box\n offset = torch.tensor([[x0, y0]], device=points.device)\n # Check if points has a channel dimension\n if len(points.shape) == 3:\n offset = offset.unsqueeze(1)\n return points + offset\n\n\ndef uncrop_masks(\n masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int\n) -> torch.Tensor:\n x0, y0, x1, y1 = crop_box\n if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h:\n return masks\n # Coordinate transform masks\n pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0)\n pad = (x0, pad_x - x0, y0, pad_y - y0)\n return torch.nn.functional.pad(masks, pad, value=0)\n\n\ndef remove_small_regions(\n mask: np.ndarray, area_thresh: float, mode: str\n) -> Tuple[np.ndarray, bool]:\n \"\"\"\n Removes small disconnected regions and holes in a mask. Returns the\n mask and an indicator of if the mask has been modified.\n \"\"\"\n import cv2 # type: ignore\n\n assert mode in [\"holes\", \"islands\"]\n correct_holes = mode == \"holes\"\n working_mask = (correct_holes ^ mask).astype(np.uint8)\n n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8)\n sizes = stats[:, -1][1:] # Row 0 is background label\n small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh]\n if len(small_regions) == 0:\n return mask, False\n fill_labels = [0] + small_regions\n if not correct_holes:\n fill_labels = [i for i in range(n_labels) if i not in fill_labels]\n # If every region is below threshold, keep largest\n if len(fill_labels) == 0:\n fill_labels = [int(np.argmax(sizes)) + 1]\n mask = np.isin(regions, fill_labels)\n return mask, True\n\n\ndef coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]:\n from pycocotools import mask as mask_utils # type: ignore\n\n h, w = uncompressed_rle[\"size\"]\n rle = mask_utils.frPyObjects(uncompressed_rle, h, w)\n rle[\"counts\"] = rle[\"counts\"].decode(\"utf-8\") # Necessary to serialize with json\n return rle\n\n\ndef batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Calculates boxes in XYXY format around masks. Return [0,0,0,0] for\n an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.\n \"\"\"\n # torch.max below raises an error on empty inputs, just skip in this case\n if torch.numel(masks) == 0:\n return torch.zeros(*masks.shape[:-2], 4, device=masks.device)\n\n # Normalize shape to CxHxW\n shape = masks.shape\n h, w = shape[-2:]\n if len(shape) > 2:\n masks = masks.flatten(0, -3)\n else:\n masks = masks.unsqueeze(0)\n\n # Get top and bottom edges\n in_height, _ = torch.max(masks, dim=-1)\n in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :]\n bottom_edges, _ = torch.max(in_height_coords, dim=-1)\n in_height_coords = in_height_coords + h * (~in_height)\n top_edges, _ = torch.min(in_height_coords, dim=-1)\n\n # Get left and right edges\n in_width, _ = torch.max(masks, dim=-2)\n in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :]\n right_edges, _ = torch.max(in_width_coords, dim=-1)\n in_width_coords = in_width_coords + w * (~in_width)\n left_edges, _ = torch.min(in_width_coords, dim=-1)\n\n # If the mask is empty the right edge will be to the left of the left edge.\n # Replace these boxes with [0, 0, 0, 0]\n empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)\n out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)\n out = out * (~empty_filter).unsqueeze(-1)\n\n # Return to original shape\n if len(shape) > 2:\n out = out.reshape(*shape[:-2], 4)\n else:\n out = out[0]\n\n return out\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/utils/misc.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport os\nimport warnings\nfrom threading import Thread\n\nimport numpy as np\nimport torch\nfrom PIL import Image\nfrom tqdm import tqdm\n\n\ndef get_sdpa_settings():\n if torch.cuda.is_available():\n old_gpu = torch.cuda.get_device_properties(0).major < 7\n # only use Flash Attention on Ampere (8.0) or newer GPUs\n use_flash_attn = torch.cuda.get_device_properties(0).major >= 8\n if not use_flash_attn:\n warnings.warn(\n \"Flash Attention is disabled as it requires a GPU with Ampere (8.0) CUDA capability.\",\n category=UserWarning,\n stacklevel=2,\n )\n # keep math kernel for PyTorch versions before 2.2 (Flash Attention v2 is only\n # available on PyTorch 2.2+, while Flash Attention v1 cannot handle all cases)\n pytorch_version = tuple(int(v) for v in torch.__version__.split(\".\")[:2])\n if pytorch_version < (2, 2):\n warnings.warn(\n f\"You are using PyTorch {torch.__version__} without Flash Attention v2 support. \"\n \"Consider upgrading to PyTorch 2.2+ for Flash Attention v2 (which could be faster).\",\n category=UserWarning,\n stacklevel=2,\n )\n math_kernel_on = pytorch_version < (2, 2) or not use_flash_attn\n else:\n old_gpu = True\n use_flash_attn = False\n math_kernel_on = True\n\n return old_gpu, use_flash_attn, math_kernel_on\n\n\ndef get_connected_components(mask):\n \"\"\"\n Get the connected components (8-connectivity) of binary masks of shape (N, 1, H, W).\n\n Inputs:\n - mask: A binary mask tensor of shape (N, 1, H, W), where 1 is foreground and 0 is\n background.\n\n Outputs:\n - labels: A tensor of shape (N, 1, H, W) containing the connected component labels\n for foreground pixels and 0 for background pixels.\n - counts: A tensor of shape (N, 1, H, W) containing the area of the connected\n components for foreground pixels and 0 for background pixels.\n \"\"\"\n from sam2 import _C\n\n return _C.get_connected_componnets(mask.to(torch.uint8).contiguous())\n\n\ndef mask_to_box(masks: torch.Tensor):\n \"\"\"\n compute bounding box given an input mask\n\n Inputs:\n - masks: [B, 1, H, W] masks, dtype=torch.Tensor\n\n Returns:\n - box_coords: [B, 1, 4], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.Tensor\n \"\"\"\n B, _, h, w = masks.shape\n device = masks.device\n xs = torch.arange(w, device=device, dtype=torch.int32)\n ys = torch.arange(h, device=device, dtype=torch.int32)\n grid_xs, grid_ys = torch.meshgrid(xs, ys, indexing=\"xy\")\n grid_xs = grid_xs[None, None, ...].expand(B, 1, h, w)\n grid_ys = grid_ys[None, None, ...].expand(B, 1, h, w)\n min_xs, _ = torch.min(torch.where(masks, grid_xs, w).flatten(-2), dim=-1)\n max_xs, _ = torch.max(torch.where(masks, grid_xs, -1).flatten(-2), dim=-1)\n min_ys, _ = torch.min(torch.where(masks, grid_ys, h).flatten(-2), dim=-1)\n max_ys, _ = torch.max(torch.where(masks, grid_ys, -1).flatten(-2), dim=-1)\n bbox_coords = torch.stack((min_xs, min_ys, max_xs, max_ys), dim=-1)\n\n return bbox_coords\n\n\ndef _load_img_as_tensor(img_path, image_size):\n img_pil = Image.open(img_path)\n img_np = np.array(img_pil.convert(\"RGB\").resize((image_size, image_size)))\n if img_np.dtype == np.uint8: # np.uint8 is expected for JPEG images\n img_np = img_np / 255.0\n else:\n raise RuntimeError(f\"Unknown image dtype: {img_np.dtype} on {img_path}\")\n img = torch.from_numpy(img_np).permute(2, 0, 1)\n video_width, video_height = img_pil.size # the original video size\n return img, video_height, video_width\n\n\nclass AsyncVideoFrameLoader:\n \"\"\"\n A list of video frames to be load asynchronously without blocking session start.\n \"\"\"\n\n def __init__(\n self,\n img_paths,\n image_size,\n offload_video_to_cpu,\n img_mean,\n img_std,\n compute_device,\n ):\n self.img_paths = img_paths\n self.image_size = image_size\n self.offload_video_to_cpu = offload_video_to_cpu\n self.img_mean = img_mean\n self.img_std = img_std\n # items in `self.images` will be loaded asynchronously\n self.images = [None] * len(img_paths)\n # catch and raise any exceptions in the async loading thread\n self.exception = None\n # video_height and video_width be filled when loading the first image\n self.video_height = None\n self.video_width = None\n self.compute_device = compute_device\n\n # load the first frame to fill video_height and video_width and also\n # to cache it (since it's most likely where the user will click)\n self.__getitem__(0)\n\n # load the rest of frames asynchronously without blocking the session start\n def _load_frames():\n try:\n for n in tqdm(range(len(self.images)), desc=\"frame loading (JPEG)\"):\n self.__getitem__(n)\n except Exception as e:\n self.exception = e\n\n self.thread = Thread(target=_load_frames, daemon=True)\n self.thread.start()\n\n def __getitem__(self, index):\n if self.exception is not None:\n raise RuntimeError(\"Failure in frame loading thread\") from self.exception\n\n img = self.images[index]\n if img is not None:\n return img\n\n img, video_height, video_width = _load_img_as_tensor(\n self.img_paths[index], self.image_size\n )\n self.video_height = video_height\n self.video_width = video_width\n # normalize by mean and std\n img -= self.img_mean\n img /= self.img_std\n if not self.offload_video_to_cpu:\n img = img.to(self.compute_device, non_blocking=True)\n self.images[index] = img\n return img\n\n def __len__(self):\n return len(self.images)\n\n\ndef load_video_frames(\n video_path,\n image_size,\n offload_video_to_cpu,\n img_mean=(0.485, 0.456, 0.406),\n img_std=(0.229, 0.224, 0.225),\n async_loading_frames=False,\n compute_device=torch.device(\"cuda\"),\n):\n \"\"\"\n Load the video frames from video_path. The frames are resized to image_size as in\n the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo.\n \"\"\"\n is_bytes = isinstance(video_path, bytes)\n is_str = isinstance(video_path, str)\n is_mp4_path = is_str and os.path.splitext(video_path)[-1] in [\".mp4\", \".MP4\"]\n if is_bytes or is_mp4_path:\n return load_video_frames_from_video_file(\n video_path=video_path,\n image_size=image_size,\n offload_video_to_cpu=offload_video_to_cpu,\n img_mean=img_mean,\n img_std=img_std,\n compute_device=compute_device,\n )\n elif is_str and os.path.isdir(video_path):\n return load_video_frames_from_jpg_images(\n video_path=video_path,\n image_size=image_size,\n offload_video_to_cpu=offload_video_to_cpu,\n img_mean=img_mean,\n img_std=img_std,\n async_loading_frames=async_loading_frames,\n compute_device=compute_device,\n )\n else:\n raise NotImplementedError(\n \"Only MP4 video and JPEG folder are supported at this moment\"\n )\n\n\ndef load_video_frames_from_jpg_images(\n video_path,\n image_size,\n offload_video_to_cpu,\n img_mean=(0.485, 0.456, 0.406),\n img_std=(0.229, 0.224, 0.225),\n async_loading_frames=False,\n compute_device=torch.device(\"cuda\"),\n):\n \"\"\"\n Load the video frames from a directory of JPEG files (\".jpg\" format).\n\n The frames are resized to image_size x image_size and are loaded to GPU if\n `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`.\n\n You can load a frame asynchronously by setting `async_loading_frames` to `True`.\n \"\"\"\n if isinstance(video_path, str) and os.path.isdir(video_path):\n jpg_folder = video_path\n else:\n raise NotImplementedError(\n \"Only JPEG frames are supported at this moment. For video files, you may use \"\n \"ffmpeg (https://ffmpeg.org/) to extract frames into a folder of JPEG files, such as \\n\"\n \"```\\n\"\n \"ffmpeg -i .mp4 -q:v 2 -start_number 0 /'%05d.jpg'\\n\"\n \"```\\n\"\n \"where `-q:v` generates high-quality JPEG frames and `-start_number 0` asks \"\n \"ffmpeg to start the JPEG file from 00000.jpg.\"\n )\n\n frame_names = [\n p\n for p in os.listdir(jpg_folder)\n if os.path.splitext(p)[-1] in [\".jpg\", \".jpeg\", \".JPG\", \".JPEG\"]\n ]\n frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))\n num_frames = len(frame_names)\n if num_frames == 0:\n raise RuntimeError(f\"no images found in {jpg_folder}\")\n img_paths = [os.path.join(jpg_folder, frame_name) for frame_name in frame_names]\n img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]\n img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]\n\n if async_loading_frames:\n lazy_images = AsyncVideoFrameLoader(\n img_paths,\n image_size,\n offload_video_to_cpu,\n img_mean,\n img_std,\n compute_device,\n )\n return lazy_images, lazy_images.video_height, lazy_images.video_width\n\n images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)\n for n, img_path in enumerate(tqdm(img_paths, desc=\"frame loading (JPEG)\")):\n images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size)\n if not offload_video_to_cpu:\n images = images.to(compute_device)\n img_mean = img_mean.to(compute_device)\n img_std = img_std.to(compute_device)\n # normalize by mean and std\n images -= img_mean\n images /= img_std\n return images, video_height, video_width\n\n\ndef load_video_frames_from_video_file(\n video_path,\n image_size,\n offload_video_to_cpu,\n img_mean=(0.485, 0.456, 0.406),\n img_std=(0.229, 0.224, 0.225),\n compute_device=torch.device(\"cuda\"),\n):\n \"\"\"Load the video frames from a video file.\"\"\"\n import decord\n\n img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]\n img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]\n # Get the original video height and width\n decord.bridge.set_bridge(\"torch\")\n video_height, video_width, _ = decord.VideoReader(video_path).next().shape\n # Iterate over all frames in the video\n images = []\n for frame in decord.VideoReader(video_path, width=image_size, height=image_size):\n images.append(frame.permute(2, 0, 1))\n\n images = torch.stack(images, dim=0).float() / 255.0\n if not offload_video_to_cpu:\n images = images.to(compute_device)\n img_mean = img_mean.to(compute_device)\n img_std = img_std.to(compute_device)\n # normalize by mean and std\n images -= img_mean\n images /= img_std\n return images, video_height, video_width\n\n\ndef fill_holes_in_mask_scores(mask, max_area):\n \"\"\"\n A post processor to fill small holes in mask scores with area under `max_area`.\n \"\"\"\n # Holes are those connected components in background with area <= self.max_area\n # (background regions are those with mask scores <= 0)\n assert max_area > 0, \"max_area must be positive\"\n\n input_mask = mask\n try:\n labels, areas = get_connected_components(mask <= 0)\n is_hole = (labels > 0) & (areas <= max_area)\n # We fill holes with a small positive mask score (0.1) to change them to foreground.\n mask = torch.where(is_hole, 0.1, mask)\n except Exception as e:\n # Skip the post-processing step on removing small holes if the CUDA kernel fails\n warnings.warn(\n f\"{e}\\n\\nSkipping the post-processing step due to the error above. You can \"\n \"still use SAM 2 and it's OK to ignore the error above, although some post-processing \"\n \"functionality may be limited (which doesn't affect the results in most cases; see \"\n \"https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).\",\n category=UserWarning,\n stacklevel=2,\n )\n mask = input_mask\n\n return mask\n\n\ndef concat_points(old_point_inputs, new_points, new_labels):\n \"\"\"Add new points and labels to previous point inputs (add at the end).\"\"\"\n if old_point_inputs is None:\n points, labels = new_points, new_labels\n else:\n points = torch.cat([old_point_inputs[\"point_coords\"], new_points], dim=1)\n labels = torch.cat([old_point_inputs[\"point_labels\"], new_labels], dim=1)\n\n return {\"point_coords\": points, \"point_labels\": labels}\n" }, { "path": "camera_pose_annotation/dynamic_mask/sam2/utils/transforms.py", "content": "# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n\n# This source code is licensed under the license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport warnings\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torchvision.transforms import Normalize, Resize, ToTensor\n\n\nclass SAM2Transforms(nn.Module):\n def __init__(\n self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0\n ):\n \"\"\"\n Transforms for SAM2.\n \"\"\"\n super().__init__()\n self.resolution = resolution\n self.mask_threshold = mask_threshold\n self.max_hole_area = max_hole_area\n self.max_sprinkle_area = max_sprinkle_area\n self.mean = [0.485, 0.456, 0.406]\n self.std = [0.229, 0.224, 0.225]\n self.to_tensor = ToTensor()\n self.transforms = torch.jit.script(\n nn.Sequential(\n Resize((self.resolution, self.resolution)),\n Normalize(self.mean, self.std),\n )\n )\n\n def __call__(self, x):\n x = self.to_tensor(x)\n return self.transforms(x)\n\n def forward_batch(self, img_list):\n img_batch = [self.transforms(self.to_tensor(img)) for img in img_list]\n img_batch = torch.stack(img_batch, dim=0)\n return img_batch\n\n def transform_coords(\n self, coords: torch.Tensor, normalize=False, orig_hw=None\n ) -> torch.Tensor:\n \"\"\"\n Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates,\n If the coords are in absolute image coordinates, normalize should be set to True and original image size is required.\n\n Returns\n Un-normalized coordinates in the range of [0, 1] which is expected by the SAM2 model.\n \"\"\"\n if normalize:\n assert orig_hw is not None\n h, w = orig_hw\n coords = coords.clone()\n coords[..., 0] = coords[..., 0] / w\n coords[..., 1] = coords[..., 1] / h\n\n coords = coords * self.resolution # unnormalize coords\n return coords\n\n def transform_boxes(\n self, boxes: torch.Tensor, normalize=False, orig_hw=None\n ) -> torch.Tensor:\n \"\"\"\n Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates,\n if the coords are in absolute image coordinates, normalize should be set to True and original image size is required.\n \"\"\"\n boxes = self.transform_coords(boxes.reshape(-1, 2, 2), normalize, orig_hw)\n return boxes\n\n def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor:\n \"\"\"\n Perform PostProcessing on output masks.\n \"\"\"\n from sam2.utils.misc import get_connected_components\n\n masks = masks.float()\n input_masks = masks\n mask_flat = masks.flatten(0, 1).unsqueeze(1) # flatten as 1-channel image\n try:\n if self.max_hole_area > 0:\n # Holes are those connected components in background with area <= self.fill_hole_area\n # (background regions are those with mask scores <= self.mask_threshold)\n labels, areas = get_connected_components(\n mask_flat <= self.mask_threshold\n )\n is_hole = (labels > 0) & (areas <= self.max_hole_area)\n is_hole = is_hole.reshape_as(masks)\n # We fill holes with a small positive mask score (10.0) to change them to foreground.\n masks = torch.where(is_hole, self.mask_threshold + 10.0, masks)\n\n if self.max_sprinkle_area > 0:\n labels, areas = get_connected_components(\n mask_flat > self.mask_threshold\n )\n is_hole = (labels > 0) & (areas <= self.max_sprinkle_area)\n is_hole = is_hole.reshape_as(masks)\n # We fill holes with negative mask score (-10.0) to change them to background.\n masks = torch.where(is_hole, self.mask_threshold - 10.0, masks)\n except Exception as e:\n # Skip the post-processing step if the CUDA kernel fails\n warnings.warn(\n f\"{e}\\n\\nSkipping the post-processing step due to the error above. You can \"\n \"still use SAM 2 and it's OK to ignore the error above, although some post-processing \"\n \"functionality may be limited (which doesn't affect the results in most cases; see \"\n \"https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).\",\n category=UserWarning,\n stacklevel=2,\n )\n masks = input_masks\n\n masks = F.interpolate(masks, orig_hw, mode=\"bilinear\", align_corners=False)\n return masks\n" }, { "path": "caption/LLM/__init__.py", "content": "" }, { "path": "caption/LLM/inference.py", "content": "import os\nimport time\nimport queue\nfrom argparse import ArgumentParser\nfrom multiprocessing import Manager\nfrom concurrent.futures import ThreadPoolExecutor\nfrom tqdm import tqdm\nimport concurrent\nimport pandas as pd\nimport numpy as np\nimport sys\nsys.path.append(os.path.abspath(os.path.join(__file__, \"../..\")))\nfrom utils.api_call import api_call\n\n\ndef get_pose(pose_dir):\n \"\"\"\n Retrieve and process pose data from extrinsics.npy file\n \"\"\"\n # Base directory for pose data\n pose_path = os.path.join(pose_dir, 'extrinsics.npy')\n assert os.path.isfile(pose_path), f\"Pose file not found: {pose_path}\"\n\n # Load and process the pose file\n poses = np.load(pose_path)\n\n # Data processing steps\n poses = poses[::5, :, 3] # Take first row for every 5 rows\n max_value = np.max(poses)\n min_value = np.min(poses)\n min_abs_value = np.min(np.abs(poses))\n\n # Normalize and convert to integers (minimize integer digits)\n poses = np.round(poses / (max_value - min_value) /\n min_abs_value).astype(int)\n\n # Keep only first 3 columns and transpose\n poses = poses[:, :3].T\n\n # Extract individual axes\n poses1, poses2, poses3 = poses[0], poses[1], poses[2]\n\n # Convert each axis to string\n poses1_str = ' '.join(map(str, poses1))\n poses2_str = ' '.join(map(str, poses2))\n poses3_str = ' '.join(map(str, poses3))\n\n # Combine into formatted string\n poses_str = f'x:{poses1_str}\\ny:{poses2_str}\\nz:{poses3_str}'\n\n return poses_str\n\n\ndef get_prompt(pose_dir, prompt_dir, vqa_caption, dist_level):\n \"\"\"\n Construct a prompt by combining content from prompt1.txt, prompt2.txt, VQA caption, and pose data\n \"\"\"\n # Read prompt components\n p1_file = os.path.join(prompt_dir, 'prompt1.txt')\n p2_file = os.path.join(prompt_dir, 'prompt2.txt')\n\n with open(p1_file, 'r', encoding='utf-8') as f:\n p1_content = f.read().strip()\n\n with open(p2_file, 'r', encoding='utf-8') as f:\n p2_content = f.read().strip()\n\n # Get pose data\n poses = get_pose(pose_dir)\n\n # Assemble final prompt\n prompt = (f\"{p1_content}\\nGiven Information:\\n{vqa_caption}\\n3.Camera Position Data:\\n{poses}\\n\"\n f\"\\n4.Motion intensity:\\n{dist_level}\\n{p2_content}\")\n\n return prompt\n\n\ndef process_single_row(args, row):\n \"\"\"\n Process a single row of data by calling API and saving the result\n \"\"\"\n # Check if VQA file exists\n vqa_path = os.path.join(args.vqa_path, f\"{row['id']}.txt\")\n assert os.path.isfile(vqa_path), f\"VQA file not found: {vqa_path}\"\n # Read VQA caption\n with open(vqa_path, \"r\") as f:\n vqa_caption = f.read()\n\n # Skip processing if file already exists\n save_file = os.path.join(args.llm_path, f\"{row['id']}.txt\")\n if os.path.exists(save_file) and os.path.getsize(save_file) > 0:\n return\n\n # Call API with retry mechanism\n pose_dir = os.path.join(args.pose_load_dir, row[\"id\"], \"reconstructions\")\n prompt_text = get_prompt(pose_dir, args.prompt_dir,\n vqa_caption, row[\"distLevel\"])\n llm_caption = api_call(prompt_text, args.model,\n args.api_key, args.base_domain)\n assert llm_caption is not None, f\"API call failed for id {row['id']}\"\n\n # Save the result with model information\n with open(save_file, 'w', encoding='utf-8') as f:\n f.write(llm_caption + f\"\\n\\n6. Qwen model: \\n{args.model}\")\n return\n\n\ndef worker(args, task_queue, pbar):\n \"\"\"\n Worker function to process tasks from the queue\n\n Args:\n task_queue: Queue containing tasks to process\n pbar: Progress bar object for tracking progress\n \"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n\n # Add delay to prevent overwhelming API\n time.sleep(args.wait_time)\n\n # Process the single row\n process_single_row(args, row)\n\n # Update progress\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"\n Parse command line arguments\n\n Returns:\n Parsed arguments object\n \"\"\"\n parser = ArgumentParser(description='VQA Processing Program')\n parser.add_argument('--csv_path', type=str, required=True,\n help='Path to CSV file')\n parser.add_argument('--pose_load_dir', type=str, required=True,\n help='Directory to load pose data')\n parser.add_argument('--output_dir', type=str, required=True,\n help='Directory to save results')\n parser.add_argument('--prompt_dir', type=str,\n default=os.path.join(os.path.dirname(\n __file__), \"vqa_prompt.txt\"),\n help='Path to prompt file')\n parser.add_argument('--model', type=str, default=\"qwen3-30b-a3b\",\n help='Model name')\n parser.add_argument('--api_key', type=str,\n default=\"sk-****\",\n help='API key')\n parser.add_argument('--num_workers', type=int, default=1,\n help='Number of worker threads')\n parser.add_argument('--wait_time', type=float, default=0.5,\n help='Time between requests in seconds')\n parser.add_argument('--base_domain', type=str, default=\"https://cn2us02.opapi.win/\",\n help='API base domain')\n return parser.parse_args()\n\n\ndef main():\n \"\"\"\n Main processing function that handles multiple rows using parallel workers\n\n Args:\n group_id (str): Identifier for the group\n prompt_dir (str): Directory containing prompt files\n model_file (str): Path to file containing model names\n api_key_file (str): Path to file containing API keys\n num_workers (int): Number of worker threads\n wait_time (float): Time to wait between requests\n base_domain (str): Base domain for API calls\n record_time (bool): Whether to record processing time\n\n Returns:\n None\n \"\"\"\n args = parse_args()\n\n # Validate temporary directory exists\n # Create LLM directory if it doesn't exist\n args.llm_path = os.path.join(args.output_dir, \"LLM\")\n if not os.path.isdir(args.llm_path):\n os.makedirs(args.llm_path, exist_ok=True)\n\n # Validate VQA directory exists\n args.vqa_path = os.path.join(args.output_dir, \"VQA\")\n assert os.path.isdir(\n args.vqa_path), f\"VQA directory not found: {args.vqa_path}\"\n\n # Read CSV file containing scene information\n df = pd.read_csv(args.csv_path)\n\n # Initialize task queue with all rows\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n # Start processing with progress bar\n with tqdm(total=len(df), desc=\"LLM Finished\") as pbar:\n with ThreadPoolExecutor(max_workers=args.num_workers) as executor:\n # Start worker threads\n futures = [executor.submit(worker, args, task_queue, pbar)\n for _ in range(args.num_workers)]\n # Wait for all workers to complete\n for future in concurrent.futures.as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "caption/LLM/prompt1.txt", "content": "You are given a video sequence with camera trajectory data representing the camera's movement through a scene.\n\nThe data consists of:\nCamera Motion Caption: A basic description of how the camera moves.\nScene Description: A detailed visual summary of the environment.\nCamera position data: Three lines, representing the sequence of the camera's x-coordinate, y-coordinate, and z-coordinate. These values are derived from normalized 3D pose data using the following formula: poses = np.round(poses / (max_value - min_value) / min_abs_value).astype(int); Each value is then multiplied by 1,000,000 and rounded to the nearest integer.\nMotion intensity: An integer that indicates the level of camera movement, where a value of 0 means the camera is static, 1 indicates slight movement, and 2 or higher represents normal or noticeable motion. In tasks such as Optimized Camera Motion Caption and Main Motion Trend Summary, this intensity value should be used to qualify the degree of motion described — for example, using \"slight forward translate\" when the intensity is 1.\n\nYour Tasks:\n\n1. Optimized Camera Motion Caption\nGenerate a refined motion caption **from the perspective of the camera itself**, using only the **camera position data** to determine movement direction and dynamics.\n\nUse the following rules to interpret motion:\n x increasing: camera moves right\n x decreasing: camera moves left\n y increasing: camera moves down\n y decreasing: camera moves up\n z increasing: camera moves forward\n z decreasing: camera moves backward\n\nAnalyze the full trajectory over time to capture acceleration, deceleration, or steady motion. Integrate scene context but prioritize accuracy based on numerical data. Avoid vague phrases like \"zoom out\" unless it's clearly due to focal length change — here, use translation terms instead.\nIf motion intensity is 0, describe the fixed viewpoint and what the camera observes from that vantage point, incorporating compositional or environmental elements from the original caption. If intensity is 1, reflect subtle movement in the description (e.g., \"slight right translate\") without exaggerating the motion. For both cases, preserve visual context while aligning with the actual movement level. Avoid mentioning data analysis or detection explicitly — let the description itself reflect the motion state.\n\nTarget Length: 50–100 words\n\n2. Scene Abstract Caption \nProvide a single-sentence summary that captures: \n- Key architectural elements \n- Overall atmosphere/style \n- Notable design features \nTarget Length: About 50 words\n\n3. Main Motion Trend Summary \nSummarize the general movement using only 1–3 short motion phrases , depending on how many are clearly present. Focus strictly on major, sustained movements — ignore minor fluctuations or brief directional changes. If only one or two movements dominate, list only those. Use directional translation terms (e.g., forward translate, left translate, upward drift)\n\n4. Scene Keywords \nExtract up to 4 keywords summarizing the key aspects of the scene. Include one term that broadly describes the scene type. Use nouns/noun phrases related to weather, place, time, lighting, scene type. Avoid adjectives/gerunds except for weather. Example: sunset, foggy, marketplace, city street, village\n\n5. Immersive Shot Summary\nBlend Optimized Camera Motion Caption and Scene Description evenly — do not focus more on the camera or the scene alone. Describe the visuals as if someone is watching a moving image unfold. Use descriptive, cinematic language that evokes imagery and emotion. Keep it concise but expressive — suitable for use in scripts, storyboards, or AI video/image generation. Target Length: 50–100 words" }, { "path": "caption/LLM/prompt2.txt", "content": "Output Format:\n\n1. Camera Motion Caption: \n[From the perspective of the camera holder, with the camera as the subject. Combine camera pose information to describe]\n\n2. Scene Abstract Caption: \n[A concise one sentence summary of the scene]\n\n3. Main Motion Trend Summary: \n[keywords separated by commas, e.g., forward translate, downward tilt]\n\n4. Scene Keywords:\n[word1, word2, word3, ...] (max 5 words)\n\n5. Immersive Shot Summary:" }, { "path": "caption/README.md", "content": "# Semantic Information Annotation\n\nThis script automates the process of generating structured text descriptions (captions) for videos through a multi-step pipeline involving Visual Question Answering (VQA), Large Language Models (LLM), result combination, and tagging.\n\n## Captioning Workflow\n\nThe video captioning process follows these sequential steps:\n\n1. **VQA Captioning**: Uses a Visual Question Answering model to analyze visual content and generate initial captions based on predefined prompts.\n2. **LLM Captioning**: Employs a Large Language Model to process pose data and generate additional descriptive captions.\n3. **Result Combination**: Merges the outputs from the VQA and LLM steps into a unified structure.\n4. **Tag Addition**: Enhances the combined results with relevant tags using a language model.\n
\n \"caption\n
\n\n## Script Explanation\n\n### Configuration Parameters\n\n- `CSV`: Path to the result CSV file generated in the annotation step\n- `SRC_DIR`: Path to the annotation output directory containing video frames and pose data\n- `OUTPUT_DIR`: Path where all output files will be saved\n- `num_workers`: Number of parallel workers to use for processing\n- `wait_time`: Waiting time between API requests (in seconds)\n\n### Step 1: VQA Captioning\n\nGenerates captions by analyzing visual content using a VQA model.\n\nParameters:\n- `--csv_path`: Path to the input CSV file\n- `--fig_load_dir`: Directory containing video frames/images\n- `--output_dir`: Directory to save VQA results\n- `--prompt_file`: Path to VQA prompt template file\n- `--model`: VQA model to use (default: gemini-2.0-flash)\n- `--api_key`: API key for accessing the VQA model service\n- `--base_domain`: API endpoint domain for the VQA model\n- `--num_workers`: Number of parallel workers\n- `--wait_time`: Waiting time between API requests\n\n### Step 2: LLM Captioning\n\nGenerates additional captions by processing pose data using a Large Language Model.\n\nParameters:\n- `--csv_path`: Path to the input CSV file\n- `--pose_load_dir`: Directory containing pose data\n- `--output_dir`: Directory to save LLM results\n- `--prompt_dir`: Directory containing LLM prompt templates\n- `--model`: LLM model to use (default: qwen3-30b-a3b)\n- `--api_key`: API key for accessing the LLM service\n- `--base_domain`: API endpoint domain for the LLM\n- `--num_workers`: Number of parallel workers\n- `--wait_time`: Waiting time between API requests\n\n### Step 3: Combine Results\n\nMerges the outputs from VQA and LLM steps into a unified format.\n\nParameters:\n- `--csv_path`: Path to the input CSV file\n- `--load_dir`: Directory containing VQA and LLM results\n- `--output_dir`: Directory to save combined results\n- `--num_workers`: Number of parallel workers\n\n### Step 4: Add Tags\n\nEnhances the combined results with relevant tags using a language model.\n\nParameters:\n- `--csv_path`: Path to the input CSV file\n- `--json_load_dir`: Directory containing combined results\n- `--prompt_file`: Path to tagging prompt template file\n- `--model`: Model to use for tagging (default: qwen3-30b-a3b)\n- `--api_key`: API key for accessing the tagging model service\n- `--base_domain`: API endpoint domain for the tagging model\n- `--num_workers`: Number of parallel workers\n- `--wait_time`: Waiting time between API requests\n\n## Usage\n\n1. Replace all placeholder values (enclosed in square brackets) with your actual paths and API keys\n2. Make the script executable: `chmod +x caption_pipeline.sh`\n3. Run the script: `./caption_pipeline.sh`\n\nThe script will execute each step sequentially, displaying start/end times and duration for each step, and save all outputs to the specified `OUTPUT_DIR`.\n\n## results example\nseveral samples of video captions generated by the model after each step." }, { "path": "caption/VQA/__init__.py", "content": "" }, { "path": "caption/VQA/inference.py", "content": "import os\nimport concurrent.futures\nfrom multiprocessing import Manager\nimport queue\nimport pandas as pd\nfrom tqdm import tqdm\nimport argparse\nimport time\nimport base64\nimport cv2\nfrom glob import glob\nimport sys\nsys.path.append(os.path.abspath(os.path.join(__file__, \"../..\")))\nfrom utils.api_call import api_call\n\n\ndef encode_image(image_path):\n \"\"\"\n Resizes an image to 640x360 and encodes it as a Base64 string with data URI prefix.\n \"\"\"\n # Read image using OpenCV\n image = cv2.imread(image_path)\n\n # Resize image to standard dimensions (640x360)\n resized_image = cv2.resize(image, (640, 360))\n\n # Encode image as JPEG and convert to Base64\n _, buffer = cv2.imencode('.jpeg', resized_image)\n base64_data = base64.b64encode(buffer).decode(\"utf-8\")\n\n # Return with data URI format for API compatibility\n return f\"data:image/jpeg;base64,{base64_data}\"\n\n\ndef get_prompt(fig_dir, prompt_text):\n \"\"\"\n Load key frames from a video, constructs a multimodal request, and calls the API.\n \"\"\"\n # Get frames from directory\n frames = sorted(glob(f\"{fig_dir}/*.jpg\"))[::5]\n\n # Construct multimodal input content\n messages_content = []\n\n # Add encoded images to request content\n for frame in frames:\n try:\n encoded_frame = encode_image(frame)\n messages_content.append({\n \"type\": \"image_url\",\n \"image_url\": {\"url\": encoded_frame}\n })\n except Exception as e:\n print(f\"Image processing error: {str(e)}\")\n return None\n\n # Add text prompt to request content\n messages_content.append({\"type\": \"text\", \"text\": prompt_text})\n\n return messages_content\n\n\ndef process_single_row(args, row):\n \"\"\"\n Process a single row: call the VQA API and save the result for one scene.\n Handles retries and error logging.\n \"\"\"\n save_path = os.path.join(args.output_dir, \"VQA\")\n if not os.path.isdir(save_path):\n os.makedirs(save_path, exist_ok=True)\n save_file = os.path.join(save_path, f\"{row['id']}.txt\")\n if os.path.exists(save_file) and os.path.getsize(save_file) > 0:\n # Skip if already exists\n return\n # Call API\n fig_dir = os.path.join(args.fig_load_dir, row['id'], \"img\")\n prompt_text = get_prompt(fig_dir, args.prompt_text)\n vqa_caption = api_call(prompt_text, args.model,\n args.api_key, args.base_domain)\n assert vqa_caption is not None, f\"API call failed for id {row['id']}\"\n # Save result\n with open(save_file, 'w', encoding='utf-8') as f:\n f.write(vqa_caption)\n\n\ndef worker(args, task_queue, pbar):\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n time.sleep(args.wait_time)\n process_single_row(args, row)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"\n Parse command line arguments for VQA batch processing.\n \"\"\"\n parser = argparse.ArgumentParser(description='VQA batch processing script')\n parser.add_argument('--csv_path', type=str, required=True,\n help='CSV file path')\n parser.add_argument('--fig_load_dir', type=str, required=True,\n help='Directory to load figures')\n parser.add_argument('--output_dir', type=str, required=True,\n help='Directory to save results')\n parser.add_argument('--prompt_file', type=str,\n default=\"vqa_prompt.txt\",\n help='Prompt file path')\n parser.add_argument('--model', type=str, default=\"gemini-2.0-flash\",\n help='Model name')\n parser.add_argument('--api_key', type=str,\n default=\"sk-****\",\n help='API key')\n parser.add_argument('--num_workers', type=int, default=4,\n help='Number of worker threads')\n parser.add_argument('--wait_time', type=float, default=0.8,\n help='Request interval for each thread (seconds)')\n parser.add_argument('--base_domain', type=str, default=\"https://cn2us02.opapi.win/\",\n help='API base domain')\n return parser.parse_args()\n\n\ndef main():\n \"\"\"\n Batch process all scenes in a group: call VQA API for each row in the CSV.\n Uses a thread pool for concurrency and supports timing.\n \"\"\"\n args = parse_args()\n\n df = pd.read_csv(args.csv_path)\n\n # Read prompt text\n with open(args.prompt_file, \"r\", encoding=\"utf-8\") as f:\n args.prompt_text = f.read().strip()\n\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n with tqdm(total=len(df), desc=f\"VQA Finished\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(max_workers=args.num_workers) as executor:\n futures = []\n for worker_id in range(args.num_workers):\n futures.append(executor.submit(worker, args, task_queue, pbar))\n for future in concurrent.futures.as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "caption/VQA/prompt.txt", "content": "You are given a sequence of video frames in chronological order. Analyze them carefully and generate two distinct captions based on the following instructions:\n\n1. Camera Motion Caption:\n\nFrom the perspective of the camera operator, describe the entire motion trajectory of the camera throughout the clip using precise cinematography terminology (e.g., static, pan, tilt, dolly, handheld, crane, aerial, zoom, etc.). \n\nDo NOT assume the camera starts in a \"static\" position just because it appears stationary in the first frame.Only describe the camera as stationary if there is no visual change across multiple consecutive frames.\n\nInstead, focus on changes between frames to infer movement. Describe motion state transitions, not frame-by-frame repetition (e.g., do not say “the camera moves forward again” if it’s continuous). For example:\n- Starting with a dolly forward along a straight path,\n- Then transitioning into a slow right-hand pan,\n- Or shifting from handheld walking movement to a stationary pivot tilt.\n\nInclude brief environmental context where relevant to clarify direction or intent (e.g., \"The camera dollies forward through a narrow alleyway, then smoothly turns left at the intersection\").\n\nKeep the final caption concise, between 50–100 words, focused only on motion and its evolution over time.\n\n2. Scene Description:\n\nProvide a rich, holistic description of the visual content. Include:\n- Main subjects and dynamic objects: who or what is present, and what they are doing (e.g., a cyclist rides past from left to right, a group of people gather near a bench),\n- Background/environment: setting (urban street, forest trail, indoor space), notable landmarks or structures,\n- Lighting and atmosphere: time of day, weather conditions, mood (e.g., golden-hour lighting, overcast sky casting soft shadows, neon-lit nighttime scene),\n- Overall tone or emotion conveyed by the scene.\n\nAvoid focusing on individual frames—describe the general impression and ongoing activity across the entire clip. Aim for around 100 words, balancing detail and conciseness.\n\nOutput Format: \nDo not include any explanations or extra text before or after your response. \nBegin directly with: \n1. Camera Motion Caption: ...\nfollowed by \n2. Scene Description: ..." }, { "path": "caption/__init__.py", "content": "" }, { "path": "caption/tagging/__init__.py", "content": "" }, { "path": "caption/tagging/inference.py", "content": "import os\nimport time\nimport json\nimport queue\nimport argparse\nimport pandas as pd\nfrom tqdm import tqdm\nfrom multiprocessing import Manager\nimport concurrent.futures\nimport sys\nsys.path.append(os.path.abspath(os.path.join(__file__, \"../..\")))\nfrom utils.api_call import api_call\n\n\ndef parse_category_tags(tag_caption):\n \"\"\"\n Parse API response to structured category data using camelCase naming convention\n \"\"\"\n lines = [line.strip()\n for line in tag_caption.strip().split('\\n') if line.strip()]\n\n # Initialize category data with default values\n category_data = {\n \"sceneType\": {\n \"first\": \"Unknown\",\n \"second\": \"Unknown\"\n },\n \"lighting\": \"Unknown\",\n \"timeOfDay\": \"Unknown\",\n \"weather\": \"Unknown\",\n \"crowdDensity\": \"Unknown\"\n }\n\n # Parse each line to extract category information\n for line in lines:\n line_lower = line.lower()\n if line_lower.startswith(\"primary scene type:\"):\n category_data[\"sceneType\"][\"first\"] = line.split(\":\", 1)[1].strip()\n elif line_lower.startswith(\"secondary scene type:\"):\n category_data[\"sceneType\"][\"second\"] = line.split(\":\", 1)[\n 1].strip()\n elif line_lower.startswith(\"lighting:\"):\n category_data[\"lighting\"] = line.split(\":\", 1)[1].strip()\n elif line_lower.startswith(\"time of day:\"):\n category_data[\"timeOfDay\"] = line.split(\":\", 1)[1].strip()\n elif line_lower.startswith(\"weather:\"):\n category_data[\"weather\"] = line.split(\":\", 1)[1].strip()\n elif line_lower.startswith(\"crowd density:\"):\n category_data[\"crowdDensity\"] = line.split(\":\", 1)[1].strip()\n\n return category_data\n\n\ndef process_single_row(args, json_file):\n \"\"\"\n Process a single JSON file to add category tags via API call\n \"\"\"\n # Check if CategoryTag field already exists\n with open(json_file, 'r') as f:\n data = json.load(f)\n # Skip if CategoryTag already exists\n if \"CategoryTag\" in data:\n return\n description = data['SceneDesc']\n\n prompt_text = args.prompt_text + description\n\n # Call API to get category tags with retry mechanism\n tag_caption = api_call(prompt_text, args.model,\n args.api_key, args.base_domain)\n assert tag_caption is not None, f\"API call failed for file {json_file}\"\n\n # Parse and add category tags to the JSON file\n category_tag = parse_category_tags(tag_caption)\n\n # Merge new data with existing data\n data[\"CategoryTag\"] = category_tag\n\n # Overwrite file with updated data\n with open(json_file, 'w', encoding='utf-8') as f:\n json.dump(data, f, ensure_ascii=False, indent=2)\n\n\ndef worker(args, task_queue, pbar):\n while True:\n try:\n index, json_file = task_queue.get(timeout=1)\n except queue.Empty:\n break\n\n time.sleep(args.wait_time)\n process_single_row(args, json_file)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments\"\"\"\n parser = argparse.ArgumentParser(\n description='Category Tag Processing Program')\n parser.add_argument('--csv_path', type=str, required=True,\n help='Path to the CSV file')\n parser.add_argument('--json_load_dir', type=str, required=True,\n help='Directory containing JSON files')\n parser.add_argument('--prompt_file', type=str,\n default=\"prompt.txt\",\n help='Path to prompt file')\n parser.add_argument('--model', type=str, default=\"qwen3-30b-a3b\",\n help='Model name')\n parser.add_argument('--api_key', type=str,\n default=\"sk-****\",\n help='API key')\n parser.add_argument('--num_workers', type=int, default=4,\n help='Number of worker threads')\n parser.add_argument('--wait_time', type=float, default=0.8,\n help='Time interval between requests per thread (seconds)')\n parser.add_argument('--base_domain', type=str, default=\"https://cn2us02.opapi.win/\",\n help='API base domain')\n return parser.parse_args()\n\n\ndef main():\n \"\"\"\n Process a group of JSON files using multiple threads to add category tags\n \"\"\"\n args = parse_args()\n\n df = pd.read_csv(args.csv_path)\n\n with open(args.prompt_file, 'r', encoding='utf-8') as f:\n args.prompt_text = f.read().strip()\n\n # Initialize task queue and add all files to process\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in df.iterrows():\n clip_id = row['id']\n json_file = os.path.join(args.json_load_dir, f\"{clip_id}.json\")\n task_queue.put((index, json_file))\n\n # Start processing with progress bar\n with tqdm(total=task_queue.qsize(), desc=\"Tags Completed\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(max_workers=args.num_workers) as executor:\n # Start worker threads\n futures = [executor.submit(worker, args, task_queue, pbar)\n for _ in range(args.num_workers)]\n\n # Wait for all workers to complete\n for future in concurrent.futures.as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "caption/tagging/prompt.txt", "content": "You are an AI assistant specialized in analyzing scene descriptions and extracting structured metadata. Your task is to read the provided scene description and infer six attributes with hierarchical classification where applicable. \n\nOutput Rules: \n1. Scene Type (Choose one from): \n - [Urban, Natural Landscape, Interior, Rural, Waterfront, Unknown] \n - Add a custom secondary tag (unrestricted) to further define the scene (e.g., Urban → \"Street Scene\", Interior → \"Library\") \n2. Lighting: [Bright / Dim/Dark / Unknown] \n3. Time of Day: [Dawn/Morning / Daytime / Dusk/Evening / Night / Unknown] \n4. Weather: [Sunny / Rainy / Foggy / Cloudy / Snowy / Unknown] \n5. Crowd Density: [Deserted / Sparse / Moderate / Crowded / Unknown] \n\nDeduction Guidelines: \n- Prioritize explicit descriptors over implied cues (e.g., \"snow scattered\" → Snowy; \"wet surfaces + cloudy\" → Cloudy) \n- If no evidence exists for any attribute, output 'Unknown' for that field. Maintain strict objectivity - never assume information beyond the text.\n\nOutput Format (strictly follow line breaks): \nPrimary Scene Type: [X]\nSecondary Scene Type: [CustomTag]\nLighting: [X]\nTime of Day: [X]\nWeather: [X]\nCrowd Density: [X]\n\nThe following is a scene description: " }, { "path": "caption/utils/__init__.py", "content": "" }, { "path": "caption/utils/api_call.py", "content": "import requests\n\n\ndef api_call(prompt_text, model, api_key, base_domain):\n \"\"\"\n Make an API call to a language model with a constructed prompt,\n handling different API formats for different model providers.\n \"\"\"\n # Determine if using Qwen model API (Aliyun)\n is_qwen = \"dashscope.aliyuncs.com\" in base_domain\n \n # Configure API endpoint and payload based on model type\n if is_qwen:\n api_url = base_domain + \"v1/chat/completions\"\n # Payload format specific to Qwen model\n payload = {\n \"model\": model,\n \"messages\": [\n # {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"}, # Optional system message\n {\"role\": \"user\", \"content\": prompt_text}\n ],\n \"enable_thinking\": False,\n \"temperature\": 0.1 # Low temperature for more deterministic output\n }\n else:\n # Payload format for other models\n api_url = base_domain + \"v1beta/openai/\"\n payload = {\n \"model\": model,\n \"messages\": [\n {\"role\": \"user\", \"content\": prompt_text}\n ],\n \"temperature\": 0.1,\n \"user\": \"User\"\n }\n\n # Configure request headers\n if is_qwen:\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {api_key}\",\n \"Accept\": \"application/json\"\n }\n else:\n headers = {\n \"Content-Type\": \"application/json\",\n \"Authorization\": f\"Bearer {api_key}\",\n \"User-Agent\": f\"({base_domain})\",\n \"Accept\": \"application/json\"\n }\n\n try:\n # Execute API request with timeout\n response = requests.post(\n api_url,\n headers=headers,\n json=payload,\n timeout=120 # 2-minute timeout\n )\n response.raise_for_status() # Raise exception for HTTP errors\n response_data = response.json()\n\n # Optional: Uncomment to log token usage\n # if 'usage' in response_data:\n # usage = response_data.get('usage', {})\n # prompt_tokens = usage.get('prompt_tokens', 0)\n # completion_tokens = usage.get('completion_tokens', 0)\n # total_tokens = usage.get('total_tokens', 0)\n # \n # print(f\"Input tokens: {prompt_tokens}\")\n # print(f\"Output tokens: {completion_tokens}\")\n # print(f\"Total tokens: {total_tokens}\")\n # else:\n # print(\"API response does not contain token usage information\")\n\n # Extract and return response content based on API format\n if is_qwen:\n return response_data.get(\"choices\", [{}])[0].get(\"message\", {}).get(\"content\", \"\")\n else:\n return response_data.get(\"choices\", [{}])[0].get(\"message\", {}).get(\"content\", \"\")\n\n except Exception as e:\n print(f\"API request error: {str(e)}\")\n return None\n" }, { "path": "caption/utils/combine.py", "content": "import os\nimport json\nimport re\nimport queue\nimport argparse\nimport pandas as pd\nfrom multiprocessing import Manager\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nfrom tqdm import tqdm\n\n\ndef parse_text_to_json(text):\n \"\"\"\n Parses text in a specific format into a JSON structure.\n \"\"\"\n # Define mapping between text labels and JSON keys\n labels = {\n \"Camera Motion Caption\": \"OptCamMotion\",\n \"Scene Abstract Caption\": \"SceneSummary\",\n \"Main Motion Trend Summary\": \"MotionTrends\",\n \"Scene Keywords\": \"SceneTags\",\n \"Immersive Shot Summary\": \"ShotImmersion\",\n \"Qwen model\": \"LLM\"\n }\n\n # Initialize result dictionary with empty values\n result = {key: \"\" for key in labels.values()}\n current_label = None\n current_content = []\n\n # Process text line by line\n lines = text.split('\\n')\n i = 0\n while i < len(lines):\n line = lines[i].strip()\n\n # Check if current line contains any label\n for label, json_key in labels.items():\n if label in line:\n # Find position of first letter after the label\n start_pos = line.find(label) + len(label)\n # Skip non-alphabet characters\n while start_pos < len(line) and not line[start_pos].isalpha():\n start_pos += 1\n\n content = line[start_pos:].strip()\n\n # Process content if it exists after the label\n if content:\n if json_key in [\"MotionTrends\", \"SceneTags\"]:\n # Split by commas (both Chinese and English), preserve spaces in phrases\n items = re.split(r'[，,]\\s*', content)\n result[json_key] = [item.strip()\n for item in items if item.strip()]\n else:\n result[json_key] = content\n\n current_label = None\n current_content = []\n else:\n # No content after label, continue reading subsequent lines\n current_label = json_key\n current_content = []\n\n break\n else:\n # If collecting content for a label\n if current_label:\n # Check if line is not empty\n if line:\n # Find first letter in line\n start_pos = 0\n while start_pos < len(line) and not line[start_pos].isalpha():\n start_pos += 1\n\n if start_pos < len(line):\n current_content.append(line[start_pos:])\n else:\n # Empty line indicates end of current label content\n content = ' '.join(current_content).strip()\n\n if current_label in [\"MotionTrends\", \"SceneTags\"]:\n # Split by commas (both Chinese and English), preserve spaces in phrases\n items = re.split(r'[，,]\\s*', content)\n result[current_label] = [item.strip()\n for item in items if item.strip()]\n else:\n result[current_label] = content\n\n current_label = None\n current_content = []\n\n i += 1\n\n # Handle Qwen model label which might extend to the end\n if current_label == \"LLM\" and current_content:\n content = ' '.join(current_content).strip()\n result[current_label] = content\n\n return result\n\n\ndef vqa_parse_text_to_json(text):\n \"\"\"\n Parses text containing Camera Motion Caption and Scene Description into JSON format.\n \"\"\"\n result = {\n \"CamMotion\": \"\",\n \"SceneDesc\": \"\"\n }\n\n # Process Camera Motion Caption - from first letter after label to newline\n camera_pattern = r'Camera Motion Caption:\\s*(\\w[\\s\\S]*?)(?=\\n|$)'\n camera_match = re.search(camera_pattern, text)\n if camera_match:\n result[\"CamMotion\"] = camera_match.group(1).strip()\n\n # Process Scene Description - from first letter after label to end of text\n scene_pattern = r'Scene Description:\\s*(\\w[\\s\\S]*)$'\n scene_match = re.search(scene_pattern, text, re.DOTALL)\n if scene_match:\n result[\"SceneDesc\"] = scene_match.group(1).strip()\n\n return result\n\n\ndef process_single_row(args, clip_id):\n \"\"\"\n Processes VQA and LLM captions for a single clip and merges them into one JSON file.\n \"\"\"\n # Define file paths\n vqa_path = os.path.join(args.load_dir, \"VQA\", f\"{clip_id}.txt\")\n assert os.path.exists(vqa_path), f\"VQA path does not exist: {vqa_path}\"\n llm_path = os.path.join(args.load_dir, \"LLM\", f\"{clip_id}.txt\")\n assert os.path.exists(llm_path), f\"LLM path does not exist: {llm_path}\"\n output_path = os.path.join(args.output_dir, f\"{clip_id}.json\")\n\n # Skip if output file already exists\n if os.path.exists(output_path) and os.path.getsize(output_path) > 0:\n return\n\n # Read VQA file content\n with open(vqa_path, 'r', encoding='utf-8') as f:\n vqa_text = f.read()\n\n # Read LLM file content\n with open(llm_path, 'r', encoding='utf-8') as f:\n llm_text = f.read()\n\n # Parse text content to JSON\n vqa_json = vqa_parse_text_to_json(vqa_text)\n llm_json = parse_text_to_json(llm_text)\n\n # Merge JSON objects\n combined_json = {**vqa_json, **llm_json}\n\n # Save merged JSON to output file\n with open(output_path, 'w', encoding='utf-8') as f:\n json.dump(combined_json, f, ensure_ascii=False, indent=2)\n\n\ndef worker(args, task_queue, pbar):\n while True:\n try:\n idx, clip_id = task_queue.get(timeout=1)\n except queue.Empty:\n break\n\n process_single_row(args, clip_id)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments.\"\"\"\n parser = argparse.ArgumentParser(\n description='Merge VQA and LLM caption data')\n parser.add_argument('--csv_path', type=str,\n required=True, help='Path to the CSV file')\n parser.add_argument('--load_dir', type=str, required=True,\n help='Directory containing caption files')\n parser.add_argument('--output_dir', type=str, required=True,\n help='Directory to save merged JSON files')\n parser.add_argument('--num_workers', type=int,\n default=32, help='Number of worker threads')\n return parser.parse_args()\n\n\ndef main():\n \"\"\"\n Processes all scenes in the specified batch.\n \"\"\"\n args = parse_args()\n\n df = pd.read_csv(args.csv_path)\n\n os.makedirs(args.output_dir, exist_ok=True)\n\n # Use multiprocessing manager for thread-safe queue\n manager = Manager()\n task_queue = manager.Queue()\n\n # Add tasks to queue\n for index, row in df.iterrows():\n task_queue.put((index, row['id']))\n\n # Start multi-threaded processing with progress bar\n with tqdm(total=len(df), desc=\"Processing progress\") as pbar:\n with ThreadPoolExecutor(max_workers=args.num_workers) as executor:\n futures = []\n for _ in range(args.num_workers):\n futures.append(executor.submit(worker, args, task_queue, pbar))\n\n # Wait for all futures to complete\n for future in as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "docker-entrypoint.sh", "content": "#!/usr/bin/env bash\n# Simple entrypoint: activate venv if present and run provided command\nset -euo pipefail\n\nif [ -f \"/workspace/venv/bin/activate\" ]; then\n echo \"Activating venv\"\n # shellcheck disable=SC1091\n source /workspace/venv/bin/activate\nfi\n\nif [ \"$#\" -gt 0 ]; then\n exec \"$@\"\nelse\n exec bash\nfi\n" }, { "path": "requirements/requirements.txt", "content": "torch==2.7.0 --index-url https://download.pytorch.org/whl/cu126\r\ntorchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu126\r\ntorchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu126\r\nopencv-python==4.11.0.86\r\ntqdm==4.67.1\r\nimageio==2.37.0\r\neinops==0.8.1\r\nscipy==1.15.2\r\nmatplotlib==3.10.0\r\nninja==1.11.1.3\r\nnumpy==1.26.4\r\npandas==2.2.3\r\nhuggingface_hub\r\n" }, { "path": "requirements/requirements_annotation.txt", "content": "wandb==0.19.8\r\ntimm==1.0.15\r\nkornia==0.8.0\r\nxformers==0.0.30\r\ntorch_scatter==2.1.2\r\ngradio_imageslider==0.0.20\r\ngradio==4.29.0\r\n# sam2\r\nhydra-core==1.3.2\r\niopath==0.1.10\r\n\r\nOpenEXR\r\n" }, { "path": "requirements/requirements_scoring.txt", "content": "ftfy==6.3.1\ndiffusers==0.29.0\naccelerate==1.4.0\nav==14.2.0\nscenedetect==0.6.5.2\ndecord==0.6.0\nimageio-ffmpeg==0.6.0\nffmpeg-python==0.2.0\nclip @ git+https://github.com/openai/CLIP.git\ncpbd==1.0.7\n# paddlepaddle-gpu==3.0.0 --index-url https://www.paddlepaddle.org.cn/packages/stable/cu126/\npaddleocr==3.0.0\nnvidia-nccl-cu12==2.26.2\nnumpy==1.26.4" }, { "path": "scoring/README.md", "content": "# Scoring\n\n## Aesthetic Score\n\nTo evaluate the aesthetic quality of videos, we use the scoring model from [CLIP+MLP Aesthetic Score Predictor](https://github.com/christophschuhmann/improved-aesthetic-predictor). This model is trained on 176K SAC (Simulacra Aesthetic Captions) pairs, 15K LAION-Logos (Logos) pairs, and 250K AVA (The Aesthetic Visual Analysis) image-text pairs.\n\nThe aesthetic score is between 1 and 10, where 5.5 can be considered as the threshold for fair aesthetics, and 6.5 for high aesthetics. Good text-to-image models can achieve a score of 7.0 or higher.\n\nFirst, download the scoring model to `./checkpoints/aesthetic.pth`. Skip this step if you already follow the installation instructions in [README](../README.md).\n\n```bash\nwget https://github.com/christophschuhmann/improved-aesthetic-predictor/raw/main/sac+logos+ava1-l14-linearMSE.pth -O checkpoints/aesthetic.pth\n```\n\nThen, run the following command to compute aesthetic scores.\n```bash\ntorchrun --nproc_per_node ${GPU_NUM} scoring/aesthetic/inference.py \\\n ${ROOT_META}/clips_info.csv \\\n --bs 16 \\\n --num_workers ${NUM_WORKERS} \\\n --fig_load_dir ${ROOT_FIG}\n```\n\n## Luminance Score\n\nLuminance was calculated for the first, middle, and last frames using the standard formula $L = 0.2126 R + 0.7152 G + 0.0722 B$, where $R$, $G$, and $B$ are the respective channel values. Clips with average luminance outside the range [20, 140], either too dark or too bright, were excluded, ensuring that only videos with proper exposure were retained.\n\nRun the following command to compute luminance scores.\n```bash\ntorchrun --nproc_per_node ${GPU_NUM} scoring/luminance/inference.py \\\n ${ROOT_META}/clips_info.csv \\\n --bs 16 \\\n --num_workers ${NUM_WORKERS} \\\n --fig_load_dir ${ROOT_FIG}\n```\n\n## Motion Score\nConventional motion analysis using optical flow is computationally expensive and less effective for videos with complex motion patterns. Inspired by Open-Sora 2.0, we adopted a lightweight VMAF-based motion analysis method integrated into FFMPEG. This method yields a motion score between 0 and 20.\n\nClips with scores outside the valid range of [2, 14], either too static (scores $<$ 2) or excessively chaotic (scores $>$ 14), were filtered out.\n\nRun the following command to compute motion scores.\n```bash\npython scoring/motion/inference.py ${ROOT_META}/clips_info.csv \\\n --temp_save_dir ${ROOT_TEMP} \\\n --num_workers $((GPU_NUM * 4)) \\\n --gpu_num ${GPU_NUM}\n```\n\n## OCR\nFor text detection, we used the latest release of PaddleOCR, which offers high accuracy and robust multilingual support. We processed the first, middle, and last frames of each clip to detect text regions, computing the ratio of text area to frame size. Clips where the text area exceeded 30% were removed, as these were considered informational rather than visual.\n\nRun the following command to compute OCR scores.\n```bash\npython scoring/ocr/inference.py ${ROOT_META}/clips_info.csv \\\n --fig_load_dir ${ROOT_FIG} \\\n --num_workers $((GPU_NUM * 4)) \\\n --gpu_num ${GPU_NUM}\n```" }, { "path": "scoring/__init__.py", "content": "" }, { "path": "scoring/aesthetic/__init__.py", "content": "" }, { "path": "scoring/aesthetic/inference.py", "content": "\"\"\"\nAesthetic scoring script for video frames using CLIP and MLP models.\nAdapted from https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/simple_inference.py\nCalculates aesthetic scores for video clips using distributed processing.\n\"\"\"\n\n# adapted from https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/main/simple_inference.py\nimport argparse\nimport gc\nimport os\nfrom glob import glob\nfrom datetime import timedelta\nfrom PIL import Image\nimport clip\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom torch.utils.data import DataLoader, DistributedSampler\nfrom tqdm import tqdm\n\n\ndef merge_scores(gathered_list: list, csv: pd.DataFrame, column):\n \"\"\"Merge aesthetic scores from all distributed processes.\"\"\"\n # Reorder results from all processes\n indices_list = list(map(lambda x: x[0], gathered_list))\n scores_list = list(map(lambda x: x[1], gathered_list))\n\n flat_indices = []\n for x in zip(*indices_list):\n flat_indices.extend(x)\n flat_scores = []\n for x in zip(*scores_list):\n flat_scores.extend(x)\n flat_indices = np.array(flat_indices)\n flat_scores = np.array(flat_scores)\n\n # Filter duplicates from distributed processing\n unique_indices, unique_indices_idx = np.unique(flat_indices, return_index=True)\n csv.loc[unique_indices, column] = flat_scores[unique_indices_idx]\n\n # Drop indices in csv not in unique_indices\n csv = csv.loc[unique_indices]\n return csv\n\n\nclass VideoTextDataset(torch.utils.data.Dataset):\n \"\"\"Dataset for loading video frames for aesthetic scoring.\"\"\"\n \n def __init__(self, csv_path, fig_load_dir, transform=None):\n self.csv_path = csv_path\n self.csv = pd.read_csv(csv_path)\n self.transform = transform\n self.fig_load_dir = fig_load_dir\n\n def __getitem__(self, index):\n \"\"\"Load and transform video frames for a single sample.\"\"\"\n sample = self.csv.iloc[index]\n\n # Load first 3 frames from video clip\n images_dir = os.path.join(self.fig_load_dir, sample[\"id\"])\n images = sorted(glob(f\"{images_dir}/img/*.jpg\"))[:3]\n\n # Apply CLIP preprocessing transforms\n images = [self.transform(Image.open(img).convert(\"RGB\")) for img in images]\n\n # Stack images into tensor\n images = torch.stack(images)\n\n return dict(index=index, images=images)\n\n def __len__(self):\n return len(self.csv)\n\n\nclass MLP(nn.Module):\n \"\"\"Multi-layer perceptron for aesthetic score prediction.\"\"\"\n \n def __init__(self, input_size):\n super().__init__()\n self.input_size = input_size\n self.layers = nn.Sequential(\n nn.Linear(self.input_size, 1024),\n nn.Dropout(0.2),\n nn.Linear(1024, 128),\n nn.Dropout(0.2),\n nn.Linear(128, 64),\n nn.Dropout(0.1),\n nn.Linear(64, 16),\n nn.Linear(16, 1),\n )\n\n def forward(self, x):\n return self.layers(x)\n\n\nclass AestheticScorer(nn.Module):\n \"\"\"Combined CLIP + MLP model for aesthetic scoring.\"\"\"\n \n def __init__(self, input_size, device):\n super().__init__()\n self.mlp = MLP(input_size)\n self.clip, self.preprocess = clip.load(\"ViT-L/14\", device=device)\n\n self.eval()\n self.to(device)\n\n def forward(self, x):\n \"\"\"Extract CLIP features and predict aesthetic scores.\"\"\"\n image_features = self.clip.encode_image(x)\n image_features = F.normalize(image_features, p=2, dim=-1).float()\n return self.mlp(image_features)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for aesthetic scoring.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the input CSV file\")\n parser.add_argument(\n \"--load_num\", type=int, default=4, help=\"Number of frames to load\"\n )\n parser.add_argument(\"--bs\", type=int, default=1024, help=\"Batch size\")\n parser.add_argument(\"--num_workers\", type=int, default=16, help=\"Number of workers\")\n parser.add_argument(\n \"--fig_load_dir\",\n type=str,\n required=True,\n help=\"Directory to load the extracted frames\",\n )\n parser.add_argument(\n \"--prefetch_factor\", type=int, default=3, help=\"Prefetch factor\"\n )\n parser.add_argument(\"--skip_if_existing\", action=\"store_true\")\n args = parser.parse_args()\n\n return args\n\n\ndef main():\n args = parse_args()\n\n csv_path = args.csv_path\n if not os.path.exists(csv_path):\n print(f\"CSV file '{csv_path}' not found. Exit.\")\n exit()\n\n wo_ext, ext = os.path.splitext(csv_path)\n out_path = f\"{wo_ext}_aes{ext}\"\n if args.skip_if_existing and os.path.exists(out_path):\n print(f\"Output CSV file '{out_path}' already exists. Exit.\")\n exit()\n\n # Initialize distributed processing\n dist.init_process_group(backend=\"nccl\", timeout=timedelta(hours=24))\n torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())\n\n # Build aesthetic scoring model\n device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n print(f\"Using device: {device}\")\n model = AestheticScorer(768, device)\n model.mlp.load_state_dict(\n torch.load(\"checkpoints/aesthetic.pth\", map_location=device)\n )\n preprocess = model.preprocess\n\n # Build dataset and dataloader\n dataset = VideoTextDataset(\n args.csv_path, transform=preprocess, fig_load_dir=args.fig_load_dir\n )\n dataloader = DataLoader(\n dataset,\n batch_size=args.bs,\n num_workers=args.num_workers,\n sampler=DistributedSampler(\n dataset,\n num_replicas=dist.get_world_size(),\n rank=dist.get_rank(),\n shuffle=False,\n drop_last=False,\n ),\n )\n\n # Compute aesthetic scores for all batches\n indices_list = []\n scores_list = []\n model.eval()\n for batch in tqdm(\n dataloader, disable=(dist.get_rank() != 0), position=dist.get_rank()\n ):\n indices = batch[\"index\"]\n images = batch[\"images\"].to(device, non_blocking=True)\n\n B = images.shape[0]\n images = rearrange(images, \"B N C H W -> (B N) C H W\")\n\n # Compute aesthetic scores using CLIP + MLP\n with torch.no_grad():\n scores = model(images)\n\n # Average scores across frames for each video\n scores = rearrange(scores, \"(B N) 1 -> B N\", B=B)\n scores = scores.mean(dim=1)\n scores_np = scores.to(torch.float32).cpu().numpy()\n\n indices_list.extend(indices.tolist())\n scores_list.extend(scores_np.tolist())\n\n # Wait for all ranks to finish data processing\n dist.barrier()\n\n # Gather results from all processes and save\n torch.cuda.empty_cache()\n gc.collect()\n gathered_list = [None] * dist.get_world_size()\n dist.all_gather_object(gathered_list, (indices_list, scores_list))\n if dist.get_rank() == 0:\n csv_new = merge_scores(gathered_list, dataset.csv, column=\"aesthetic score\")\n csv_new.to_csv(out_path, index=False)\n print(f\"New csv with aesthetic scores saved to '{out_path}'.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scoring/luminance/__init__.py", "content": "" }, { "path": "scoring/luminance/inference.py", "content": "\"\"\"\nLuminance analysis script for video frames using distributed processing.\nCalculates mean, min, and max luminance scores for video clips using PyTorch distributed computing.\n\"\"\"\n\nimport argparse\nimport os\nimport gc\nfrom glob import glob\nfrom datetime import timedelta\nfrom PIL import Image\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.distributed as dist\nfrom torch.utils.data import DataLoader, DistributedSampler\nfrom torchvision.transforms.functional import pil_to_tensor\nfrom tqdm import tqdm\n\n\ndef merge_scores(gathered_list: list, csv: pd.DataFrame):\n \"\"\"Merge luminance scores from all distributed processes.\"\"\"\n # Reorder results from all processes\n indices_list = list(map(lambda x: x[0], gathered_list))\n mean_scores_list = list(map(lambda x: x[1], gathered_list))\n min_scores_list = list(map(lambda x: x[2], gathered_list))\n max_scores_list = list(map(lambda x: x[3], gathered_list))\n\n flat_indices = []\n for x in zip(*indices_list):\n flat_indices.extend(x)\n flat_mean_scores = []\n for x in zip(*mean_scores_list):\n flat_mean_scores.extend(x)\n flat_min_scores = []\n for x in zip(*min_scores_list):\n flat_min_scores.extend(x)\n flat_max_scores = []\n for x in zip(*max_scores_list):\n flat_max_scores.extend(x)\n flat_indices = np.array(flat_indices)\n flat_mean_scores = np.array(flat_mean_scores)\n flat_min_scores = np.array(flat_min_scores)\n flat_max_scores = np.array(flat_max_scores)\n\n # Filter duplicates from distributed processing\n unique_indices, unique_indices_idx = np.unique(flat_indices, return_index=True)\n csv.loc[unique_indices, \"luminance mean\"] = flat_mean_scores[unique_indices_idx]\n csv.loc[unique_indices, \"luminance min\"] = flat_min_scores[unique_indices_idx]\n csv.loc[unique_indices, \"luminance max\"] = flat_max_scores[unique_indices_idx]\n\n # Drop indices in csv not in unique_indices\n csv = csv.loc[unique_indices]\n return csv\n\n\nclass VideoDataset(torch.utils.data.Dataset):\n \"\"\"Dataset to handle video luminance computation.\"\"\"\n\n def __init__(self, csv_path, fig_load_dir):\n self.csv_path = csv_path\n self.csv = pd.read_csv(csv_path)\n self.fig_load_dir = fig_load_dir\n\n def __getitem__(self, index):\n \"\"\"Get video frames and compute luminance for a single sample.\"\"\"\n sample = self.csv.iloc[index]\n\n # Load first 3 frames from video clip\n images_dir = os.path.join(self.fig_load_dir, sample[\"id\"])\n images = sorted(glob(f\"{images_dir}/img/*.jpg\"))[:3]\n\n # Transform images to tensors\n images = torch.stack(\n [pil_to_tensor(Image.open(img).convert(\"RGB\")) for img in images]\n )\n\n return {\"index\": index, \"images\": images}\n\n def __len__(self):\n return len(self.csv)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for luminance analysis.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the input CSV file\")\n parser.add_argument(\"--bs\", type=int, default=4, help=\"Batch size\")\n parser.add_argument(\"--num_workers\", type=int, default=16, help=\"Number of workers\")\n parser.add_argument(\n \"--fig_load_dir\",\n type=str,\n required=True,\n help=\"Directory to load the extracted frames\",\n )\n parser.add_argument(\"--skip_if_existing\", action=\"store_true\")\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n csv_path = args.csv_path\n\n if not os.path.exists(csv_path):\n print(f\"csvdata file '{csv_path}' not found. Exiting.\")\n return\n\n output_path = csv_path.replace(\".csv\", \"_lum.csv\")\n if args.skip_if_existing and os.path.exists(output_path):\n print(f\"Output '{output_path}' already exists. Exiting.\")\n return\n\n # Initialize distributed processing\n dist.init_process_group(backend=\"nccl\", timeout=timedelta(hours=24))\n (\n torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())\n if torch.cuda.is_available()\n else None\n )\n\n # Setup dataset and distributed dataloader\n dataset = VideoDataset(csv_path, fig_load_dir=args.fig_load_dir)\n dataloader = DataLoader(\n dataset,\n batch_size=args.bs,\n num_workers=args.num_workers,\n sampler=DistributedSampler(\n dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank()\n ),\n )\n\n # Process batches and calculate luminance scores\n indices_list = []\n mean_scores_list = []\n max_scores_list = []\n min_scores_list = []\n device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n for batch in tqdm(\n dataloader, disable=(dist.get_rank() != 0), position=dist.get_rank()\n ):\n indices = batch[\"index\"]\n images = batch[\"images\"].to(device, non_blocking=True) # [B, N, C, H, W]\n\n # Calculate luminance using standard RGB weights\n R, G, B = images[:, :, 0], images[:, :, 1], images[:, :, 2]\n luminance = 0.2126 * R + 0.7152 * G + 0.0722 * B\n scores = luminance.mean(dim=[2, 3])\n\n # Compute statistics across frames\n mean_scores = scores.mean(dim=1).cpu().numpy()\n max_scores = scores.max(dim=1)[0].cpu().numpy()\n min_scores = scores.min(dim=1)[0].cpu().numpy()\n\n indices_list.extend(indices.tolist())\n mean_scores_list.extend(mean_scores.tolist())\n max_scores_list.extend(max_scores.tolist())\n min_scores_list.extend(min_scores.tolist())\n\n # Wait for all ranks to finish data processing\n dist.barrier()\n\n # Gather results from all processes and save\n torch.cuda.empty_cache()\n gc.collect()\n gathered_list = [None] * dist.get_world_size()\n dist.all_gather_object(\n gathered_list,\n (indices_list, mean_scores_list, min_scores_list, max_scores_list),\n )\n if dist.get_rank() == 0:\n csv_new = merge_scores(gathered_list, dataset.csv)\n csv_new.to_csv(output_path, index=False)\n print(f\"New csv with luminance scores saved to '{output_path}'\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scoring/motion/INSTALL.md", "content": "# Compiling FFmpeg with NVIDIA GPU Acceleration and VMAF on Ubuntu\n\nThis guide provides a comprehensive walkthrough for compiling FFmpeg from source on an Ubuntu system equipped with an NVIDIA GPU. The resulting build will support NVIDIA's hardware encoding/decoding (NVENC/DEC), NPP filters (NVIDIA Performance Primitives), and CUDA-based VMAF (Video Multi-Method Assessment Fusion) for video quality assessment.\n\n\n## Environment and Versions\n\nBefore you begin, ensure your system environment is similar to the configuration below. Version matching is crucial for a successful compilation.\nThe GPU needs to support HEVC; refer to the [NVIDIA NVDEC Support Matrix](https://en.wikipedia.org/wiki/NVIDIA_Video_Coding_Engine#NVDEC).\n\n- **GPU**: NVIDIA GeForce RTX 4090 or other compatible models\n- **OS**: Ubuntu 22.04 \n- **NVIDIA Driver Version**: A version compatible with CUDA 12.6\n- **CUDA Version (from `nvidia-smi`)**: `12.x`\n- **CUDA Toolkit Version**: `12.6` (This is the version used for compilation)\n- **Target FFmpeg Version**: `6.1`\n\n**Key Tip**: The version of the `NVIDIA Codec Headers` (`ffnvcodec`) must be compatible with the `CUDA Toolkit` version installed on your system and the version of `FFmpeg` you intend to compile.\n\n## Compilation Steps\n\nPlease follow these steps in order.\n\n\n### Step 1: Install System Dependencies\n\nUpdate system packages and install required development tools and libraries:\n\n```bash\nsudo apt-get update\nsudo DEBIAN_FRONTEND=noninteractive apt-get install -y \\\n libopenjp2-7-dev \\\n ninja-build \\\n cmake \\\n git \\\n python3 \\\n python3-pip \\\n nasm \\\n xxd \\\n pkg-config \\\n curl \\\n unzip \\\n ca-certificates \\\n libnuma-dev \\\n libsm6 \\\n libxext6 \\\n libxrender1 \\\n libgl1 \\\n vim \\\n nvidia-cuda-toolkit\n```\n\n\n### Step 2: Clone Required Repositories\n\n```bash\n# Create a working directory (custom path allowed)\nmkdir -p ~/ffmpeg-build && cd ~/ffmpeg-build\n\n# Clone nv-codec-headers (NVIDIA codec headers)\ngit clone https://github.com/FFmpeg/nv-codec-headers.git\n\n# Clone libvmaf (video quality assessment library)\ngit clone https://github.com/Netflix/vmaf.git\ncd vmaf && git checkout master # Switch to master branch (modify version if needed)\ncd ..\n\n# Clone FFmpeg source code\ngit clone https://github.com/FFmpeg/FFmpeg.git\ncd FFmpeg && git checkout master # Switch to master branch (modify version if needed)\ncd ..\n```\n\n\n### Step 3: Install nv-codec-headers\n\n```bash\ncd nv-codec-headers\nmake\nsudo make install\ncd ..\n```\n\n\n### Step 4: Compile and Install libvmaf (with CUDA Support)\n\n1. Install the meson build tool:\n ```bash\n python3 -m pip install meson\n ```\n\n2. Compile and install libvmaf:\n ```bash\n cd vmaf\n meson libvmaf/build libvmaf \\\n -Denable_cuda=true \\\n -Denable_avx512=true \\\n --buildtype release\n ninja -vC libvmaf/build\n sudo ninja -vC libvmaf/build install\n cd ..\n ```\n\n3. Update system library cache:\n ```bash\n export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/x86_64-linux-gnu/\n sudo ldconfig\n ```\n\n\n### Step 5: Compile and Install FFmpeg (with NVIDIA and libvmaf Support)\n\n```bash\ncd FFmpeg\n\n# Configure compilation options (enable CUDA, NVENC, NVDEC, and libvmaf)\n./configure \\\n --enable-libnpp \\\n --enable-nonfree \\\n --enable-nvdec \\\n --enable-nvenc \\\n --enable-cuvid \\\n --enable-cuda \\\n --enable-cuda-nvcc \\\n --enable-libvmaf \\\n --enable-ffnvcodec \\\n --disable-stripping \\\n --extra-cflags=\"-I/usr/local/cuda/include\" \\\n --extra-ldflags=\"-L/usr/local/cuda/lib64 -L/usr/local/cuda/lib64/stubs/\"\n\n# Compile (adjust the number after -j based on CPU cores for faster compilation)\nmake -j$(nproc)\n\n# Install\nsudo make install\n\ncd ..\n```\n\n\n### Step 6: Configure Python Environment\n\n1. Upgrade pip and set up links:\n ```bash\n sudo ln -sf /usr/bin/python3 /usr/bin/python\n python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel\n ```\n\n2. Install Python dependencies (assuming project code is cloned locally; replace with actual path):\n ```bash\n # Navigate to the project root directory\n cd /path/to/your/project\n\n # Install dependencies\n python3 -m pip --no-cache-dir install -r requirements/requirements.txt\n python3 -m pip --no-cache-dir install -r requirements/requirements_scoring.txt || true\n python3 -m pip --no-cache-dir install -r requirements/requirements_annotation.txt || true\n ```\n\n\n### Step 7: Verify Installation\n\n1. Check FFmpeg version and configuration:\n ```bash\n ffmpeg -version\n ffmpeg -encoders | grep nvenc # Verify NVENC support\n ffmpeg -decoders | grep nvdec # Verify NVDEC support\n ffmpeg -filters | grep vmaf # Verify libvmaf support\n ```\n\n2. If all the above commands output corresponding content correctly, the installation is successful.\n\n## Troubleshooting\n\n### Issue 1: VMAF compilation fails with `vcs_version.h: No such file or directory`\n\n- **Cause**: This error typically occurs if you downloaded the VMAF source code as a ZIP archive instead of using `git clone`. The build script relies on the `.git` directory to generate version header files.\n- **Solution**: Always use `git clone` to get the source code.\n ```bash\n git clone https://github.com/Netflix/vmaf.git\n ```\n\n### Issue 2: FFmpeg `configure` fails with error about Video Codec SDK version being too low\n\n- **Error Message**: Something like `ERROR: nvenc requested, but NVIDIA Video Codec SDK 12.1 or later is required.` (The version number may vary).\n- **Cause**: This means the version of `nv-codec-headers` you checked out is not compatible with your NVIDIA driver, CUDA Toolkit, or the version of FFmpeg you are building.\n- **Solution**:\n 1. Carefully re-check your [NVIDIA Driver](https://www.nvidia.com/Download/index.aspx) and [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit-archive) versions.\n 2. Go back to [Step 3: Install NVIDIA Codec Headers](#step-3-install-nvidia-codec-headers) and ensure you `git checkout` the branch that best matches your environment (e.g., `sdk/12.6`).\n 3. Consult the [Official NVIDIA FFmpeg Guide](https://docs.nvidia.com/video-technologies/video-codec-sdk/ffmpeg-with-nvidia-gpu/index.html) or the `nv-codec-headers` repository to confirm version compatibility.\n\n## References\n\n- [VMAF on GitHub](https://github.com/Netflix/vmaf)\n- [FFmpeg Official Source](https://github.com/FFmpeg/FFmpeg/tree/release/6.1)\n- [NVIDIA Codec Headers Source](https://github.com/FFmpeg/nv-codec-headers/tree/sdk/12.6)\n- [Official NVIDIA Guide for Compiling FFmpeg](https://docs.nvidia.com/video-technologies/video-codec-sdk/ffmpeg-with-nvidia-gpu/index.html)\n" }, { "path": "scoring/motion/__init__.py", "content": "" }, { "path": "scoring/motion/inference.py", "content": "\"\"\"\nMotion analysis script for video quality assessment using FFmpeg and VMAF.\nCalculates motion scores for video clips using hardware acceleration when available.\n\"\"\"\n\nimport os\nimport argparse\nimport pandas as pd\nimport subprocess\nfrom multiprocessing import Manager\nimport queue\nimport concurrent.futures\nfrom tqdm import tqdm\n\nFFMPEG_PATH = \"/usr/local/bin/ffmpeg\"\n\n\ndef get_ffmpeg_acceleration():\n \"\"\"\n Auto detect the best acceleration method.\n Priority: NVIDIA GPU > CPU.\n \"\"\"\n try:\n # Get the list of ffmpeg configuration\n output = subprocess.check_output(\n [FFMPEG_PATH, \"-version\"], stderr=subprocess.DEVNULL\n ).decode(\"utf-8\")\n\n if \"--enable-cuda-nvcc\" in output and \"--enable-libvmaf\" in output:\n return \"nvidia\"\n else:\n return \"cpu\" # Use CPU\n except Exception as e:\n print(f\"FFmpeg acceleration detection failed: {e}\")\n return \"cpu\"\n\n\nACCELERATION_TYPE = get_ffmpeg_acceleration()\nprint(f\"FFmpeg acceleration type: {ACCELERATION_TYPE}\")\n\n\ndef process_single_row(video_path, args, process_id):\n \"\"\"Process a single video to generate motion analysis CSV using FFmpeg.\"\"\"\n path = os.path.join(\n args.temp_save_dir, os.path.basename(video_path).split(\".\")[0] + \".csv\"\n )\n\n # Build FFmpeg command with appropriate acceleration\n command = [FFMPEG_PATH]\n if ACCELERATION_TYPE == \"nvidia\":\n command += [\n \"-hwaccel\",\n \"cuda\",\n \"-hwaccel_output_format\",\n \"cuda\",\n \"-hwaccel_device\",\n f\"{process_id % args.gpu_num}\",\n ]\n command += [\"-i\", f\"{video_path}\"]\n if ACCELERATION_TYPE == \"nvidia\":\n command += [\n \"-hwaccel\",\n \"cuda\",\n \"-hwaccel_output_format\",\n \"cuda\",\n \"-hwaccel_device\",\n f\"{process_id % args.gpu_num}\",\n ]\n command += [\"-i\", f\"{video_path}\"]\n if ACCELERATION_TYPE == \"nvidia\":\n command += [\n \"-filter_complex\",\n f\"[0:v]scale_cuda=format=yuv420p[dis],[1:v]scale_cuda=format=yuv420p[ref],[dis][ref]libvmaf_cuda=log_fmt=csv:log_path={path}\",\n ]\n else:\n command += [\"-lavfi\", f\"libvmaf=log_fmt=csv:log_path={path}\"]\n command += [\"-f\", \"null\", \"-\"]\n try:\n result = subprocess.run(command, capture_output=True, text=True, check=True)\n except subprocess.CalledProcessError as e:\n print(f\"Error: {e.stderr}\")\n\n\ndef calculate_score(row, args):\n \"\"\"Calculate motion score for a specific video clip segment.\"\"\"\n csv_path = os.path.join(args.temp_save_dir, f'{row[\"id_ori\"]}.csv')\n df = pd.read_csv(csv_path)\n df = df[(df[\"Frame\"] >= row[\"frame_start\"]) & (df[\"Frame\"] <= row[\"frame_end\"])]\n mean_value = df[\"integer_motion2\"].mean()\n return mean_value\n\n\ndef worker1(task_queue, progress_queue, args, process_id):\n \"\"\"Worker function for processing videos in parallel.\"\"\"\n while True:\n try:\n video_path = task_queue.get(timeout=1)\n except queue.Empty:\n break\n process_single_row(video_path, args, process_id)\n progress_queue.put(video_path)\n task_queue.task_done()\n\n\ndef worker2(task_queue, results_queue, args):\n \"\"\"Worker function for calculating motion scores in parallel.\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n value = calculate_score(row, args)\n results_queue.put((index, value))\n task_queue.task_done()\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for motion analysis.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--csv_path\", type=str, required=True, help=\"Path to the CSV file\")\n parser.add_argument(\n \"--temp_save_dir\",\n type=str,\n required=True,\n help=\"Directory to save the temporary files\",\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=None, help=\"#workers for concurrent.futures\"\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"disable parallel processing\"\n )\n parser.add_argument(\"--gpu_num\", type=int, default=1, help=\"gpu number\")\n parser.add_argument(\"--skip_if_existing\", action=\"store_true\")\n args = parser.parse_args()\n return args\n\n\ndef main():\n args = parse_args()\n\n wo_ext, ext = os.path.splitext(args.csv_path)\n out_path = f\"{wo_ext}_motion{ext}\"\n if args.skip_if_existing and os.path.exists(out_path):\n print(f\"Output CSV file '{out_path}' already exists. Exit.\")\n exit()\n\n df = pd.read_csv(args.csv_path)\n video_paths = df[\"video_path\"].unique()\n\n if args.disable_parallel:\n # Sequential processing\n results = []\n for video_path in tqdm(video_paths, desc=\"Processing videos\"):\n result = process_single_row(video_path, args, 0)\n results.append(result)\n\n for index, row in tqdm(\n df.iterrows(), total=len(df), desc=\"Calculating scores\"\n ):\n result = calculate_score(row, args)\n df.at[index, \"motion\"] = result\n else:\n # Parallel processing\n if args.num_workers is not None:\n num_workers = args.num_workers\n else:\n num_workers = os.cpu_count() or 1\n\n # First phase: process videos to generate CSV files\n manager = Manager()\n task_queue = manager.Queue()\n progress_queue = manager.Queue()\n\n for video_path in video_paths:\n task_queue.put(video_path)\n\n with concurrent.futures.ProcessPoolExecutor(\n max_workers=num_workers\n ) as executor:\n futures = []\n for id in range(num_workers):\n futures.append(\n executor.submit(worker1, task_queue, progress_queue, args, id)\n )\n\n processed = 0\n total_video_tasks = len(video_paths)\n with tqdm(total=total_video_tasks, desc=\"Processing videos\") as pbar:\n while processed < total_video_tasks:\n try:\n progress_queue.get(timeout=1)\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and progress_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n # Second phase: calculate motion scores\n result_queue = manager.Queue()\n task_queue = manager.Queue()\n\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n with concurrent.futures.ProcessPoolExecutor(\n max_workers=num_workers\n ) as executor:\n futures = []\n for _ in range(num_workers):\n futures.append(executor.submit(worker2, task_queue, result_queue, args))\n\n results = []\n processed = 0\n total_score_tasks = len(df)\n with tqdm(total=total_score_tasks, desc=\"Calculating scores\") as pbar:\n while processed < total_score_tasks:\n try:\n results.append(result_queue.get(timeout=1))\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and result_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n # Collect and sort results\n while not result_queue.empty():\n results.append(result_queue.get())\n results.sort(key=lambda x: x[0])\n results = list(map(lambda x: x[1], results))\n df[\"motion score\"] = results\n\n df.to_csv(out_path, index=False)\n print(f\"New df with motion scores saved to '{out_path}'.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scoring/ocr/__init__.py", "content": "" }, { "path": "scoring/ocr/inference.py", "content": "\"\"\"\nOCR analysis script for video frames using PaddleOCR.\nCalculates text area ratios for video clips using distributed processing.\n\"\"\"\n\nimport os\nfrom glob import glob\nimport argparse\nimport pandas as pd\nfrom multiprocessing import Manager\nimport queue\nimport concurrent.futures\nfrom tqdm import tqdm\nimport cv2\nfrom paddleocr import PaddleOCR\n\n\ndef process_single_row(row, args, model):\n \"\"\"Process a single row to calculate OCR text area ratio.\"\"\"\n img_dir = os.path.join(args.fig_load_dir, row[\"id\"])\n img_list = sorted(glob(f\"{img_dir}/img/*.jpg\"))[:3]\n # Load images\n images = [cv2.imread(img_path) for img_path in img_list]\n images = [img for img in images if img is not None]\n if not images:\n return 0.0\n\n result = model.predict(input=images)\n area = images[0].shape[0] * images[0].shape[1] # Image area\n\n area_list = []\n for res in result:\n total_text_area = 0 # Initialize total text area\n for rec_box in res[\"rec_boxes\"]:\n x_min, y_min, x_max, y_max = (\n float(rec_box[0]),\n float(rec_box[1]),\n float(rec_box[2]),\n float(rec_box[3]),\n ) # Extract top-left and bottom-right coordinates\n text_area = (x_max - x_min) * (y_max - y_min) # Calculate text area\n total_text_area += text_area\n ratio = total_text_area / area\n area_list.append(ratio)\n return (\n max(area_list) if area_list else 0.0\n ) # Return max area ratio, 0.0 if no text detected\n\n\ndef worker(task_queue, result_queue, args, id):\n \"\"\"Worker function for multiprocessing OCR inference.\"\"\"\n gpu_id = id % args.gpu_num\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id) # Bind to specific GPU\n device = \"gpu:0\" # if torch.cuda.is_available() else \"cpu\"\n\n # Initialize PaddleOCR model with disabled orientation and unwarping features\n model = PaddleOCR(\n device=device,\n use_doc_orientation_classify=False, # Disable document orientation classification\n use_doc_unwarping=False, # Disable text image correction\n use_textline_orientation=False, # Disable text line orientation classification\n )\n\n while True:\n try:\n index, row = task_queue.get_nowait()\n except queue.Empty:\n break\n\n area_list = process_single_row(row, args, model)\n result_queue.put((index, area_list))\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for OCR inference.\"\"\"\n parser = argparse.ArgumentParser(description=\"SAM2 Image Predictor\")\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\n \"--fig_load_dir\", type=str, default=\"img\", help=\"Directory containing images\"\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=16, help=\"#workers for concurrent.futures\"\n )\n parser.add_argument(\"--gpu_num\", type=int, default=1, help=\"gpu number\")\n parser.add_argument(\"--skip_if_existing\", action=\"store_true\")\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n if not os.path.exists(args.csv_path):\n print(f\"csv file '{args.csv_path}' not found. Exit.\")\n return\n\n wo_ext, ext = os.path.splitext(args.csv_path)\n out_path = f\"{wo_ext}_ocr{ext}\"\n if args.skip_if_existing and os.path.exists(out_path):\n print(f\"Output csv file '{out_path}' already exists. Exit.\")\n exit()\n\n df = pd.read_csv(args.csv_path)\n results = []\n\n if args.disable_parallel:\n # Sequential processing\n model = PaddleOCR(\n device=\"gpu:0\", # if torch.cuda.is_available() else \"cpu\"\n use_doc_orientation_classify=False, # Disable document orientation classification\n use_doc_unwarping=False, # Disable text image correction\n use_textline_orientation=False, # Disable text line orientation classification\n )\n ocr_scores = []\n for index, row in tqdm(df.iterrows(), total=len(df), desc=\"Processing rows\"):\n score = process_single_row(row, args, model)\n ocr_scores.append(score)\n results.append((index, score))\n else:\n # Set up multiprocessing queues\n manager = Manager()\n task_queue = manager.Queue()\n result_queue = manager.Queue()\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n # Process tasks with multiple workers\n with concurrent.futures.ProcessPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for id in range(args.num_workers):\n futures.append(\n executor.submit(worker, task_queue, result_queue, args, id)\n )\n\n processed = 0\n total_tasks = len(df)\n with tqdm(total=total_tasks, desc=\"Processing rows\") as pbar:\n while processed < total_tasks:\n try:\n results.append(result_queue.get(timeout=1))\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and result_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n # Collect and sort results\n while not result_queue.empty():\n index, area_list = result_queue.get()\n results.append((index, area_list))\n\n results.sort(key=lambda x: x[0])\n df[\"ocr score\"] = [x[1] for x in results]\n\n df.to_csv(out_path, index=False)\n print(f\"New csv (shape={df.shape}) with ocr results saved to '{out_path}'.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "scripts/annotation.sh", "content": "#!/bin/bash\nCSV=[Replace with the path to the CSV file generated in the scoring step]\nOUTPUT_DIR=[Replace with the path to your output directory]\nmkdir -p ${OUTPUT_DIR}\n\nCUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nGPU_NUM=8\nENHANCED=true # Set to true to enable enhanced instruction generation\n\nmeasure_time() {\n local step_number=$1\n shift\n local green=\"\\e[32m\"\n local red=\"\\e[31m\"\n local no_color=\"\\e[0m\"\n local yellow=\"\\e[33m\"\n \n start_time=$(date +%s)\n echo -e \"${green}Step ${step_number} started at: $(date)${no_color}\"\n\n \"$@\"\n\n end_time=$(date +%s)\n echo -e \"${red}Step ${step_number} finished at: $(date)${no_color}\"\n echo -e \"${yellow}Duration: $((end_time - start_time)) seconds${no_color}\"\n echo \"---------------------------------------\"\n}\n\n# 1. Extract frames\nmeasure_time 1 python utils/extract_frames.py \\\n --csv_path ${CSV} \\\n --output_dir ${OUTPUT_DIR} \\\n --num_workers $((GPU_NUM * 2)) \\\n --target_size \"1280*720\" \\\n --backend \"opencv\" \\\n --interval 0.2\n\n# 2.1 Depth Estimation with Depth-Anything\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 2.1 torchrun --standalone --nproc_per_node ${GPU_NUM} camera_pose_annotation/depth_estimation/Depth-Anything/inference_batch.py \\\n --csv_path ${CSV} \\\n --encoder vitl \\\n --checkpoints_path checkpoints \\\n --output_dir ${OUTPUT_DIR} \\\n --bs 16 \\\n --num_workers ${GPU_NUM}\n\n# 2.2 Depth Estimation with UniDepth\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 2.2 torchrun --standalone --nproc_per_node ${GPU_NUM} camera_pose_annotation/depth_estimation/UniDepth/inference_batch.py \\\n --csv_path ${CSV} \\\n --output_dir ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints \\\n --bs 32 \\\n --num_workers ${GPU_NUM}\n\n# 3. Camera Tracking\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 3 python camera_pose_annotation/camera_tracking/inference_batch.py \\\n --csv_path ${CSV} \\\n --dir_path ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints \\\n --gpu_id ${CUDA_VISIBLE_DEVICES} \\\n --num_workers $((GPU_NUM * 2))\n\n# 4.1 CVD Optimization Preprocess\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 4.1 python camera_pose_annotation/cvd_opt/preprocess/inference_batch.py \\\n --csv_path ${CSV} \\\n --dir_path ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints \\\n --gpu_id ${CUDA_VISIBLE_DEVICES} \\\n --num_workers $((GPU_NUM * 2))\n\n# 4.2 CVD Optimization\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 4.2 python camera_pose_annotation/cvd_opt/inference_batch.py \\\n --csv_path ${CSV} \\\n --dir_path ${OUTPUT_DIR} \\\n --gpu_id ${CUDA_VISIBLE_DEVICES} \\\n --num_workers $((GPU_NUM * 2))\n # --only_depth\n\n# 5. Dynamic Mask Prediction\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 5 python camera_pose_annotation/dynamic_mask/inference_batch.py \\\n --csv_path ${CSV} \\\n --dir_path ${OUTPUT_DIR} \\\n --checkpoints_path checkpoints \\\n --gpu_num ${GPU_NUM} \\\n --num_workers $((GPU_NUM * 2))\n\n# 6. Evaluation of the results\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 6 python utils/evaluation.py \\\n --csv_path ${CSV} \\\n --dir_path ${OUTPUT_DIR} \\\n --gpu_num ${GPU_NUM} \\\n --num_workers $((GPU_NUM * 2)) \\\n --output_path ${OUTPUT_DIR}/final_results.csv\n\n# 7. Get motion instructions\nif [ \"$ENHANCED\" = false ] ; then\n CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 7 python utils/get_instructions.py \\\n --csv_path ${CSV} \\\n --dir_path ${OUTPUT_DIR} \\\n --interval 2 \\\n --num_workers $((GPU_NUM * 2))\nelse\n echo \"Standard instruction generation is enabled.\"\n CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 7 python utils/get_instructions_enhanced.py \\\n --csv_path ${OUTPUT_DIR}/final_results.csv \\\n --dir_path ${OUTPUT_DIR} \\\n --num_workers $((GPU_NUM * 2))\nfi\n\n# 8. Normalize the intrinsics\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 8 python utils/normalize_intrinsics.py \\\n --csv_path ${CSV} \\\n --dir_path ${OUTPUT_DIR} \\\n --num_workers $((GPU_NUM * 2))\n\n# [Optional] Convert the output poses.npy into a c2w/w2c matrix\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 9 python utils/quat_to_mat.py \\\n --csv_path ${CSV} \\\n --format c2w \\\n --dir_path ${OUTPUT_DIR} \\\n --num_workers $((GPU_NUM * 2))\n" }, { "path": "scripts/caption.sh", "content": "#!/bin/bash\nCSV=[Replace with the path to the result CSV file generated in the annotation step]\nSRC_DIR=[Replace with the path to the annotation output directory]\nOUTPUT_DIR=[Replace with the path to your output directory]\nmkdir -p ${OUTPUT_DIR}\n\nnum_workers=8\nwait_time=1\n\n# VQA\nvqa_prompt_file=caption/VQA/prompt.txt\nvqa_model=gemini-2.0-flash\nvqa_api_key=[Replace with your api key]\nvqa_base_domain=https://generativelanguage.googleapis.com/\n\n# LLM\nllm_prompt_dir=caption/LLM\nllm_model=qwen3-30b-a3b\nllm_api_key=[Replace with your api key]\nllm_base_domain=https://dashscope.aliyuncs.com/compatible-mode/\n\n# Tagging\ntag_prompt_file=caption/tagging/prompt.txt\ntag_model=qwen3-30b-a3b\ntag_api_key=[Replace with your api key]\ntag_base_domain=https://dashscope.aliyuncs.com/compatible-mode/\n\nmeasure_time() {\n local step_number=$1\n shift\n local green=\"\\e[32m\"\n local red=\"\\e[31m\"\n local no_color=\"\\e[0m\"\n local yellow=\"\\e[33m\"\n \n start_time=$(date +%s)\n echo -e \"${green}Step $step_number started at: $(date)${no_color}\"\n\n \"$@\"\n\n end_time=$(date +%s)\n echo -e \"${red}Step $step_number finished at: $(date)${no_color}\"\n echo -e \"${yellow}Duration: $((end_time - start_time)) seconds${no_color}\"\n echo \"---------------------------------------\"\n}\n\n# 1. VQA caption\nmeasure_time 1 python caption/VQA/inference.py \\\n --csv_path ${CSV} \\\n --fig_load_dir ${SRC_DIR} \\\n --output_dir ${OUTPUT_DIR} \\\n --prompt_file ${vqa_prompt_file} \\\n --model ${vqa_model} \\\n --api_key ${vqa_api_key} \\\n --base_domain ${vqa_base_domain} \\\n --num_workers ${num_workers} \\\n --wait_time ${wait_time}\n\n# 2. LLM caption\nmeasure_time 2 python caption/LLM/inference.py \\\n --csv_path $CSV \\\n --pose_load_dir $SRC_DIR \\\n --output_dir $OUTPUT_DIR \\\n --prompt_dir $llm_prompt_dir \\\n --model $llm_model \\\n --api_key $llm_api_key \\\n --num_workers $num_workers \\\n --base_domain $llm_base_domain \\\n --wait_time $wait_time\n\n# 3. Combine results\nmeasure_time 3 python caption/utils/combine.py \\\n --csv_path $CSV \\\n --load_dir $OUTPUT_DIR \\\n --output_dir $OUTPUT_DIR/results \\\n --num_workers $num_workers\n\n# 4. Add tags\npython caption/tagging/inference.py \\\n --csv_path $CSV \\\n --json_load_dir $OUTPUT_DIR/results \\\n --prompt_file $tag_prompt_file \\\n --model $tag_model \\\n --api_key $tag_api_key \\\n --num_workers $num_workers \\\n --base_domain $tag_base_domain \\\n --wait_time $wait_time" }, { "path": "scripts/docker_prepulls.sh", "content": "#!/usr/bin/env bash\n# This script pre-pulls and tags GPU-related Docker images from specified registries.\n\nset -euo pipefail\n\n# Minimal script: pre-pull three images (builder/runtime/buildkit) and tag them to\n# canonical names so downstream scripts can rely on the expected tags.\n\n# You can override these by setting the env vars before running this script.\nBUILDER_IMAGE=${BUILDER_IMAGE:-swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04}\nRUNTIME_IMAGE=${RUNTIME_IMAGE:-swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/nvidia/cuda:12.6.3-runtime-ubuntu22.04}\nBUILDKIT_IMAGE=${BUILDKIT_IMAGE:-swr.cn-north-4.myhuaweicloud.com/ddn-k8s/docker.io/moby/buildkit:buildx-stable-1}\n\nretry_pull() {\n local img=\"$1\"\n for i in 1 2 3; do\n echo \"pull attempt $i for ${img}...\"\n if docker pull \"${img}\"; then\n echo \"pulled ${img}\"\n return 0\n fi\n sleep $((i * 2))\n done\n echo \"Failed to pull ${img} after retries\" >&2\n return 1\n}\n\necho \"Pre-pulling images...\"\necho \"- builder: ${BUILDER_IMAGE}\"\necho \"- runtime: ${RUNTIME_IMAGE}\"\necho \"- buildkit: ${BUILDKIT_IMAGE}\"\n\nretry_pull \"${BUILDER_IMAGE}\" || true\nretry_pull \"${RUNTIME_IMAGE}\" || true\nretry_pull \"${BUILDKIT_IMAGE}\" || true\n\nCANONICAL_BUILDKIT_TAG=\"moby/buildkit:buildx-stable-1\"\nif docker image inspect \"${BUILDKIT_IMAGE}\" >/dev/null 2>&1; then\n echo \"Tagging ${BUILDKIT_IMAGE} -> ${CANONICAL_BUILDKIT_TAG} (local only)\"\n docker tag \"${BUILDKIT_IMAGE}\" \"${CANONICAL_BUILDKIT_TAG}\" || true\nfi\n\n# Also tag the mirrored CUDA images to the original docker.io names expected by\n# Dockerfiles and other scripts. This lets downstream tooling refer to\n# docker.io/nvidia/cuda:12.6.3-... even when images were pulled from a mirror.\nORIG_BUILDER_TAG=\"docker.io/nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04\"\nORIG_RUNTIME_TAG=\"docker.io/nvidia/cuda:12.6.3-runtime-ubuntu22.04\"\nif docker image inspect \"${BUILDER_IMAGE}\" >/dev/null 2>&1; then\n echo \"Tagging ${BUILDER_IMAGE} -> ${ORIG_BUILDER_TAG}\"\n docker tag \"${BUILDER_IMAGE}\" \"${ORIG_BUILDER_TAG}\" || true\nfi\nif docker image inspect \"${RUNTIME_IMAGE}\" >/dev/null 2>&1; then\n echo \"Tagging ${RUNTIME_IMAGE} -> ${ORIG_RUNTIME_TAG}\"\n docker tag \"${RUNTIME_IMAGE}\" \"${ORIG_RUNTIME_TAG}\" || true\nfi\n\necho \"Done pulling/tagging images.\"\necho \"You can now run downstream build steps that expect these images to exist locally.\"\n" }, { "path": "scripts/download_checkpoints.sh", "content": "mkdir -p ./checkpoints/\ncd ./checkpoints/\n\n# aesthetic\nwget https://github.com/christophschuhmann/improved-aesthetic-predictor/raw/main/sac+logos+ava1-l14-linearMSE.pth -O aesthetic.pth\n\n# megasam\nwget https://github.com/mega-sam/mega-sam/blob/main/checkpoints/megasam_final.pth -O megasam_final.pth\n\n# raft\ngdown -c https://drive.google.com/uc?id=1MqDajR89k-xLV0HIrmJ0k-n8ZpG6_suM -O raft-things.pth\n\n# depth anything\nhuggingface-cli download --resume-download depth-anything/Depth-Anything-V2-Large --local-dir Depth-Anything\n\n# unidepth\nhuggingface-cli download --resume-download lpiccinelli/unidepth-v2-vitl14 --local-dir UniDepth\n\n# sam\nhuggingface-cli download --resume-download facebook/sam2.1-hiera-large --local-dir SAM2\n" }, { "path": "scripts/scoring.sh", "content": "#!/bin/bash\nVIDEO_DIR=[Replace with the path to your video files]\nOUTPUT_DIR=[Replace with the path to your output directory]\nmkdir -p ${OUTPUT_DIR}\n\n# Choose whether to cut the clips precisely based on the timestamps or to cut them fast based on keyframes.\n# The precise cutting will be slower but more accurate, while the fast cutting will be faster but may not be as accurate.\nFAST_CUT=False\n\nGPU_NUM=8\nCUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nNUM_WORKERS=$((GPU_NUM * 2))\n\nROOT_CLIPS=${OUTPUT_DIR}/clip\nROOT_META=${OUTPUT_DIR}/meta\nROOT_FIG=${OUTPUT_DIR}/fig\nROOT_TEMP=${OUTPUT_DIR}/temp\n\nfor dir in ${ROOT_CLIPS} ${ROOT_META} ${ROOT_FIG} ${ROOT_TEMP}; do\n if [ ! -d ${dir} ]; then\n mkdir -p ${dir}\n fi\ndone\n\nmeasure_time() {\n local step_number=$1\n shift\n local green=\"\\e[32m\"\n local red=\"\\e[31m\"\n local no_color=\"\\e[0m\"\n local yellow=\"\\e[33m\"\n \n start_time=$(date +%s)\n echo -e \"${green}Step ${step_number} started at: $(date)${no_color}\"\n\n \"$@\"\n\n end_time=$(date +%s)\n echo -e \"${red}Step ${step_number} finished at: $(date)${no_color}\"\n echo -e \"${yellow}Duration: $((end_time - start_time)) seconds${no_color}\"\n echo \"---------------------------------------\"\n}\n\n# 1.1 Create a meta file from a video folder. This should output ${ROOT_META}/meta.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 1.1 python utils/convert.py \\\n --video_dir ${VIDEO_DIR} \\\n --output ${ROOT_META}/meta.csv\n\n# 1.2 Get video information and remove broken videos. This should output ${ROOT_META}/meta_info_fmin${fmin_1}.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 1.2 python utils/get_info.py \\\n --csv_path ${ROOT_META}/meta.csv \\\n --csv_save_path ${ROOT_META}/meta_info.csv \\\n --backend \"opencv\" \\\n --num_workers 16\n\n# 2.1 Detect scenes. This should output ${ROOT_META}/meta_info_fmin${fmin_1}_timestamp.csv\n# Also, you can set the params like \"--start-remove-sec 0.5 --end-remove-sec 0.5\"\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 2.1 python utils/scene_detect.py \\\n --csv_path ${ROOT_META}/meta_info.csv \\\n --backend \"opencv\" \\\n --num_workers 64 \\\n --frame_skip 2\\\n --start_remove_sec 0.3 \\\n --end_remove_sec 0.3 \\\n --min_seconds 3 \\\n --max_seconds 15\n\n# 2.2 Get clips. This should output ${ROOT_META}/clips_info.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 2.2 python utils/get_clip.py \\\n --csv_path ${ROOT_META}/meta_info_timestamp.csv \\\n --csv_save_dir ${ROOT_META} \\\n --num_workers $((GPU_NUM * 4))\n\n# 2.3 Extract frames for scoring.\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 2.3 python utils/extract_frames.py \\\n --csv_path ${ROOT_META}/clips_info.csv \\\n --output_dir ${ROOT_FIG} \\\n --num_workers 64 \\\n --target_size \"640*360\" \\\n --backend \"opencv\"\n\n# 3.1 Predict aesthetic scores. This should output ${ROOT_META}/clips_info_aes.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 3.1 torchrun --nproc_per_node ${GPU_NUM} scoring/aesthetic/inference.py \\\n --csv_path ${ROOT_META}/clips_info.csv \\\n --bs 16 \\\n --num_workers ${NUM_WORKERS} \\\n --fig_load_dir ${ROOT_FIG}\n\n# 3.2 Predict luminance scores. This should output ${ROOT_META}/clips_info_lum.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 3.2 torchrun --nproc_per_node ${GPU_NUM} scoring/luminance/inference.py \\\n --csv_path ${ROOT_META}/clips_info.csv \\\n --bs 16 \\\n --num_workers ${NUM_WORKERS} \\\n --fig_load_dir ${ROOT_FIG}\n\n# 3.3 get motion score. This should output ${ROOT_META}/clips_info_motion.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 3.3 python scoring/motion/inference.py \\\n --csv_path ${ROOT_META}/clips_info.csv \\\n --temp_save_dir ${ROOT_TEMP} \\\n --num_workers $((GPU_NUM * 4)) \\\n --gpu_num ${GPU_NUM}\n \n# 3.4 get text by OCR using PaddleOCR, this should output ${ROOT_META}/clips_info_ocr.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 3.4 python scoring/ocr/inference.py \\\n --csv_path ${ROOT_META}/clips_info.csv \\\n --fig_load_dir ${ROOT_FIG} \\\n --num_workers $((GPU_NUM * 4)) \\\n --gpu_num ${GPU_NUM}\n\n# 4 merge all the scores. This should output ${ROOT_META}/clips_with_score.csv\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 4 python utils/merge_tables.py \\\n --csv_dir ${ROOT_META} \\\n --output ${ROOT_META}/clips_scores.csv\n\n# 5 Filter the clips.\nCUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 5 python utils/filter.py \\\n --csv_path ${ROOT_META}/clips_scores.csv \\\n --csv_save_path ${ROOT_META}/filtered_clips.csv \\\n --aes_min 4 \\\n --lum_min 20 \\\n --lum_max 140 \\\n --motion_min 2 \\\n --motion_max 14 \\\n --ocr_max 0.3\n\n# 6 Cut the clips.\nif [ \"$FAST_CUT\" = False ]; then\n echo \"Using precise cutting based on timestamps.\"\n CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 6 python utils/cut.py \\\n --csv_path ${ROOT_META}/filtered_clips.csv \\\n --csv_save_path ${OUTPUT_DIR}/results.csv \\\n --video_save_dir ${ROOT_CLIPS} \\\n --num_workers $((GPU_NUM * 4)) \\\n --gpu_num $GPU_NUM \\\n # --keep_audio \nelse\n echo \"Using fast cutting based on keyframes.\"\n CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} measure_time 6 python utils/cut_fast.py \\\n --csv_path ${ROOT_META}/filtered_clips.csv \\\n --csv_save_path ${OUTPUT_DIR}/results.csv \\\n --video_save_dir ${ROOT_CLIPS} \\\n --num_workers $((GPU_NUM * 4)) \\\n # --keep_audio\nfi" }, { "path": "utils/README.md", "content": "# Utils\n\n- [`convert.py`](convert.py): convert all the paths of videos in a directory to a specific format, like csv.\n- [`cut.py`](cut.py): cut videos into clips.\n- [`download_SpatialVID.py`](download_SpatialVID.py): download the SpatialVID dataset.\n- [`download_YouTube.py`](download_YouTube.py): download videos from YouTube.\n- [`evaluate.py`](evaluate.py): evaluate the quality of video reconstructions.\n- [`expand_npz.py`](expand_npz.py): get dynamic masks compressed in a npz file.\n- [`extract_frames.py`](extract_frames.py): extract frames from videos.\n- [`filter.py`](filter.py): filter video clips based on score.\n- [`get_clip.py`](get_clip.py): get the clips separated from the video.\n- [`get_info.py`](get_info.py): get video information, such as duration and resolution.\n- [`get_instructions.py`](get_instructions.py): get motion instructions from camera poses.\n- [`get_instructions_enhanced.py`](get_instructions_enhanced.py): an enhanced version to get more detailed and accurate motion instructions from camera poses.\n- [`merge_tables.py`](merge_tables.py): merge multiple csv tables into one.\n- [`normalize_intrinsics.py`](normalize_intrinsics.py): normalize camera intrinsics.\n- [`pack_clip_assets.py`](pack_clip_assets.py): pack all the output files into an npz file for visualization.\n- [`quat_to_mat.py`](quat_to_mat.py): convert camera parameters to camera-to-world or world-to-camera matrices.\n- [`read_video.py`](read_video.py): read videos using opencv or av.\n- [`scene_detect.py`](scene_detect.py): separate videos into clips.\n" }, { "path": "utils/__init__.py", "content": "" }, { "path": "utils/convert.py", "content": "\"\"\"\nVideo file conversion utility for the SpatialVID project.\n\nThis module provides functionality to scan directories for video files,\nprocess them, and generate CSV metadata files containing video information.\n\"\"\"\n\nimport argparse\nimport os\nimport time\nimport pandas as pd\n\n\n# Supported video file extensions\nVID_EXTENSIONS = (\".mp4\", \".avi\", \".mov\", \".mkv\", \".m2ts\", \".webm\")\n\n\ndef scan_recursively(root):\n \"\"\"\n Recursively scan a directory tree and yield all entries.\n \"\"\"\n num = 0\n for entry in os.scandir(root):\n if entry.is_file():\n yield entry\n elif entry.is_dir():\n num += 1\n if num % 100 == 0:\n print(f\"Scanned {num} directories.\")\n yield from scan_recursively(entry.path)\n\n\ndef get_filelist(file_path, exts=None):\n \"\"\"\n Get a list of files from a directory tree, optionally filtered by extensions.\n \"\"\"\n filelist = []\n time_start = time.time()\n\n # Use recursive scanning to find all files\n obj = scan_recursively(file_path)\n for entry in obj:\n if entry.is_file():\n ext = os.path.splitext(entry.name)[-1].lower()\n if exts is None or ext in exts:\n filelist.append(entry.path)\n\n time_end = time.time()\n print(f\"Scanned {len(filelist)} files in {time_end - time_start:.2f} seconds.\")\n return filelist\n\n\ndef split_by_capital(name):\n \"\"\"\n Split a camelCase or PascalCase string by capital letters.\n \"\"\"\n new_name = \"\"\n for i in range(len(name)):\n if name[i].isupper() and i != 0:\n new_name += \" \"\n new_name += name[i]\n return new_name\n\n\ndef process_general_videos(root, output):\n \"\"\"\n Process video files in a directory and generate a CSV metadata file.\n \"\"\"\n # Expand user path (e.g., ~ to home directory)\n root = os.path.expanduser(root)\n if not os.path.exists(root):\n return\n\n # Get list of video files with supported extensions\n path_list = get_filelist(root, VID_EXTENSIONS)\n # Note: In some cases (like realestate dataset), you might want to use:\n # path_list = get_filelist(root) # without extension filtering\n\n path_list = list(set(path_list)) # Remove duplicate entries\n\n # Extract filename without extension as ID\n fname_list = [os.path.splitext(os.path.basename(x))[0] for x in path_list]\n # Get relative paths from root directory\n relpath_list = [os.path.relpath(x, root) for x in path_list]\n\n # Create DataFrame with video metadata\n df = pd.DataFrame(dict(video_path=path_list, id=fname_list, relpath=relpath_list))\n\n # Ensure output directory exists\n os.makedirs(os.path.dirname(output), exist_ok=True)\n df.to_csv(output, index=False)\n print(f\"Saved {len(df)} samples to {output}.\")\n\n\nif __name__ == \"__main__\":\n # Set up command line argument parser\n parser = argparse.ArgumentParser(\n description=\"Convert video directory structure to CSV metadata file\"\n )\n parser.add_argument(\"--video_dir\", type=str, help=\"Root directory containing video files\")\n parser.add_argument(\"--split\", type=str, default=\"train\", help=\"Dataset split name\")\n parser.add_argument(\"--info\", type=str, default=None, help=\"Additional info file\")\n parser.add_argument(\n \"--output\", type=str, default=None, required=True, help=\"Output CSV file path\"\n )\n args = parser.parse_args()\n\n # Process videos and generate metadata CSV\n process_general_videos(args.video_dir, args.output)\n" }, { "path": "utils/cut.py", "content": "\"\"\"\nPrecise frame-level video cutting tool\nStrategy: Two-phase seek + forced keyframe alignment output\n\"\"\"\n\nimport argparse\nimport os\nimport concurrent.futures\nfrom functools import partial\nimport pandas as pd\nimport subprocess\nfrom scenedetect import FrameTimecode\nfrom tqdm import tqdm\n\n\nFFMPEG_PATH = \"/usr/local/bin/ffmpeg\"\n\n\ndef get_ffmpeg_acceleration():\n try:\n output = subprocess.check_output(\n [FFMPEG_PATH, \"-encoders\"], stderr=subprocess.DEVNULL\n ).decode(\"utf-8\")\n if \"hevc_nvenc\" in output:\n return \"nvidia\"\n return \"cpu\"\n except Exception as e:\n print(f\"FFmpeg acceleration detection failed: {e}\")\n return \"cpu\"\n\n\nACCELERATION_TYPE = get_ffmpeg_acceleration()\nprint(f\"FFmpeg acceleration type: {ACCELERATION_TYPE}\")\n\n\n# ════════════════════════════════════════════════════════════\n# Core Utility Functions\n# ════════════════════════════════════════════════════════════\n\ndef seconds_to_timecode(seconds: float) -> str:\n \"\"\"\n Convert seconds to FFmpeg precise timecode string.\n Keep enough decimal places to ensure frame accuracy.\n\n Example: 1.033333 -> \"0:00:01.033333\"\n \"\"\"\n hours = int(seconds // 3600)\n minutes = int((seconds % 3600) // 60)\n secs = seconds % 60\n # Keep 6 decimal places (microsecond-level precision)\n return f\"{hours}:{minutes:02d}:{secs:09.6f}\"\n\n\ndef build_precise_cut_cmd(\n video_path: str,\n start_sec: float,\n end_sec: float,\n save_path: str,\n args,\n process_id: int,\n shorter_size: int | None,\n) -> list[str]:\n \"\"\"\n Build frame-precise FFmpeg cut command.\n\n Strategy: Two-phase seek\n ┌──────────────────────────────────────────────────────────┐\n │ -ss (pre, coarse seek) │\n │ -> Jump to nearest keyframe before start_sec │\n │ -> Avoid decoding from file start (speed optimize) │\n │ │\n │ -i input │\n │ │\n │ -ss (post, fine seek) │\n │ -> Decode from coarse point to exact start_sec │\n │ -> value = start_sec - coarse_seek (always positive) │\n │ │\n │ -t duration │\n │ -> Exact duration │\n │ │\n │ Force re-encode (cannot use -c copy, otherwise │\n │ start frame won't be precise) │\n └──────────────────────────────────────────────────────────┘\n \"\"\"\n duration = end_sec - start_sec\n\n if duration <= 0:\n raise ValueError(f\"Invalid duration {duration:.4f}s (start={start_sec}, end={end_sec})\")\n\n # ==== Phase 1: Coarse seek (pre seek) ====\n # Safety margin: ensure coarse point is before start_sec keyframe\n # Too little -> may land after start_sec (seek ineffective)\n # Too much -> decode more frames (slightly slower)\n # Experience: max(GOP_size, 5s) covers most videos\n GOP_SAFETY_MARGIN = 5.0\n coarse_seek = max(0.0, start_sec - GOP_SAFETY_MARGIN)\n\n # Offset for post precise seek = target time - coarse time\n fine_seek = start_sec - coarse_seek\n\n cmd = [FFMPEG_PATH, \"-nostdin\", \"-y\"]\n\n # ==== GPU hardware acceleration (decode phase) ====\n if ACCELERATION_TYPE == \"nvidia\":\n cmd += [\n \"-hwaccel\", \"cuda\",\n \"-hwaccel_output_format\", \"cuda\",\n \"-hwaccel_device\", str(process_id % args.gpu_num),\n ]\n\n # ==== Phase 1: Coarse seek (pre, fast jump to GOP boundary) ====\n cmd += [\"-ss\", seconds_to_timecode(coarse_seek)]\n\n # ==== Input file ====\n cmd += [\"-i\", video_path]\n\n # ==== Phase 2: Precise seek (post, decode from GOP boundary to exact frame) ====\n # Only need post seek when fine_seek > 0\n # When coarse_seek == 0, fine_seek == start_sec, still correct\n if fine_seek > 0.001: # Ignore errors less than 1ms\n cmd += [\"-ss\", seconds_to_timecode(fine_seek)]\n\n # ==== Exact duration ====\n cmd += [\"-t\", seconds_to_timecode(duration)]\n\n # ==== Video filters (scale + fps) ====\n filters = _build_video_filters(shorter_size, args, ACCELERATION_TYPE)\n if filters:\n cmd += [\"-vf\", \",\".join(filters)]\n\n # ==== Encoder (must re-encode to ensure frame precision) ====\n cmd += _build_encoder_args(ACCELERATION_TYPE)\n\n # ==== Frame rate ====\n if args.target_fps is not None:\n cmd += [\"-r\", str(args.target_fps)]\n\n # ==== Audio ====\n if args.keep_audio:\n cmd += [\"-map\", \"0:v\", \"-map\", \"0:a?\", \"-c:a\", \"aac\", \"-b:a\", \"128k\"]\n else:\n cmd += [\"-map\", \"0:v\", \"-an\"]\n\n # ==== Output: force keyframe at first frame for easy concatenation/playback ====\n cmd += [\n \"-force_key_frames\", \"expr:gte(t,0)\", # Force keyframe at second 0\n save_path,\n ]\n\n return cmd\n\n\ndef _build_video_filters(shorter_size, args, accel_type) -> list[str]:\n \"\"\"Build video filter list\"\"\"\n filters = []\n\n if shorter_size is not None:\n if accel_type == \"nvidia\":\n # CUDA scale filter\n scale = (\n f\"scale_cuda=\"\n f\"'if(gt(iw,ih),-2,{shorter_size})':\"\n f\"'if(gt(iw,ih),{shorter_size},-2)'\"\n )\n else:\n # Software scale: lanczos best quality, bicubic next\n scale = (\n f\"scale=\"\n f\"'if(gt(iw,ih),-2,{shorter_size})':\"\n f\"'if(gt(iw,ih),{shorter_size},-2)'\"\n f\":flags=lanczos\"\n )\n filters.append(scale)\n\n if args.target_fps is not None:\n # fps filter more accurate than -r parameter (-r sometimes drops frames)\n filters.append(f\"fps={args.target_fps}\")\n\n return filters\n\n\ndef _build_encoder_args(accel_type) -> list[str]:\n \"\"\"Build encoder arguments\"\"\"\n if accel_type == \"nvidia\":\n return [\n \"-c:v\", \"hevc_nvenc\",\n \"-preset\", \"p4\", # p4=quality/speed balance, p7=slowest best\n \"-rc\", \"vbr\",\n \"-cq\", \"24\", # Quality factor, smaller is better (like CRF)\n \"-b:v\", \"0\", # No bitrate limit in VBR mode\n ]\n else:\n return [\n \"-c:v\", \"libx264\",\n \"-preset\", \"fast\", # fast is best speed/quality for precise cutting\n \"-crf\", \"18\", # High quality (0=lossless, 23=default, 18=visually lossless)\n \"-pix_fmt\", \"yuv420p\", # Most compatible pixel format\n ]\n\n\n# ════════════════════════════════════════════════════════════\n# Single Row Processing (maintains compatibility with original interface)\n# ════════════════════════════════════════════════════════════\n\ndef process_single_row(row, args, process_id):\n \"\"\"\n Precise frame-level cutting of a single segment.\n\n Returns:\n (row_values_list, valid, error_message)\n \"\"\"\n video_path = row[\"video_path\"]\n save_dir = args.video_save_dir\n\n #\n # ==== Scale size calculation ====\n shorter_size = args.shorter_size\n if (shorter_size is not None) and (\"height\" in row) and (\"width\" in row):\n min_size = min(row[\"height\"], row[\"width\"])\n if min_size <= shorter_size:\n shorter_size = None # Already small enough, skip scaling (no upsample)\n\n # ==== Timestamp parsing ====\n try:\n seg_start = FrameTimecode(timecode=row[\"timestamp_start\"], fps=row[\"fps\"])\n seg_end = FrameTimecode(timecode=row[\"timestamp_end\"], fps=row[\"fps\"])\n except Exception as e:\n error_msg = f\"Invalid timestamp for id={row.get('id', '?')}: {e}\"\n print(error_msg)\n return row.values.tolist(), False, error_msg\n\n start_sec = seg_start.get_seconds()\n end_sec = seg_end.get_seconds()\n duration = end_sec - start_sec\n\n if duration <= 0:\n error_msg = (\n f\"Invalid duration {duration:.4f}s for id={row.get('id','?')} \"\n f\"(start={start_sec:.4f}, end={end_sec:.4f})\"\n )\n print(error_msg)\n return row.values.tolist(), False, error_msg\n\n clip_id = row[\"id\"]\n save_path = os.path.join(save_dir, f\"{clip_id}.mp4\")\n\n # ==== Skip if already exists ====\n if os.path.exists(save_path) and os.path.getsize(save_path) > 0:\n row = row.copy()\n row[\"video_path\"] = save_path\n return row.values.tolist(), True, \"\"\n\n # ==== Source file check ====\n if not os.path.exists(video_path):\n error_msg = f\"Source video not found: {video_path} (id={clip_id})\"\n print(error_msg)\n return row.values.tolist(), False, error_msg\n\n # ==== Build precise cut command ====\n try:\n cmd = build_precise_cut_cmd(\n video_path = video_path,\n start_sec = start_sec,\n end_sec = end_sec,\n save_path = save_path,\n args = args,\n process_id = process_id,\n shorter_size = shorter_size,\n )\n except ValueError as e:\n error_msg = f\"Command build failed for id={clip_id}: {e}\"\n print(error_msg)\n return row.values.tolist(), False, error_msg\n\n # ==== Execute FFmpeg ====\n try:\n subprocess.run(cmd, check=True, stderr=subprocess.PIPE)\n except subprocess.CalledProcessError as e:\n stderr_text = e.stderr.decode(\"utf-8\", errors=\"replace\") if e.stderr else str(e)\n error_msg = f\"FFmpeg failed for id={clip_id}:\\n{stderr_text}\"\n print(error_msg)\n _cleanup(save_path)\n return row.values.tolist(), False, error_msg\n except Exception as e:\n error_msg = f\"Unexpected error for id={clip_id}: {e}\"\n print(error_msg)\n _cleanup(save_path)\n return row.values.tolist(), False, error_msg\n\n # ==== Basic integrity check ====\n if not os.path.exists(save_path) or os.path.getsize(save_path) == 0:\n _cleanup(save_path)\n error_msg = f\"FFmpeg produced empty/missing output for id={clip_id}\"\n print(error_msg)\n return row.values.tolist(), False, error_msg\n\n row = row.copy()\n row[\"video_path\"] = save_path\n return row.values.tolist(), True, \"\"\n\n\ndef _cleanup(path: str):\n \"\"\"Safely delete file\"\"\"\n try:\n if os.path.exists(path):\n os.remove(path)\n except OSError:\n pass\n\n\n# ════════════════════════════════════════════════════════════\n# Argument Parsing\n# ════════════════════════════════════════════════════════════\n\ndef parse_args():\n parser = argparse.ArgumentParser(\n description=\"Precise frame-level video cutting tool\",\n formatter_class=argparse.ArgumentDefaultsHelpFormatter,\n )\n\n # ==== Input/Output ====\n parser.add_argument(\"--csv_path\", type=str, required=True,\n help=\"Input CSV file path\")\n parser.add_argument(\"--csv_save_path\", type=str, required=True,\n help=\"Output CSV file path (success records)\")\n parser.add_argument(\"--video_save_dir\", type=str, required=True,\n help=\"Directory to save cut segments\")\n\n # ==== Video parameters ====\n parser.add_argument(\"--target_fps\", type=int, default=None,\n help=\"Target frame rate (None=keep source frame rate)\")\n parser.add_argument(\"--shorter_size\", type=int, default=None,\n help=\"Short edge target size (maintain aspect ratio, no upsample)\")\n parser.add_argument(\"--keep_audio\", action=\"store_true\",\n help=\"Keep audio track (default: discard)\")\n\n # ==== Parallel control ====\n parser.add_argument(\"--num_workers\", type=int, default=None,\n help=\"Number of parallel workers (None=auto=CPU cores)\")\n parser.add_argument(\"--disable_parallel\", action=\"store_true\",\n help=\"Disable parallel processing (for debugging)\")\n parser.add_argument(\"--gpu_num\", type=int, default=1,\n help=\"Number of available GPUs\")\n\n # ==== Result handling ====\n parser.add_argument(\"--drop_invalid_timestamps\", action=\"store_true\",\n help=\"Filter invalid timestamps and save corrected CSV\")\n\n return parser.parse_args()\n\n\n# ════════════════════════════════════════════════════════════\n# Parallel Worker\n# ════════════════════════════════════════════════════════════\n\ndef _worker_fn(task: tuple, args, process_id: int) -> tuple:\n \"\"\"\n Top-level worker function for ProcessPoolExecutor (must be serializable).\n\n Args:\n task: (index, row_dict) <- Use dict instead of Series to avoid serialization issues\n\n Returns:\n (index, row_values, valid, error_msg)\n \"\"\"\n index, row_dict = task\n\n # Restore dict to pandas Series (process_single_row depends on Series interface)\n row = pd.Series(row_dict)\n return (index,) + tuple(process_single_row(row, args, process_id)[0:3])\n # Note: process_single_row returns (row_values, valid, error_msg)\n # Packed here as (index, row_values, valid, error_msg)\n\n\n# ════════════════════════════════════════════════════════════\n# Result Saving\n# ════════════════════════════════════════════════════════════\n\ndef save_results(all_results: list, csv: pd.DataFrame, args):\n \"\"\"\n Save processing results to success/failure CSVs separately.\n\n Success CSV: Remove timestamp helper columns, update video_path to cut path\n Failure CSV: Keep all original columns, add error column\n \"\"\"\n columns = csv.columns.tolist()\n\n success_rows, failed_rows, failed_errors = [], [], []\n\n for index, row_values, valid, error_msg in all_results:\n if valid:\n success_rows.append(row_values)\n else:\n failed_rows.append(row_values)\n failed_errors.append(error_msg)\n\n # ==== Save success records ====\n if success_rows:\n success_df = pd.DataFrame(success_rows, columns=columns)\n\n # Remove cutting process helper columns (not needed by downstream)\n drop_cols = [\n c for c in [\"timestamp_start\", \"timestamp_end\", \"frame_start\", \"frame_end\"]\n if c in success_df.columns\n ]\n if drop_cols:\n success_df = success_df.drop(columns=drop_cols)\n\n success_df.to_csv(args.csv_save_path, index=False)\n print(f\"\\n[OK] Success: {len(success_df)} records -> {args.csv_save_path}\")\n else:\n print(\"\\n[X] No success records\")\n\n # ==== Save failure records ====\n if failed_rows:\n base, ext = os.path.splitext(args.csv_save_path)\n failed_csv_path = f\"{base}_failed{ext}\"\n\n failed_df = pd.DataFrame(failed_rows, columns=columns)\n failed_df[\"error\"] = failed_errors\n failed_df.to_csv(failed_csv_path, index=False)\n print(f\"[X] Failed: {len(failed_df)} records -> {failed_csv_path}\")\n\n # ==== Save corrected timestamps (optional) ====\n if args.drop_invalid_timestamps and failed_rows:\n valid_indices = [r[0] for r in all_results if r[2]]\n filtered_csv = csv.iloc[valid_indices]\n assert args.csv_path.endswith(\"timestamp.csv\"), \\\n \"--drop_invalid_timestamps only supports *timestamp.csv files\"\n corrected_path = args.csv_path.replace(\"timestamp.csv\", \"correct_timestamp.csv\")\n filtered_csv.to_csv(corrected_path, index=False)\n print(f\"[OK] Corrected timestamps -> {corrected_path}\")\n\n\n# ════════════════════════════════════════════════════════════\n# Main Function\n# ════════════════════════════════════════════════════════════\n\ndef main():\n args = parse_args()\n\n # ==== Pre-check ====\n if not os.path.exists(args.csv_path):\n print(f\"[ERROR] CSV file does not exist: {args.csv_path}\")\n return\n\n os.makedirs(args.video_save_dir, exist_ok=True)\n\n csv = pd.read_csv(args.csv_path)\n total = len(csv)\n print(f\"Total {total} records to process\")\n\n all_results = []\n\n # ==== Serial mode ====\n if args.disable_parallel:\n for index, row in tqdm(csv.iterrows(), total=total, desc=\"Cutting progress\"):\n row_values, valid, error_msg = process_single_row(row, args, process_id=0)\n all_results.append((index, row_values, valid, error_msg))\n\n # ==== Parallel mode ====\n else:\n num_workers = args.num_workers or (os.cpu_count() or 1)\n num_workers = min(num_workers, total) # worker count not exceeding task count\n\n # Convert row to dict to avoid pandas Series serialization issues\n tasks = [\n (index, row.to_dict())\n for index, row in csv.iterrows()\n ]\n\n with concurrent.futures.ProcessPoolExecutor(\n max_workers=num_workers\n ) as executor:\n # Use enumerate to round-robin process_id (GPU rotation)\n futures = {\n executor.submit(\n _worker_fn,\n task,\n args,\n task_idx % max(args.gpu_num, 1), # GPU rotation\n ): task_idx\n for task_idx, task in enumerate(tasks)\n }\n\n with tqdm(total=total, desc=\"Cutting progress\") as pbar:\n for future in concurrent.futures.as_completed(futures):\n try:\n result = future.result() # (index, row_values, valid, error_msg)\n all_results.append(result)\n except Exception as e:\n task_idx = futures[future]\n index, _ = tasks[task_idx]\n row_values = csv.iloc[index].values.tolist()\n all_results.append((index, row_values, False, str(e)))\n print(f\"\\n[ERROR] Worker exception (task_idx={task_idx}): {e}\")\n finally:\n pbar.update(1)\n\n # ==== Sort by original order ====\n all_results.sort(key=lambda x: x[0])\n\n # ==== Statistics summary ====\n success_count = sum(1 for r in all_results if r[2])\n failed_count = total - success_count\n print(f\"\\n{'='*50}\")\n print(f\"Processing complete: Total={total}, Success={success_count}, Failed={failed_count}\")\n print(f\"{'='*50}\")\n\n # ==== Save results ====\n save_results(all_results, csv, args)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/cut_fast.py", "content": "\"\"\"\nHigh-speed video cutting utility using FFmpeg stream copy.\n\nFeatures:\n- No re-encoding: uses `-c copy`\n- Optional audio: use --keep_audio to retain audio tracks\n- Group tasks by source video_path for better efficiency\n- Parallel processing by video group\n- Per-clip progress bar\n- Save successful and failed CSVs\n\nNotes:\n- This method is very fast, but not always frame-accurate.\n- Clip boundaries may align to nearby keyframes depending on source encoding.\n\"\"\"\n\nimport argparse\nimport os\nimport queue\nimport subprocess\nimport concurrent.futures\nfrom multiprocessing import Manager\n\nimport pandas as pd\nfrom scenedetect import FrameTimecode\nfrom tqdm import tqdm\n\n\nFFMPEG_PATH = \"/usr/local/bin/ffmpeg\"\n\n\ndef process_single_row(row, save_dir, keep_audio=False):\n \"\"\"\n Cut one clip from source video using ffmpeg stream copy.\n\n Args:\n row: DataFrame row with clip metadata\n save_dir: directory to save output clips\n keep_audio: if True, copy audio streams; if False, drop audio\n\n Returns:\n (row_values_list, valid, error_message)\n \"\"\"\n video_path = row[\"video_path\"]\n sample_id = row[\"id\"]\n save_path = os.path.join(save_dir, f\"{sample_id}.mp4\")\n\n # Already exists -> treat as success\n if os.path.exists(save_path) and os.path.getsize(save_path) > 0:\n row = row.copy()\n row[\"video_path\"] = save_path\n return row.values.tolist(), True, \"\"\n\n if not os.path.exists(video_path):\n error_msg = f\"Source video not found: {video_path} (id={sample_id})\"\n return row.values.tolist(), False, error_msg\n\n # Parse timestamps\n try:\n fps = row[\"fps\"]\n seg_start = FrameTimecode(timecode=row[\"timestamp_start\"], fps=fps)\n seg_end = FrameTimecode(timecode=row[\"timestamp_end\"], fps=fps)\n\n start_sec = float(seg_start.get_seconds())\n end_sec = float(seg_end.get_seconds())\n duration = end_sec - start_sec\n\n if duration <= 0:\n error_msg = f\"Non-positive duration for id={sample_id}: {duration}\"\n return row.values.tolist(), False, error_msg\n\n except Exception as e:\n error_msg = f\"Invalid timestamp for id={sample_id}: {e}\"\n return row.values.tolist(), False, error_msg\n\n try:\n # Build stream mapping and audio arguments based on keep_audio flag.\n # '0:a?' uses '?' so FFmpeg silently skips if no audio track exists.\n if keep_audio:\n map_args = [\"-map\", \"0:v:0\", \"-map\", \"0:a?\"]\n audio_args = [\"-c:a\", \"copy\"]\n else:\n map_args = [\"-map\", \"0:v:0\"]\n audio_args = [\"-an\"]\n\n # Fast seek + stream copy; explicitly specify video codec to avoid ambiguity.\n cmd = [\n FFMPEG_PATH,\n \"-nostdin\",\n \"-y\",\n \"-ss\",\n str(start_sec),\n \"-t\",\n str(duration),\n \"-i\",\n video_path,\n *map_args,\n *audio_args,\n \"-c:v\",\n \"copy\",\n \"-avoid_negative_ts\",\n \"make_zero\",\n save_path,\n ]\n\n subprocess.run(\n cmd,\n check=True,\n stdout=subprocess.DEVNULL,\n stderr=subprocess.PIPE,\n )\n\n # Verify output exists and non-empty\n if not os.path.exists(save_path) or os.path.getsize(save_path) == 0:\n if os.path.exists(save_path):\n os.remove(save_path)\n error_msg = f\"FFmpeg produced empty/missing output for id={sample_id}\"\n return row.values.tolist(), False, error_msg\n\n row = row.copy()\n row[\"video_path\"] = save_path\n return row.values.tolist(), True, \"\"\n\n except subprocess.CalledProcessError as e:\n stderr_text = e.stderr.decode(\"utf-8\", errors=\"ignore\") if e.stderr else str(e)\n error_msg = f\"FFmpeg failed for id={sample_id}: {stderr_text}\"\n if os.path.exists(save_path):\n os.remove(save_path)\n return row.values.tolist(), False, error_msg\n\n except Exception as e:\n error_msg = f\"Unexpected error for id={sample_id}: {e}\"\n if os.path.exists(save_path):\n os.remove(save_path)\n return row.values.tolist(), False, error_msg\n\n\ndef process_video_group(group_df, save_dir, keep_audio=False):\n \"\"\"\n Process all clips from the same source video.\n\n Args:\n group_df: DataFrame containing rows from one source video_path\n save_dir: output clip directory\n keep_audio: passed through to process_single_row\n\n Returns:\n list of tuples: (index, row_values, valid, error_msg)\n \"\"\"\n results = []\n\n # Sort by start timestamp to make access pattern a bit more sequential\n if \"timestamp_start\" in group_df.columns:\n group_df = group_df.sort_values(by=\"timestamp_start\")\n\n for index, row in group_df.iterrows():\n row_values, valid, error_msg = process_single_row(\n row, save_dir, keep_audio=keep_audio\n )\n results.append((index, row_values, valid, error_msg))\n\n return results\n\n\ndef worker(task_queue, results_queue, video_save_dir, keep_audio=False):\n \"\"\"\n Worker that processes one video group at a time.\n \"\"\"\n while True:\n try:\n video_path, group_df = task_queue.get(timeout=1)\n except queue.Empty:\n break\n\n try:\n group_results = process_video_group(\n group_df, video_save_dir, keep_audio=keep_audio\n )\n for item in group_results:\n results_queue.put(item)\n finally:\n task_queue.task_done()\n\n\ndef parse_args():\n parser = argparse.ArgumentParser(\n description=\"Fast video cutting utility using FFmpeg stream copy\"\n )\n parser.add_argument(\"--csv_path\", type=str, required=True, help=\"Input CSV path\")\n parser.add_argument(\n \"--csv_save_path\", type=str, required=True, help=\"Output CSV path\"\n )\n parser.add_argument(\n \"--video_save_dir\", type=str, required=True, help=\"Directory to save clips\"\n )\n parser.add_argument(\n \"--num_workers\",\n type=int,\n default=None,\n help=\"Number of parallel workers (defaults to CPU count)\",\n )\n parser.add_argument(\n \"--disable_parallel\",\n action=\"store_true\",\n help=\"Disable parallel processing\",\n )\n parser.add_argument(\n \"--drop_invalid_timestamps\",\n action=\"store_true\",\n help=\"Drop invalid timestamp rows and save corrected CSV\",\n )\n parser.add_argument(\n \"--keep_audio\",\n action=\"store_true\",\n help=\"Retain audio tracks in output clips (dropped by default)\",\n )\n return parser.parse_args()\n\n\ndef main():\n args = parse_args()\n\n if not os.path.exists(args.csv_path):\n print(f\"csv file '{args.csv_path}' not found. Exit.\")\n return\n\n os.makedirs(args.video_save_dir, exist_ok=True)\n\n csv = pd.read_csv(args.csv_path)\n if len(csv) == 0:\n print(\"Input CSV is empty. Exit.\")\n return\n\n required_cols = [\"id\", \"video_path\", \"timestamp_start\", \"timestamp_end\", \"fps\"]\n missing_cols = [c for c in required_cols if c not in csv.columns]\n if missing_cols:\n raise ValueError(f\"Missing required columns: {missing_cols}\")\n\n results = []\n\n # Group by source video\n grouped_items = list(csv.groupby(\"video_path\", sort=False))\n total_tasks = len(csv)\n\n if args.disable_parallel:\n success_cnt = 0\n fail_cnt = 0\n\n with tqdm(total=total_tasks, desc=\"Processing clips\", dynamic_ncols=True) as pbar:\n for video_path, group_df in grouped_items:\n group_results = process_video_group(\n group_df, args.video_save_dir, keep_audio=args.keep_audio\n )\n for item in group_results:\n results.append(item)\n _, _, valid, _ = item\n if valid:\n success_cnt += 1\n else:\n fail_cnt += 1\n pbar.update(1)\n pbar.set_postfix(success=success_cnt, fail=fail_cnt)\n\n else:\n manager = Manager()\n task_queue = manager.Queue()\n results_queue = manager.Queue()\n\n for video_path, group_df in grouped_items:\n task_queue.put((video_path, group_df))\n\n num_workers = args.num_workers if args.num_workers else os.cpu_count()\n num_workers = max(1, num_workers)\n\n with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:\n futures = []\n for _ in range(num_workers):\n futures.append(\n executor.submit(\n worker,\n task_queue,\n results_queue,\n args.video_save_dir,\n args.keep_audio, # Forward keep_audio flag to each worker\n )\n )\n\n finished = 0\n success_cnt = 0\n fail_cnt = 0\n\n with tqdm(total=total_tasks, desc=\"Processing clips\", dynamic_ncols=True) as pbar:\n while finished < total_tasks:\n try:\n item = results_queue.get(timeout=1)\n except queue.Empty:\n continue\n\n results.append(item)\n finished += 1\n\n _, _, valid, _ = item\n if valid:\n success_cnt += 1\n else:\n fail_cnt += 1\n\n pbar.update(1)\n pbar.set_postfix(success=success_cnt, fail=fail_cnt)\n\n for future in futures:\n future.result()\n\n # Sort back by original row index\n results.sort(key=lambda x: x[0])\n\n # Separate successful and failed\n success_rows = []\n failed_rows = []\n failed_errors = []\n\n for index, row_values, valid, error_msg in results:\n if valid:\n success_rows.append(row_values)\n else:\n failed_rows.append(row_values)\n failed_errors.append(error_msg)\n\n # Optional corrected timestamp CSV\n if args.drop_invalid_timestamps:\n valid_indices = [r[0] for r in results if r[2]]\n filtered_csv = csv.iloc[valid_indices]\n\n if args.csv_path.endswith(\"timestamp.csv\"):\n corrected_path = args.csv_path.replace(\"timestamp.csv\", \"correct_timestamp.csv\")\n else:\n base, ext = os.path.splitext(args.csv_path)\n corrected_path = f\"{base}_corrected{ext}\"\n\n filtered_csv.to_csv(corrected_path, index=False)\n print(f\"Corrected timestamp file saved to '{corrected_path}'\")\n\n columns = csv.columns\n\n # Save successful clips CSV\n if success_rows:\n success_df = pd.DataFrame(success_rows, columns=columns)\n\n for col in [\"timestamp_start\", \"timestamp_end\", \"frame_start\", \"frame_end\"]:\n if col in success_df.columns:\n success_df = success_df.drop(columns=[col])\n\n success_df.to_csv(args.csv_save_path, index=False)\n print(f\"Saved {len(success_df)} successful clip(s) to {args.csv_save_path}.\")\n else:\n print(\"No successful clips were generated.\")\n\n # Save failed clips CSV\n if failed_rows:\n base, ext = os.path.splitext(args.csv_save_path)\n failed_csv_path = f\"{base}_failed{ext}\"\n\n failed_df = pd.DataFrame(failed_rows, columns=columns)\n failed_df[\"error\"] = failed_errors\n failed_df.to_csv(failed_csv_path, index=False)\n print(f\"Saved {len(failed_df)} failed record(s) to {failed_csv_path}.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/download_SpatialVID.py", "content": "import argparse\nfrom huggingface_hub import hf_hub_download, snapshot_download\n\n\ndef main():\n # Setup command line arguments\n parser = argparse.ArgumentParser(\n description=\"Download SpatialVID dataset from Hugging Face Hub.\"\n )\n parser.add_argument(\n \"--repo_id\",\n type=str,\n choices=[\"SpatialVID\", \"SpatialVID-HQ\"],\n required=True,\n help=\"Dataset type to download (SpatialVID or SpatialVID-HQ)\",\n )\n parser.add_argument(\n \"--type\",\n type=str,\n choices=[\"videos\", \"annotations\", \"depths\", \"metadata\", \"all\"],\n required=True,\n help=\"Type of data to download (videos, annotations, metadata, all)\",\n )\n parser.add_argument(\n \"--group_id\",\n type=int,\n help=\"Specific group ID to download (e.g., 'group_1'). If not provided, downloads all groups.\",\n default=None,\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n help=\"Local directory to save dataset\",\n default=\"./SpatialVID_data\",\n )\n args = parser.parse_args()\n\n repo_id = f\"SpatialVID/{args.repo_id}\"\n\n # Download csv metadata\n if args.type == \"metadata\":\n hub_path = f\"data/train/{args.repo_id.replace('-', '_')}_metadata.csv\"\n hf_hub_download(\n repo_id=repo_id,\n repo_type=\"dataset\",\n filename=hub_path,\n local_dir=args.output_dir,\n resume_download=True,\n )\n print(f\"Downloaded file '{hub_path}' from {repo_id} to {args.output_dir}\")\n # Download specific group\n elif args.group_id:\n hub_path = f\"{args.type}/group_{args.group_id:04d}.tar.gz\"\n hf_hub_download(\n repo_id=repo_id,\n repo_type=\"dataset\",\n filename=hub_path,\n local_dir=args.output_dir,\n resume_download=True,\n )\n print(f\"Downloaded file '{hub_path}' from {repo_id} to {args.output_dir}\")\n # Download entire type directory\n elif args.type == \"all\":\n snapshot_download(\n repo_id=repo_id,\n repo_type=\"dataset\",\n local_dir=args.output_dir,\n resume_download=True,\n )\n print(f\"Downloaded entire dataset from {repo_id} to {args.output_dir}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/download_YouTube.py", "content": "\"\"\"\nUtility script to download YouTube videos using yt-dlp with support for concurrency and sharding.\nAdapted from https://huggingface.co/Ligeng-Zhu/panda70m-download\n\nrunning script: python download_YouTube.py --csv=\"$csv_file\" # this csv file must contains 'url' column\n\nif you want to download a specific youtube video, consider using:\n- yt-dlp -F --list-formats https://www.youtube.com/watch\\?v\\=omP01s7RUSA # --proxy 127.0.0.1:xxxx --cookies cookies.txt\n\nrun 'ls -l /path/to/folder/*.json | wc -l' for counting the videos already downloaded\n\nCustomization Guide: \n For customizing download settings (such as video format, cookie configurations like automatic Chrome cookie retrieval or custom cookie file usage), \n refer to the official documentation at https://github.com/yt-dlp/yt-dlp-wiki.\n\"\"\"\n\nimport sys, os, os.path as osp\nimport yt_dlp\nimport asyncio\nfrom concurrent.futures import ProcessPoolExecutor\nimport fire\nimport pandas as pd\nimport json\nimport time\n\n\ndef ytb_download(url, json_info, output_dir=\"ytb_videos/\"):\n \"\"\"\n Download a specified YouTube video using yt-dlp and save related metadata.\n \"\"\"\n os.makedirs(output_dir, exist_ok=True)\n uid = url.split(\"?v=\")[-1]\n yt_opts = {\n \"format\": \"bv[height=720][ext=mp4]\"\n # \"format\": \"bv[height=720]\", # Download the best quality available\n # \"format\": \"bv[height=720][ext=mp4][vcodec!^=av]\",\n # \"proxy\": \"127.0.0.1:xxxx\",\n \"outtmpl\": osp.join(output_dir, f\"{uid}.%(ext)s\"), # Output template\n # \"cookiesfrombrowser\": \"chrome\", # Use Chrome's cookies automatically\n # \"cookiefile\": \"cookies.txt\", # Use a custom cookies file\n # \"postprocessors\": [\n # {\n # \"key\": \"FFmpegVideoConvertor\",\n # \"preferedformat\": \"mp4\", # Convert video to mp4 format (slow)\n # }\n # ],\n # \"verbose\" : True,\n \"abort-on-error\": True, # Abort downloading when an error occurs\n \"retries\": 60, # Number of retries\n \"ffmpeg_location\": \"/usr/bin/ffmpeg\", # Path to ffmpeg\n \"quiet\": True, # Suppress output\n \"sleep-requested\": 5, # Sleep for 1.25 seconds between requests\n \"min-sleep-interval\": 60,\n \"max-sleep-interval\": 90,\n }\n video_path_mp4 = osp.join(output_dir, f\"{uid}.mp4\")\n video_path_webm = osp.join(output_dir, f\"{uid}.webm\")\n meta_path = osp.join(output_dir, f\"{uid}.json\")\n if (osp.exists(video_path_mp4) or osp.exists(video_path_webm)) and osp.exists(\n meta_path\n ):\n print(f\"\\033[91m{uid} already labeled.\\033[0m\")\n return 0\n try:\n with yt_dlp.YoutubeDL(yt_opts) as ydl:\n ydl.download([url])\n with open(meta_path, \"w\") as fp:\n json.dump(json_info, fp, indent=2)\n return 0\n # exception logs\n except Exception as e:\n print(f\"\\033[91mError downloading {url}: {e}\\033[0m\")\n err_map = {\n \"Requested format is not available\": \"z0322_dld_format_noavailable.log\",\n \"removed by\": \"z0322_dld_removed_by.log\",\n \"Private video\": \"z0322_dld_private_video.log\",\n }\n for key, log_file in err_map.items():\n if key in str(e):\n with open(osp.join(output_dir, f\"{log_file}\"), \"a\") as f:\n f.write(f\"{url}\\n\")\n break\n else:\n with open(osp.join(output_dir, f\"z0322_dld_othererr.log\"), \"a\") as f:\n f.write(f\"{url}, {str(e)}\\n\")\n return -1\n\n\nasync def main(csv_path, output_dir, max_workers=10, shards=0, total=-1, limit=False):\n \"\"\"\n Batch download YouTube videos specified in a CSV file, supporting sharding and concurrency.\n \"\"\"\n PPE = ProcessPoolExecutor(max_workers=max_workers)\n loop = asyncio.get_event_loop()\n df = pd.read_csv(csv_path)\n csv_path = os.path.basename(csv_path)\n output_dir = f'{output_dir}/{csv_path.split(\".\")[0]}'\n data_list = list(df.iterrows())\n if total > 0:\n chunk = len(data_list) // total\n begin_idx = shards * chunk\n end_idx = (shards + 1) * chunk if shards < total - 1 else len(data_list)\n data_list = data_list[begin_idx:end_idx]\n print(f\"download total {len(data_list)} videos\")\n tasks = []\n for idx, (index, row) in enumerate(data_list):\n video_url = row[\"url\"]\n # json_info = {\"caption\": row[\"caption\"]}\n json_info = {\"caption\": ''} # for file checking.\n tasks.append(\n loop.run_in_executor(PPE, ytb_download, video_url, json_info, output_dir)\n )\n if limit and idx >= 20:\n break\n res = await asyncio.gather(*tasks)\n print(f\"[{sum(res)} / {len(res)}]\")\n\n\ndef entry(\n csv=\"meta_data_sample_500.csv\",\n output_dir=\"path/to/output\",\n shards=0,\n total=-1,\n limit=False,\n max_workers=2,\n):\n \"\"\"\n Command line entry function, supports fire invocation.\n \"\"\"\n print(csv, output_dir, shards, total, max_workers)\n start_time = time.time()\n print(\n f\"\\033[92mStarting execution at {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}\\033[0m\"\n )\n asyncio.run(\n main(\n csv,\n output_dir,\n max_workers=max_workers,\n shards=shards,\n total=total,\n limit=limit,\n )\n )\n end_time = time.time()\n print(\n f\"\\033[92mFinished execution at {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))}\\033[0m\"\n )\n print(f\"\\033[92mTotal execution time: {end_time - start_time:.2f} seconds\\033[0m\")\n\n\ndef add_download(csv_path):\n \"\"\"\n Download missing videos according to the new_vid_path field in the CSV file.\n \"\"\"\n data = pd.read_csv(csv_path)\n unique_ids = data['YouTube id'].unique()\n\n for uid in unique_ids:\n video_url = f\"https://www.youtube.com/watch?v={uid}\"\n ytb_download(video_url, json_info={}, output_dir=\"videos/\")\n print(f\"Downloaded {video_url}\")\n\n\nif __name__ == \"__main__\":\n # Call entry function via command line arguments\n \n fire.Fire(entry)\n # for supplement download: add_download(csv_path='xxx.csv')\n" }, { "path": "utils/evaluation.py", "content": "\"\"\"\nCamera trajectory evaluation utility with anomaly detection and motion analysis.\n\"\"\"\n\nimport os\nimport argparse\nimport pandas as pd\nimport numpy as np\nimport torch\nimport concurrent.futures\nimport multiprocessing as mp\nfrom multiprocessing import Manager\nimport queue\nfrom tqdm import tqdm\nfrom scipy.signal import find_peaks\nfrom scipy.ndimage import gaussian_filter\n\n# Import mask utility functions\nfrom expand_npz import expand\n\n\ndef load_file(cam_pos_file, mask_file, device):\n \"\"\"Load camera parameters and dynamic masks from files\"\"\"\n try:\n # Load camera parameters and split into position and rotation\n params = torch.from_numpy(np.load(cam_pos_file)).float().to(device)\n cam_pos = params[:, :3] # Position coordinates\n cam_rotate = params[:, 3:] # Rotation quaternions\n time_steps = params.shape[0]\n\n # Load and expand dynamic masks\n masks = torch.from_numpy(expand(np.load(mask_file))).to(device)\n except FileNotFoundError:\n print(f\"Error: File not found - {cam_pos_file}\")\n exit()\n except Exception as e:\n print(f\"Error processing {cam_pos_file}: {e}\")\n exit()\n\n return cam_pos, cam_rotate, time_steps, masks\n\n\ndef anomaly_detection(cam_pos, time_steps, threshold, device):\n \"\"\"Detect trajectory anomalies using linear prediction with acceleration\"\"\"\n if time_steps < 4:\n return True # Not enough data\n\n preds = torch.zeros((time_steps, 3), dtype=torch.float32, device=device)\n error_count = 0\n\n # Linear prediction with acceleration\n for t in range(0, time_steps - 3):\n # Calculate velocity and acceleration\n v1 = cam_pos[t + 2] - cam_pos[t + 1]\n v2 = cam_pos[t + 1] - cam_pos[t]\n acceleration = v1 - v2\n\n # Predict next position\n preds[t + 3] = cam_pos[t + 2] + v1 + 0.5 * acceleration\n\n # Check prediction error\n error = torch.sqrt(torch.sum((preds[t + 3] - cam_pos[t + 3]) ** 2))\n if error > 0.03:\n error_count += 1\n if error_count >= threshold:\n return True\n else:\n error_count = 0\n\n return False\n\n\ndef move_distance(cam_pos, time_steps, device):\n \"\"\"Calculate total movement distance and classify into levels\"\"\"\n total_distance = torch.tensor(0., dtype=torch.float32, device=device)\n\n # Distance thresholds for classification\n thresholds = [0.08, 0.28, 0.92, 2.41]\n \n # Calculate cumulative distance\n for i in range(0, time_steps - 1):\n total_distance += torch.norm(cam_pos[i + 1] - cam_pos[i])\n\n # Determine movement level\n distance_val = total_distance.item()\n level = sum(1 for threshold in thresholds if distance_val >= threshold)\n\n return distance_val, level\n\n\ndef quaternion_multiply(q1, q2):\n \"\"\"Multiply two quaternions\"\"\"\n # Extract components (q in [x, y, z, w] format)\n w1, x1, y1, z1 = q1[3], q1[0], q1[1], q1[2]\n w2, x2, y2, z2 = q2[3], q2[0], q2[1], q2[2]\n\n # Quaternion multiplication\n matrix = torch.tensor([\n [w1, -z1, y1, x1],\n [z1, w1, -x1, y1],\n [-y1, x1, w1, z1],\n [-x1, -y1, -z1, w1]\n ], dtype=q1.dtype, device=q1.device)\n\n vector = torch.tensor([x2, y2, z2, w2], dtype=q2.dtype, device=q2.device)\n result = torch.matmul(matrix, vector)\n\n return result\n\n\ndef rotation_angle(cam_rotate, time_steps, device):\n \"\"\"Calculate total rotation angle between consecutive frames\"\"\"\n total_radians = torch.tensor(0.0, device=device)\n\n for i in range(0, time_steps - 1):\n q1 = cam_rotate[i]\n q2 = cam_rotate[i + 1]\n\n # Calculate relative rotation\n q1_inverse = torch.stack([-q1[0], -q1[1], -q1[2], q1[3]], dim=0)\n q_relative = quaternion_multiply(q2, q1_inverse)\n w = torch.clamp(q_relative[3], -1.0, 1.0)\n\n # Convert to angle\n rotation_angle_rad = 2 * torch.arccos(w)\n total_radians += rotation_angle_rad\n\n return total_radians.item()\n\n\ndef trajectory_turns(cam_pos, time_steps, device, threshold=0.45):\n \"\"\"Detect significant turns in camera trajectory\"\"\"\n if time_steps < 3:\n return [], 0\n\n angles = []\n # Calculate angles between trajectory segments\n for t in range(1, time_steps - 1):\n v1 = cam_pos[t] - cam_pos[0]\n v2 = cam_pos[time_steps - 1] - cam_pos[t]\n\n # Avoid division by zero\n v1_norm = torch.norm(v1)\n v2_norm = torch.norm(v2)\n if v1_norm < 1e-8 or v2_norm < 1e-8:\n continue\n\n # Calculate angle between vectors\n cos_theta = torch.dot(v1, v2) / (v1_norm * v2_norm)\n cos_theta = torch.clamp(cos_theta, -1.0, 1.0)\n angle = torch.arccos(cos_theta)\n angles.append(angle.item())\n\n # Smooth and find peaks\n angles = gaussian_filter(angles, sigma=5)\n peaks, _ = find_peaks(angles, height=threshold, distance=5)\n peaks_values = [angles[i] for i in peaks]\n\n # Include maximum angle if significant\n max_angle = max(angles)\n if max_angle > threshold and max_angle not in peaks_values:\n peaks_values.append(max_angle)\n\n return len(peaks_values)\n\n\ndef dynamic_ratio(masks):\n \"\"\"Calculate ratio of dynamic pixels in video frames\"\"\"\n # Downsample for efficiency\n masks = masks[::5, :, :]\n \n dynamic_pixels = torch.sum(masks)\n total_pixels = masks.shape[1] * masks.shape[2] * masks.shape[0]\n \n return (dynamic_pixels / total_pixels).item()\n\n\ndef process_single_row(row, index, args, device):\n \"\"\"Process a single video row to extract trajectory metrics\"\"\"\n video_id = row['id']\n rec_path = os.path.join(args.dir_path, video_id, \"reconstructions\")\n cam_pos_file = os.path.join(rec_path, \"poses.npy\")\n mask_file = os.path.join(rec_path, \"dyn_masks.npz\")\n\n # Check file existence\n if not os.path.exists(cam_pos_file) or not os.path.exists(mask_file):\n print(f\"File not found: {cam_pos_file} or {mask_file}\")\n return False, False, -1, -1, -1, -1, -1\n\n # Load and process data\n cam_pos, cam_rotate, time_steps, masks = load_file(cam_pos_file, mask_file, device)\n\n # Calculate metrics\n anomaly = anomaly_detection(cam_pos, time_steps, args.anomaly_threshold, device)\n move_dist, dist_level = move_distance(cam_pos, time_steps, device)\n rot_angle = rotation_angle(cam_rotate, time_steps, device)\n traj_turns = trajectory_turns(cam_pos, time_steps, device)\n dyn_ratio = dynamic_ratio(masks)\n\n return True, anomaly, move_dist, dist_level, rot_angle, traj_turns, dyn_ratio\n\n\ndef worker(task_queue, result_queue, args, worker_id):\n \"\"\"Worker function for parallel processing\"\"\"\n # Assign GPU based on worker ID\n device = torch.device(\n f\"cuda:{worker_id % args.gpu_num}\"\n if torch.cuda.is_available() else \"cpu\"\n )\n\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n\n result = process_single_row(row, index, args, device)\n result_queue.put((index, result))\n task_queue.task_done()\n\n\ndef parse_args():\n \"\"\"Parse command line arguments\"\"\"\n parser = argparse.ArgumentParser(description=\"Camera Trajectory Evaluation\")\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to input CSV file\")\n parser.add_argument(\"--dir_path\", type=str, default=\"./outputs\", help=\"Base directory with reconstruction data\")\n parser.add_argument(\"--output_path\", type=str, default=\"./outputs/evaluation_results.csv\", help=\"Output CSV path\")\n parser.add_argument(\"--anomaly_threshold\", type=int, default=2, help=\"Anomaly detection threshold\")\n parser.add_argument('--gpu_num', type=int, default=1, help='Number of GPUs to use')\n parser.add_argument(\"--num_workers\", type=int, default=4, help=\"Number of parallel workers\")\n parser.add_argument(\"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\")\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n # Setup multiprocessing\n mp.set_start_method('spawn')\n args = parse_args()\n\n # Load input data\n df = pd.read_csv(args.csv_path)\n\n results = []\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(df.iterrows(), total=len(df), desc=\"Processing videos\"):\n device = torch.device(f\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n result = process_single_row(row, index, args, device)\n results.append((index, result))\n else:\n # Parallel processing\n manager = Manager()\n task_queue = manager.Queue()\n\n # Add tasks to queue\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n result_queue = manager.Queue()\n\n # Run workers\n with concurrent.futures.ProcessPoolExecutor(max_workers=args.num_workers) as executor:\n futures = []\n for worker_id in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, result_queue, args, worker_id))\n\n processed = 0\n total_tasks = len(df)\n with tqdm(total=total_tasks, desc=\"Processing videos\") as pbar:\n while processed < total_tasks:\n try:\n index, result = result_queue.get(timeout=1)\n results.append((index, result))\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and result_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n # Collect results\n while not result_queue.empty():\n index, result = result_queue.get()\n results.append((index, result))\n\n # Sort and save results\n results.sort(key=lambda x: x[0])\n\n df['success'] = [result[1][0] for result in results]\n df['anomaly'] = [result[1][1] for result in results]\n df['moveDist'] = [result[1][2] for result in results]\n df['distLevel'] = [result[1][3] for result in results]\n df['rotAngle'] = [result[1][4] for result in results]\n df['trajTurns'] = [result[1][5] for result in results]\n df['dynamicRatio'] = [result[1][6] for result in results]\n\n df.to_csv(args.output_path, index=False)\n print(f\"Results saved to {args.output_path}\")\n" }, { "path": "utils/expand_npz.py", "content": "\"\"\"\nMask utility functions for processing sparse matrix data.\n\"\"\"\n\nimport numpy as np\nfrom scipy.sparse import csr_matrix\n\n\ndef expand(loaded_data):\n \"\"\"\n Reconstruct 3D mask from sparse matrix data.\n \n Args:\n loaded_data (dict): Dictionary containing sparse matrix data with keys:\n - 'shape': Original matrix dimensions\n - 'f_{i}_data': Sparse matrix data for frame i\n - 'f_{i}_indices': Sparse matrix indices for frame i \n - 'f_{i}_indptr': Sparse matrix index pointers for frame i\n \n Returns:\n np.ndarray: 3D array with shape (frames, height, width)\n \"\"\"\n reconstructed_sparse_matrices = []\n num_frames = (len(loaded_data) - 1) // 3 # Calculate number of frames\n matrix_shape = loaded_data['shape'] # Get original matrix dimensions\n\n # Reconstruct sparse matrix for each frame\n for i in range(num_frames):\n data = loaded_data[f'f_{i}_data']\n indices = loaded_data[f'f_{i}_indices']\n indptr = loaded_data[f'f_{i}_indptr']\n reconstructed_matrix = csr_matrix((data, indices, indptr), shape=matrix_shape)\n reconstructed_sparse_matrices.append(reconstructed_matrix)\n\n # Stack all frames into a 3D array (frames, height, width)\n reconstructed_mask_3d = np.stack([m.toarray() for m in reconstructed_sparse_matrices], axis=0)\n return reconstructed_mask_3d\n" }, { "path": "utils/extract_frames.py", "content": "\"\"\"\nVideo frame extraction utility with parallel processing support.\n\"\"\"\n\nimport os\nimport sys\nimport cv2\nimport av\nimport glob\nimport argparse\nimport pandas as pd\nimport queue\nimport concurrent.futures\nfrom multiprocessing import Manager\nfrom tqdm import tqdm\n\n\ndef extract_frames_opencv(\n video_path, output_dir, interval, frame_start, num_frames, target_size=None\n):\n \"\"\"Extract frames from video at specified intervals\"\"\"\n # Create output directory\n if not os.path.exists(output_dir):\n os.makedirs(output_dir)\n\n # Open video file\n cap = cv2.VideoCapture(video_path)\n cap.set(cv2.CAP_PROP_POS_FRAMES, frame_start)\n\n if not cap.isOpened():\n print(f\"Error: Could not open video file {video_path}\")\n sys.exit(1)\n\n # Extract frames\n for frame in range(num_frames):\n ret, image = cap.read()\n if not ret:\n break\n\n # Save frame at specified intervals\n if frame % interval == 0:\n frame_filename = os.path.join(output_dir, f\"frame_{frame:06d}.jpg\")\n if target_size is not None:\n h, w = image.shape[:2]\n # Adaptively adjust target size based on video orientation\n # For portrait videos (height > width), swap width and height of target size\n if h > w: # Portrait video\n target_w, target_h = target_size[1], target_size[0]\n else: # Landscape video\n target_w, target_h = target_size[0], target_size[1]\n image = cv2.resize(image, (target_w, target_h))\n cv2.imwrite(frame_filename, image)\n\n cap.release()\n\n\ndef extract_frames_av(\n video_path, output_dir, interval, frame_start, num_frames, target_size=None\n):\n \"\"\"\n Extract frames from video at specified intervals using PyAV backend.\n \"\"\"\n # Create output directory\n if not os.path.exists(output_dir):\n os.makedirs(output_dir)\n\n try:\n # Open video file\n container = av.open(video_path)\n stream = container.streams.video[0]\n stream.thread_type = 'AUTO'\n except Exception as e:\n print(f\"Error: Could not open video file {video_path}. Reason: {e}\")\n return\n\n # Get video properties\n fps = float(stream.average_rate)\n time_base = stream.time_base\n \n target_sec = frame_start / fps\n # Set a small tolerance (e.g., half a frame time) to prevent frame loss due to floating-point precision issues\n epsilon = 0.5 / fps \n\n # Seek to the target start time\n if frame_start > 0:\n target_pts = int(target_sec / time_base)\n container.seek(target_pts, stream=stream, backward=True)\n\n count = 0\n for packet in container.demux(stream):\n try:\n for frame in packet.decode():\n if frame.pts is None:\n continue\n\n current_sec = frame.pts * time_base\n if current_sec < (target_sec - epsilon):\n continue\n\n if count >= num_frames:\n break\n\n if count % interval == 0:\n image = frame.to_ndarray(format='bgr24')\n frame_filename = os.path.join(output_dir, f\"frame_{count:06d}.jpg\")\n if target_size is not None:\n if isinstance(target_size, str):\n w, h = map(int, target_size.split('*'))\n target_size = (w, h)\n h, w = image.shape[:2]\n if h > w:\n target_w, target_h = target_size[1], target_size[0]\n else:\n target_w, target_h = target_size[0], target_size[1]\n image = cv2.resize(image, (target_w, target_h))\n cv2.imwrite(frame_filename, image)\n\n count += 1\n if count >= num_frames:\n break\n except av.error.InvalidDataError:\n continue # 跳过损坏的包\n\n container.close()\n\n\ndef _calc_expected_frames(num_frames, interval):\n \"\"\"Calculate the expected number of output frames based on total frames and interval.\"\"\"\n if interval <= 0:\n return num_frames\n # Frames at indices 0, interval, 2*interval, ... that are < num_frames\n return (num_frames - 1) // interval + 1\n\n\ndef _verify_frames(img_dir, expected_frames):\n \"\"\"Check if img_dir has enough valid (non-empty) frame files.\n\n Returns True if the directory exists and contains at least `expected_frames`\n non-zero-byte frame_*.jpg files.\n \"\"\"\n if not os.path.isdir(img_dir):\n return False\n frame_files = glob.glob(os.path.join(img_dir, \"frame_*.jpg\"))\n if len(frame_files) < expected_frames:\n return False\n if any(os.path.getsize(f) == 0 for f in frame_files):\n return False\n return True\n\n\ndef process_single_row(row, row_index, args):\n \"\"\"Process a single video row to extract frames.\n\n Returns:\n True if processing succeeded or was skipped (already done),\n False if an error occurred.\n \"\"\"\n video_path = row[\"video_path\"]\n frame_start = row.get(\"frame_start\", 0)\n num_frames = row[\"num_frames\"]\n output_dir = os.path.join(args.output_dir, row[\"id\"])\n\n img_dir = os.path.join(output_dir, \"img\")\n\n # Calculate frame extraction interval\n if args.interval is None:\n interval = row[\"num_frames\"] // 3 # Extract 3 frames by default\n elif args.interval == 0:\n interval = 1 # Extract every frame\n else:\n interval = int(args.interval * row[\"fps\"])\n\n expected_frames = _calc_expected_frames(num_frames, interval)\n\n # --- Skip logic: already has enough valid frames ---\n if _verify_frames(img_dir, expected_frames):\n return True\n\n if not os.path.exists(output_dir):\n os.makedirs(output_dir)\n\n try:\n if args.backend == \"opencv\":\n extract_frames_opencv(\n video_path, img_dir, interval, frame_start, num_frames, args.target_size\n )\n elif args.backend == \"av\":\n extract_frames_av(\n video_path, img_dir, interval, frame_start, num_frames, args.target_size\n )\n\n # Post-extraction verification\n if not _verify_frames(img_dir, expected_frames):\n actual_count = len(glob.glob(os.path.join(img_dir, \"frame_*.jpg\")))\n print(\n f\"[Verify FAIL] {row['id']}: expected {expected_frames} frames, \"\n f\"got {actual_count} (or contains empty files).\"\n )\n return False\n\n return True\n except Exception as e:\n print(f\"Error: Could not extract frames from video {video_path}. Reason: {e}\")\n return False\n\n\ndef worker(task_queue, progress_queue, failed_indices, args):\n \"\"\"Worker function for parallel frame extraction\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n success = process_single_row(row, index, args)\n if not success:\n failed_indices.append(index)\n progress_queue.put(index)\n task_queue.task_done()\n\n\ndef parse_args():\n \"\"\"Parse command line arguments\"\"\"\n parser = argparse.ArgumentParser(description=\"Extract frames from video files\")\n parser.add_argument(\n \"--csv_path\", type=str, help=\"Path to CSV file with video csvdata\"\n )\n parser.add_argument(\n \"--output_dir\",\n type=str,\n default=\"extract_frames\",\n help=\"Output directory for extracted frames\",\n )\n parser.add_argument(\n \"--interval\",\n type=float,\n default=None,\n help=\"Frame extraction interval in seconds (set to 0 to extract every frame)\",\n )\n parser.add_argument(\n \"--target_size\",\n type=str,\n default=None,\n help=\"Resize frames to size (width*height). For portrait videos (h>w), dimensions will be automatically swapped to (height*width) to maintain correct orientation.\",\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=None, help=\"Number of parallel workers\"\n )\n parser.add_argument(\n \"--backend\",\n type=str,\n default=\"opencv\",\n choices=[\"opencv\", \"av\"],\n help=\"Backend for video reading\",\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\ndef main():\n \"\"\"Main function to process frame extraction\"\"\"\n args = parse_args()\n\n # Parse target size if provided\n if args.target_size is not None:\n args.target_size = tuple(map(int, args.target_size.split(\"*\")))\n\n # Create output directory\n if not os.path.exists(args.output_dir):\n os.makedirs(args.output_dir)\n\n # Load video csvdata\n csv = pd.read_csv(args.csv_path)\n failed_indices = []\n\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(\n csv.iterrows(), total=len(csv), desc=\"Processing videos\"\n ):\n success = process_single_row(row, index, args)\n if not success:\n failed_indices.append(index)\n else:\n # Parallel processing\n num_workers = args.num_workers if args.num_workers else os.cpu_count() or 1\n\n manager = Manager()\n task_queue = manager.Queue()\n progress_queue = manager.Queue()\n shared_failed_indices = manager.list()\n\n # Add tasks to queue\n for index, row in csv.iterrows():\n task_queue.put((index, row))\n\n # Execute workers\n with concurrent.futures.ProcessPoolExecutor(\n max_workers=num_workers\n ) as executor:\n futures = []\n for _ in range(num_workers):\n future = executor.submit(worker, task_queue, progress_queue, shared_failed_indices, args)\n futures.append(future)\n\n processed = 0\n total_tasks = len(csv)\n with tqdm(total=total_tasks, desc=\"Processing videos\") as pbar:\n while processed < total_tasks:\n try:\n progress_queue.get(timeout=1)\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and progress_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n failed_indices = list(shared_failed_indices)\n\n # Save failed rows to a separate CSV; keep only successful rows in the original CSV\n if failed_indices:\n failed_csv = csv.loc[failed_indices]\n base, ext = os.path.splitext(args.csv_path)\n failed_csv_path = f\"{base}_failed{ext}\"\n failed_csv.to_csv(failed_csv_path, index=False)\n\n csv = csv.drop(index=failed_indices)\n csv.to_csv(args.csv_path, index=False)\n\n print(f\"\\n{len(failed_indices)} video(s) failed. Saved to: {failed_csv_path}\")\n print(f\"Original CSV updated. Remaining rows: {len(csv)}\")\n else:\n print(\"\\nAll videos processed successfully.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/filter.py", "content": "\"\"\"\nDataset filtering utility for video metadata with various quality metrics.\n\"\"\"\n\nimport argparse\nimport os\nimport random\nfrom glob import glob\nimport numpy as np\nimport pandas as pd\n\n\ndef main(args):\n \"\"\"Apply filtering criteria to dataset\"\"\"\n # Load data\n data = pd.read_csv(args.csv_path)\n\n # Apply filters based on various metrics\n if args.frames_min is not None:\n assert \"num_frames\" in data.columns\n data = data[data[\"num_frames\"] >= args.frames_min]\n if args.frames_max is not None:\n assert \"num_frames\" in data.columns\n data = data[data[\"num_frames\"] <= args.frames_max]\n if args.fps_max is not None:\n assert \"fps\" in data.columns\n data = data[(data[\"fps\"] <= args.fps_max) | np.isnan(data[\"fps\"])]\n if args.fps_min is not None:\n assert \"fps\" in data.columns\n data = data[(data[\"fps\"] >= args.fps_min) | np.isnan(data[\"fps\"])]\n if args.resolution_max is not None:\n if \"resolution\" not in data.columns:\n height = data[\"height\"]\n width = data[\"width\"]\n data[\"resolution\"] = height * width\n data = data[data[\"resolution\"] <= args.resolution_max]\n if args.aes_min is not None:\n assert \"aesthetic score\" in data.columns\n data = data[data[\"aesthetic score\"] >= args.aes_min]\n if args.ocr_max is not None:\n assert \"ocr score\" in data.columns\n data = data[data[\"ocr score\"] <= args.ocr_max]\n if args.ocr_min is not None:\n assert \"ocr score\" in data.columns\n data = data[data[\"ocr score\"] >= args.ocr_min]\n if args.lum_min is not None:\n assert \"luminance mean\" in data.columns\n data = data[data[\"luminance mean\"] >= args.lum_min]\n if args.lum_max is not None:\n assert \"luminance mean\" in data.columns\n data = data[data[\"luminance mean\"] <= args.lum_max]\n if args.motion_min is not None:\n assert \"motion score\" in data.columns\n data = data[data[\"motion score\"] >= args.motion_min]\n if args.motion_max is not None:\n assert \"motion score\" in data.columns\n data = data[data[\"motion score\"] <= args.motion_max]\n\n # Save filtered data\n data.to_csv(args.csv_save_path, index=False)\n print(f\"Saved {len(data)} samples to {args.csv_save_path}.\")\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for dataset filtering\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Filter video dataset by quality metrics\"\n )\n parser.add_argument(\n \"--csv_path\", type=str, required=True, help=\"Path to input CSV file\"\n )\n parser.add_argument(\n \"--csv_save_path\", type=str, default=None, help=\"Path to save output CSV file\"\n )\n parser.add_argument(\"--seed\", type=int, default=42, help=\"Random seed\")\n\n # Video property filters\n parser.add_argument(\n \"--frames_min\", type=int, default=None, help=\"Minimum number of frames\"\n )\n parser.add_argument(\n \"--frames_max\", type=int, default=None, help=\"Maximum number of frames\"\n )\n parser.add_argument(\n \"--resolution_max\", type=int, default=None, help=\"Maximum resolution\"\n )\n parser.add_argument(\"--fps_max\", type=float, default=None, help=\"Maximum FPS\")\n parser.add_argument(\"--fps_min\", type=float, default=None, help=\"Minimum FPS\")\n\n # Quality metric filters\n parser.add_argument(\n \"--aes_min\", type=float, default=None, help=\"Minimum aesthetic score\"\n )\n parser.add_argument(\n \"--flow_min\", type=float, default=None, help=\"Minimum optical flow score\"\n )\n parser.add_argument(\n \"--flow_max\", type=float, default=None, help=\"Maximum optical flow score\"\n )\n parser.add_argument(\"--ocr_max\", type=float, default=None, help=\"Maximum OCR score\")\n parser.add_argument(\"--ocr_min\", type=float, default=None, help=\"Minimum OCR score\")\n parser.add_argument(\n \"--lum_min\", type=float, default=None, help=\"Minimum luminance score\"\n )\n parser.add_argument(\n \"--lum_max\", type=float, default=None, help=\"Maximum luminance score\"\n )\n parser.add_argument(\n \"--blur_max\", type=float, default=None, help=\"Maximum blur score\"\n )\n parser.add_argument(\n \"--motion_min\", type=float, default=None, help=\"Minimum motion score\"\n )\n parser.add_argument(\n \"--motion_max\", type=float, default=None, help=\"Maximum motion score\"\n )\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n # Set random seeds for reproducibility\n if args.seed is not None:\n random.seed(args.seed)\n np.random.seed(args.seed)\n main(args)\n" }, { "path": "utils/get_clip.py", "content": "\"\"\"\nVideo clip information extraction utility with timestamp parsing.\n\"\"\"\n\nimport argparse\nimport os\nimport queue\nimport concurrent.futures\nfrom functools import partial\nimport pandas as pd\nfrom scenedetect import FrameTimecode\nimport re\nfrom tqdm import tqdm\n\n\ndef process_single_row(row, args):\n \"\"\"Process a single video row to extract clip information\"\"\"\n video_path = row[\"video_path\"]\n new_rows = []\n\n try:\n if \"timestamp\" in row:\n timestamp_str = row[\"timestamp\"]\n # Parse timestamps using regex\n timestamp_pattern = (\n r\"\\('(\\d{2}:\\d{2}:\\d{2}\\.\\d+)', '(\\d{2}:\\d{2}:\\d{2}\\.\\d+)'\\)\"\n )\n matches = re.findall(timestamp_pattern, timestamp_str)\n scene_list = [\n (FrameTimecode(s, fps=row[\"fps\"]), FrameTimecode(t, fps=row[\"fps\"]))\n for s, t in matches\n ]\n else:\n scene_list = [None]\n if args.drop_invalid_timestamps:\n return new_rows, True\n except Exception as e:\n if args.drop_invalid_timestamps:\n return new_rows, False\n\n height = row[\"height\"]\n width = row[\"width\"]\n fps = row[\"fps\"]\n\n # Extract clip information for each scene\n for idx, scene in enumerate(scene_list):\n if scene is not None:\n s, t = scene # FrameTimecode objects\n\n fname = os.path.basename(video_path)\n fname_wo_ext = os.path.splitext(fname)[0]\n\n # Calculate clip metrics\n num_frames = t.frame_num - s.frame_num\n aspect_ratio = width / height if height != 0 else 0\n resolution = f\"{width}x{height}\"\n timestamp_start = s.get_timecode()\n timestamp_end = t.get_timecode()\n frame_start = s.frame_num\n frame_end = t.frame_num\n id_ori = row[\"id\"] if \"id\" in row else \"\"\n id = f\"{fname_wo_ext}_{idx}\"\n new_rows.append(\n [\n video_path,\n id,\n num_frames,\n height,\n width,\n aspect_ratio,\n fps,\n resolution,\n timestamp_start,\n timestamp_end,\n frame_start,\n frame_end,\n id_ori,\n ]\n )\n\n return (new_rows, True)\n\n\ndef worker(task_queue, results_queue, args):\n \"\"\"Worker function for parallel processing\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n result = process_single_row(row, args)\n results_queue.put((index, result))\n task_queue.task_done()\n\n\ndef parse_args():\n \"\"\"Parse command line arguments\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Extract video clip information from csvdata\"\n )\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the input CSV file\")\n parser.add_argument(\n \"--csv_save_dir\",\n type=str,\n required=True,\n help=\"Directory to save output CSV file\",\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=None, help=\"Number of parallel workers\"\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n parser.add_argument(\n \"--drop_invalid_timestamps\",\n action=\"store_true\",\n help=\"Drop rows with invalid timestamps\",\n )\n args = parser.parse_args()\n return args\n\n\ndef main():\n args = parse_args()\n csv_path = args.csv_path\n if not os.path.exists(csv_path):\n print(f\"csv file '{csv_path}' not found. Exit.\")\n return\n\n os.makedirs(args.csv_save_dir, exist_ok=True)\n\n # Load csvdata\n csv = pd.read_csv(args.csv_path)\n\n # Setup multiprocessing\n from multiprocessing import Manager\n\n manager = Manager()\n task_queue = manager.Queue()\n results_queue = manager.Queue()\n\n for index, row in csv.iterrows():\n task_queue.put((index, row))\n\n if args.disable_parallel:\n # Sequential processing\n results = []\n for index, row in tqdm(\n csv.iterrows(), total=len(csv), desc=\"Processing rows\"\n ):\n result = process_single_row(row, args)\n results.append((index, result))\n else:\n # Parallel processing\n num_workers = args.num_workers if args.num_workers else os.cpu_count() or 1\n\n with concurrent.futures.ProcessPoolExecutor(\n max_workers=num_workers\n ) as executor:\n futures = []\n for _ in range(num_workers):\n future = executor.submit(worker, task_queue, results_queue, args)\n futures.append(future)\n\n # Per-row progress is more informative than per-worker completion.\n results = []\n processed = 0\n total_tasks = len(csv)\n with tqdm(total=total_tasks, desc=\"Processing rows\") as pbar:\n while processed < total_tasks:\n try:\n results.append(results_queue.get(timeout=1))\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and results_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n while not results_queue.empty():\n results.append(results_queue.get())\n\n # Process results\n results.sort(key=lambda x: x[0])\n new_rows = []\n valid_rows = []\n for index, (rows, valid) in results:\n if valid:\n valid_rows.append(index)\n new_rows.extend(rows)\n\n # Save corrected timestamps if needed\n if args.drop_invalid_timestamps:\n csv = csv[valid_rows]\n assert args.csv_path.endswith(\"timestamp.csv\"), \"Only support *timestamp.csv\"\n csv.to_csv(\n args.csv_path.replace(\"timestamp.csv\", \"correct_timestamp.csv\"),\n index=False,\n )\n print(\n f\"Corrected timestamp file saved to '{args.csv_path.replace('timestamp.csv', 'correct_timestamp.csv')}'\"\n )\n\n # Create and save clip information DataFrame\n columns = [\n \"video_path\",\n \"id\",\n \"num_frames\",\n \"height\",\n \"width\",\n \"aspect_ratio\",\n \"fps\",\n \"resolution\",\n \"timestamp_start\",\n \"timestamp_end\",\n \"frame_start\",\n \"frame_end\",\n \"id_ori\",\n ]\n new_df = pd.DataFrame(new_rows, columns=columns)\n new_csv_path = os.path.join(args.csv_save_dir, \"clips_info.csv\")\n new_df.to_csv(new_csv_path, index=False)\n print(f\"Saved {len(new_df)} clip information to {new_csv_path}.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/get_info.py", "content": "\"\"\"\nVideo information extraction utility supporting multiple backends (OpenCV, TorchVision, AV).\n\"\"\"\n\nimport argparse\nimport os\nimport random\nimport cv2\nimport av\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport concurrent.futures\n\n\ndef get_video_length(cap, method=\"header\"):\n \"\"\"Get video frame count using different methods\"\"\"\n assert method in [\"header\", \"set\"]\n if method == \"header\":\n length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n else:\n cap.set(cv2.CAP_PROP_POS_AVI_RATIO, 1)\n length = int(cap.get(cv2.CAP_PROP_POS_FRAMES))\n return length\n\n\ndef get_video_info(args):\n \"\"\"Extract video information using specified backend\"\"\"\n idx, path, backend = args\n try:\n if backend == \"torchvision\":\n from tools.utils.read_video import read_video\n\n vframes, infos = read_video(path)\n num_frames, height, width = (\n vframes.shape[0],\n vframes.shape[2],\n vframes.shape[3],\n )\n fps = (\n float(infos.get(\"video_fps\", np.nan))\n if isinstance(infos, dict)\n else np.nan\n )\n elif backend == \"opencv\":\n cap = cv2.VideoCapture(path)\n if not cap.isOpened():\n raise ValueError(\"Video open failed\")\n num_frames = get_video_length(cap, method=\"header\")\n height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n fps = float(cap.get(cv2.CAP_PROP_FPS))\n cap.release()\n elif backend == \"av\":\n container = av.open(path)\n stream = container.streams.video[0]\n num_frames = int(stream.frames)\n height = int(stream.height)\n width = int(stream.width)\n if stream.average_rate is not None:\n fps = float(stream.average_rate)\n elif stream.guessed_rate is not None:\n fps = float(stream.guessed_rate)\n else:\n fps = np.nan\n else:\n raise ValueError(\"Unknown backend\")\n\n # Calculate derived metrics\n hw = height * width\n aspect_ratio = height / width if width > 0 else np.nan\n return (idx, True, num_frames, height, width, aspect_ratio, hw, fps)\n except Exception:\n return (idx, False, 0, 0, 0, np.nan, np.nan, np.nan)\n\n\ndef main(args):\n \"\"\"Main function to extract video information\"\"\"\n # Load data\n data = pd.read_csv(args.csv_path)\n if data.empty:\n data.to_csv(args.csv_save_path, index=False)\n print(f\"Input CSV is empty. Saved 0 samples to {args.csv_save_path}.\")\n return\n\n tasks = [(index, row[\"video_path\"], args.backend) for index, row in data.iterrows()]\n num_workers = args.num_workers if args.num_workers else os.cpu_count() or 1\n\n # Process videos with a per-video progress bar (more intuitive than per-worker)\n if args.disable_parallel or num_workers <= 1:\n ret = [\n get_video_info(task)\n for task in tqdm(tasks, total=len(tasks), desc=\"Processing videos\")\n ]\n else:\n with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:\n ret = list(\n tqdm(\n executor.map(get_video_info, tasks, chunksize=16),\n total=len(tasks),\n desc=\"Processing videos\",\n )\n )\n\n ret.sort(key=lambda x: x[0])\n (\n _idx_list,\n success_list,\n num_frames_list,\n height_list,\n width_list,\n aspect_ratio_list,\n hw_list,\n fps_list,\n ) = zip(*ret)\n\n # Add extracted information to DataFrame\n data[\"success\"] = success_list\n data[\"num_frames\"] = num_frames_list\n data[\"height\"] = height_list\n data[\"width\"] = width_list\n data[\"aspect_ratio\"] = aspect_ratio_list\n data[\"resolution\"] = hw_list\n data[\"fps\"] = fps_list\n\n # Filter existing files if requested\n if args.ext:\n assert \"video_path\" in data.columns\n data = data[data[\"video_path\"].apply(os.path.exists)]\n\n # Sort by frame count\n if \"num_frames\" in data.columns:\n data = data.sort_values(by=\"num_frames\", ascending=True)\n\n data.to_csv(args.csv_save_path, index=False)\n print(f\"Saved {len(data)} samples to {args.csv_save_path}.\")\n\n\ndef parse_args():\n \"\"\"Parse command line arguments\"\"\"\n parser = argparse.ArgumentParser(\n description=\"Extract video information using multiple backends\"\n )\n parser.add_argument(\n \"--csv_path\", type=str, required=True, help=\"Path to input CSV file\"\n )\n parser.add_argument(\n \"--csv_save_path\", type=str, default=None, help=\"Path to save output CSV file\"\n )\n parser.add_argument(\n \"--backend\",\n type=str,\n default=\"opencv\",\n help=\"Video backend\",\n choices=[\"opencv\", \"torchvision\", \"av\"],\n )\n parser.add_argument(\n \"--disable-parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=None, help=\"Number of parallel workers\"\n )\n parser.add_argument(\"--seed\", type=int, default=42, help=\"Random seed\")\n\n # File existence checking\n parser.add_argument(\"--ext\", action=\"store_true\", help=\"Check if video files exist\")\n\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n # Set random seeds for reproducibility\n if args.seed is not None:\n random.seed(args.seed)\n np.random.seed(args.seed)\n main(args)\n" }, { "path": "utils/get_instructions.py", "content": "\"\"\"\nThis module processes camera pose sequences and generates movement instructions.\n\"\"\"\n\nimport argparse\nimport numpy as np\nfrom scipy.spatial.transform import Rotation as R\nimport os\nimport pandas as pd\nfrom multiprocessing import Manager\nimport concurrent.futures\nimport queue\nfrom tqdm import tqdm\nimport json\n\n\ndef filter_poses(poses_array, alpha):\n \"\"\"\n Filter pose sequences using exponential moving average.\n - Position: Exponential moving average (EMA)\n - Orientation (quaternion): NLERP-based EMA with hemisphere flip handling\n\n Args:\n poses_array: Array of poses [position(3) + quaternion(4)]\n alpha: Smoothing factor (0 < alpha < 1)\n\n Returns:\n Filtered pose array with same shape as input\n \"\"\"\n\n positions = poses_array[:, :3]\n quaternions = poses_array[:, 3:]\n\n filtered_positions = np.zeros_like(positions)\n filtered_quaternions = np.zeros_like(quaternions)\n\n # Initialize with first frame\n filtered_positions[0] = positions[0]\n filtered_quaternions[0] = quaternions[0]\n\n for i in range(1, len(poses_array)):\n filtered_positions[i] = (\n alpha * positions[i] + (1 - alpha) * filtered_positions[i - 1]\n )\n\n # quaternion filtering with hemisphere check\n last_q = filtered_quaternions[i - 1]\n current_q = quaternions[i]\n\n # 1. Check hemisphere to ensure interpolation takes the \"shortest path\"\n if np.dot(last_q, current_q) < 0:\n current_q = -current_q\n\n # 2. Linear interpolation\n interp_q = (1 - alpha) * last_q + alpha * current_q\n\n # 3. Re-normalize to ensure unit quaternion\n filtered_quaternions[i] = interp_q / np.linalg.norm(interp_q)\n\n return np.hstack([filtered_positions, filtered_quaternions])\n\n\ndef poses_to_multi_instructions(poses_array, translation_thresh, rotation_thresh_deg):\n \"\"\"\n Convert camera pose sequence to concurrent movement instruction sequence.\n \"\"\"\n\n # Convert NumPy array to Scipy Rotation objects for easier computation\n poses = []\n for row in poses_array:\n pos = row[:3]\n rot = R.from_quat(row[3:])\n poses.append((pos, rot))\n\n command_sequence = []\n rotation_thresh_rad = np.deg2rad(rotation_thresh_deg)\n\n for i in range(len(poses) - 1):\n # Calculate local relative movement\n pos_t_w2c, rot_t_w2c = poses[i]\n pos_t1_w2c, rot_t1_w2c = poses[i+1]\n delta_rot = rot_t1_w2c * rot_t_w2c.inv()\n\n pos_t_c2w = -rot_t_w2c.inv().apply(pos_t_w2c)\n pos_t1_c2w = -rot_t1_w2c.inv().apply(pos_t1_w2c)\n local_delta_pos = rot_t_w2c.apply(pos_t1_c2w - pos_t_c2w)\n\n dx, dy, dz = local_delta_pos\n euler_angles_rad = delta_rot.as_euler(\n \"yxz\"\n ) # 'y' for yaw, 'x' for pitch, 'z' for roll\n yaw_change, pitch_change, roll_change = euler_angles_rad\n\n instructions = []\n\n # Translation movements\n if dz < -translation_thresh:\n instructions.append(\"Dolly Out\")\n elif dz > translation_thresh:\n instructions.append(\"Dolly In\")\n if dx > translation_thresh:\n instructions.append(\"Truck Right\")\n elif dx < -translation_thresh:\n instructions.append(\"Truck Left\")\n if dy > translation_thresh:\n instructions.append(\"Pedestal Down\")\n elif dy < -translation_thresh:\n instructions.append(\"Pedestal Up\")\n\n # Rotation movements\n if yaw_change > rotation_thresh_rad:\n instructions.append(\"Pan Left\")\n elif yaw_change < -rotation_thresh_rad:\n instructions.append(\"Pan Right\")\n\n if pitch_change > rotation_thresh_rad:\n instructions.append(\"Tilt Down\")\n elif pitch_change < -rotation_thresh_rad:\n instructions.append(\"Tilt Up\")\n\n if roll_change > rotation_thresh_rad:\n instructions.append(\"Roll CCW\")\n elif roll_change < -rotation_thresh_rad:\n instructions.append(\"Roll CW\")\n\n if not instructions:\n instructions.append(\"Stay\")\n\n command_sequence.append(instructions)\n\n return command_sequence\n\n\ndef process_single_row(args, row):\n \"\"\"Process a single video row to generate camera movement instructions.\"\"\"\n npy_path = os.path.join(args.dir_path, row[\"id\"], \"reconstructions\", \"poses.npy\")\n\n # Load and subsample poses, then apply filtering\n raw_poses = np.load(npy_path)[:: args.interval]\n filtered_poses = filter_poses(raw_poses, alpha=args.alpha)\n\n # Generate movement instructions\n instructions = poses_to_multi_instructions(\n filtered_poses, args.translation_threshold, args.rotation_threshold\n )\n\n json_file = os.path.join(args.dir_path, row[\"id\"], \"instructions.json\")\n if os.path.exists(json_file) and os.path.getsize(json_file) > 0:\n return\n\n # Merge consecutive identical instructions\n merged_instructions = {}\n start = 0\n prev_cmd = instructions[0]\n for i in range(1, len(instructions)):\n if instructions[i] == prev_cmd:\n continue\n else:\n key = f\"{start}->{i}\"\n merged_instructions[key] = prev_cmd\n start = i\n prev_cmd = instructions[i]\n # Add final segment\n key = f\"{start}->{len(instructions)}\"\n merged_instructions[key] = prev_cmd\n\n # Save instructions to JSON file\n with open(json_file, \"w\") as f:\n json.dump(merged_instructions, f, ensure_ascii=False, indent=2)\n\n\ndef worker(task_queue, args, pbar):\n \"\"\"Worker function for parallel processing of video rows.\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n\n process_single_row(args, row)\n\n task_queue.task_done()\n pbar.update(1)\n\n\ndef args_parser():\n \"\"\"Parse command line arguments.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"--csv_path\", type=str, default=\"outputs.csv\", help=\"Path to the input CSV file\"\n )\n parser.add_argument(\"--dir_path\", type=str, default=\"./outputs\")\n parser.add_argument(\n \"--interval\", type=int, default=2, help=\"Interval for computing instructions\"\n )\n parser.add_argument(\n \"--alpha\",\n type=float,\n default=0.1,\n help=\"Smoothing factor for filtering (0 < alpha < 1)\",\n )\n parser.add_argument(\n \"--translation_threshold\",\n type=float,\n default=0.02,\n help=\"Translation threshold for command generation\",\n )\n parser.add_argument(\n \"--rotation_threshold\",\n type=float,\n default=0.5,\n help=\"Rotation threshold for command generation\",\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=8, help=\"Number of parallel workers\"\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\ndef main():\n args = args_parser()\n csv = pd.read_csv(args.csv_path)\n\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(csv.iterrows(), total=len(csv)):\n process_single_row(args, row)\n else:\n # Parallel processing using ThreadPoolExecutor\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in csv.iterrows():\n task_queue.put((index, row))\n\n with tqdm(total=len(csv), desc=\"Finished tasks\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for _ in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, args, pbar))\n for future in concurrent.futures.as_completed(futures):\n future.result()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/get_instructions_enhanced.py", "content": "import argparse\nfrom math import sqrt\nimport numpy as np\nfrom scipy.spatial.transform import Rotation as R\nimport os\nimport pandas as pd\nfrom multiprocessing import Manager\nimport concurrent.futures\nimport queue\nfrom tqdm import tqdm\nimport json\nfrom collections import defaultdict, Counter\nimport itertools\n\n\ndef filter_poses(poses_array, alpha):\n \"\"\"\n Smooth pose sequence using Exponential Moving Average (EMA).\n - Position: Standard EMA\n - Quaternion: EMA with hemisphere check (shortest path interpolation)\n \"\"\"\n positions = poses_array[:, :3]\n quaternions = poses_array[:, 3:]\n\n filtered_pos = np.zeros_like(positions)\n filtered_quat = np.zeros_like(quaternions)\n\n # Initialize with first frame\n filtered_pos[0], filtered_quat[0] = positions[0], quaternions[0]\n\n for i in range(1, len(poses_array)):\n # Position smoothing\n filtered_pos[i] = alpha * positions[i] + \\\n (1 - alpha) * filtered_pos[i-1]\n\n # Quaternion smoothing with hemisphere correction\n last_q, curr_q = filtered_quat[i-1], quaternions[i]\n if np.dot(last_q, curr_q) < 0: # Flip to shortest interpolation path\n curr_q = -curr_q\n\n interp_q = (1 - alpha) * last_q + alpha * curr_q\n filtered_quat[i] = interp_q / \\\n np.linalg.norm(interp_q) # Keep unit quaternion\n\n return np.hstack([filtered_pos, filtered_quat])\n\n\ndef poses_to_multi_instructions(poses_array, translation_thresh, rotation_thresh_deg, interval=1):\n \"\"\"\n Convert pose sequence to motion instructions (e.g., Dolly, Pan).\n Calculates pose difference between frame i and i+interval (convolution-like).\n \"\"\"\n # Convert to (position, Rotation object) pairs\n poses = [(row[:3], R.from_quat(row[3:])) for row in poses_array]\n command_seq = []\n\n # Adjust thresholds by interval (scaling for longer gaps)\n rotation_thresh_deg *= sqrt(interval) / 1.8\n rotation_thresh_rad = np.deg2rad(rotation_thresh_deg)\n translation_thresh *= sqrt(interval)\n stride = int(sqrt(interval) + 1)\n\n i = 0\n while True:\n if i + interval >= len(poses): # Ensure valid frame pair\n break\n\n # Calculate relative motion (local coordinate system)\n pos_t_w2c, rot_t_w2c = poses[i]\n pos_t1_w2c, rot_t1_w2c = poses[i+interval]\n delta_rot = rot_t1_w2c * rot_t_w2c.inv()\n\n pos_t_c2w = -rot_t_w2c.inv().apply(pos_t_w2c)\n pos_t1_c2w = -rot_t1_w2c.inv().apply(pos_t1_w2c)\n local_delta_pos = rot_t_w2c.apply(pos_t1_c2w - pos_t_c2w)\n\n dx, dy, dz = local_delta_pos\n yaw, pitch, roll = delta_rot.as_euler(\n \"yxz\") # Yaw:Pan, Pitch:Tilt, Roll:Rotate\n\n instructions = []\n # Translation commands\n if dz < -translation_thresh:\n instructions.append(\"Dolly Out\")\n elif dz > translation_thresh:\n instructions.append(\"Dolly In\")\n if dx > translation_thresh:\n instructions.append(\"Truck Right\")\n elif dx < -translation_thresh:\n instructions.append(\"Truck Left\")\n if dy > translation_thresh:\n instructions.append(\"Pedestal Down\")\n elif dy < -translation_thresh:\n instructions.append(\"Pedestal Up\")\n\n # Rotation commands\n if yaw > rotation_thresh_rad:\n instructions.append(\"Pan Left\")\n elif yaw < -rotation_thresh_rad:\n instructions.append(\"Pan Right\")\n if pitch > rotation_thresh_rad:\n instructions.append(\"Tilt Down\")\n elif pitch < -rotation_thresh_rad:\n instructions.append(\"Tilt Up\")\n if roll > rotation_thresh_rad:\n instructions.append(\"Roll CCW\")\n elif roll < -rotation_thresh_rad:\n instructions.append(\"Roll CW\")\n\n command_seq.append(instructions if instructions else [\"Stay\"])\n i += stride\n return command_seq\n\n\ndef calculate_relative_scale(total_distance, num_poses, f_translation, min_threshold=0.001):\n \"\"\"\n Calculate relative translation threshold (dynamic scaling by total motion).\n \"\"\"\n if num_poses <= 1:\n return min_threshold\n base_scale = total_distance / num_poses # Base scale per frame\n return max(base_scale / f_translation, min_threshold)\n\n\ndef voter(args, row, interval, alpha):\n \"\"\"\n Process single video with specific (interval, alpha) parameter pair.\n \"\"\"\n # Locate pose file\n npy_path = os.path.join(\n args.dir_path, row[\"id\"], \"reconstructions\", \"poses.npy\"\n )\n\n try:\n raw_poses = np.load(npy_path)\n filtered_poses = filter_poses(raw_poses, alpha)\n\n # Calculate dynamic thresholds\n translation_thresh = calculate_relative_scale(\n row[\"moveDist\"], len(\n filtered_poses), args.f_translation, args.min_threshold_translation\n )\n rotation_thresh = args.rotation_threshold\n\n return poses_to_multi_instructions(\n filtered_poses, translation_thresh, rotation_thresh, interval\n )\n except Exception as e:\n print(f\"Error processing {row['id']}: {e}\")\n return None\n\n\ndef collect_all_results(args, row, param_combinations):\n \"\"\"Collect instruction results for all (interval, alpha) pairs.\"\"\"\n results = []\n for interval, alpha in param_combinations:\n res = voter(args, row, interval, alpha)\n if res is not None:\n results.append(res)\n return results\n\n\n# ------------------------------ Voting Logic ------------------------------\ndef get_mutually_exclusive_groups():\n \"\"\"Return groups of conflicting instructions (cannot coexist).\"\"\"\n return [\n [\"Dolly In\", \"Dolly Out\"], [\"Truck Left\", \"Truck Right\"],\n [\"Pedestal Up\", \"Pedestal Down\"], [\"Pan Left\", \"Pan Right\"],\n [\"Tilt Up\", \"Tilt Down\"], [\"Roll CW\", \"Roll CCW\"]\n ]\n\n\ndef remove_conflicting_instructions(instructions, conflict_groups):\n \"\"\"Remove conflicting instructions (keep higher-voted ones).\"\"\"\n selected = []\n selected_set = set()\n for inst, count in instructions:\n conflict = False\n for group in conflict_groups:\n if inst in group and any(s in group for s in selected_set):\n conflict = True\n break\n if not conflict:\n selected.append((inst, count))\n selected_set.add(inst)\n return selected\n\n\ndef smart_instruction_selection(non_conflicting_inst):\n \"\"\"\n Smart instruction selection based on vote distribution:\n - Keep leading votes (3x threshold for断层)\n - Max 4 instructions\n - Prioritize non-\"Stay\" commands\n \"\"\"\n if not non_conflicting_inst:\n return [\"Stay\"]\n if len(non_conflicting_inst) == 1:\n return [non_conflicting_inst[0][0]]\n\n # Separate Stay and other instructions\n stay = [i for i in non_conflicting_inst if i[0] == \"Stay\"]\n others = [i for i in non_conflicting_inst if i[0] != \"Stay\"]\n if not others:\n return [\"Stay\"]\n\n votes = [c for _, c in others]\n max_vote = votes[0]\n selected = []\n\n # Check for vote gap (3x threshold)\n if len(others) >= 2 and max_vote >= votes[1] * 3:\n selected = [i[0] for i in others if i[1] == max_vote]\n else:\n # Select up to 4 leading instructions\n gap_thresh = max_vote * 0.5\n selected = [i[0] for i in others if i[1] >= gap_thresh][:4]\n\n # Ensure minimum 2 instructions if no large gap\n if len(selected) < 2 and len(others) >= 2 and max_vote < votes[1] * 3:\n selected = [i[0] for i in others[:2]]\n\n return selected if selected else [\"Stay\"]\n\n\ndef collect_interval_based_votes(all_results, param_combinations):\n \"\"\"\n Vote by time interval: collect all instructions covering (start_frame->end_frame).\n Handles overlapping segments from different (interval, alpha) pairs.\n \"\"\"\n if not all_results:\n return {}\n\n # Get max frame covered by any parameter pair\n max_frames = 0\n for index, res in enumerate(all_results):\n interval = param_combinations[index][0]\n stride = int(sqrt(interval) + 1)\n if res:\n last_start = (len(res)-1) * stride\n max_frames = max(max_frames, last_start + interval)\n\n interval_votes = {}\n for start in range(max_frames):\n end = start + 1\n vote_counter = Counter()\n # Check all parameter results for coverage of (start->end)\n for res_index, res in enumerate(all_results):\n interval, _ = param_combinations[res_index]\n stride = int(sqrt(interval) + 1)\n for seg_index, seg in enumerate(res):\n seg_start = seg_index * stride\n seg_end = seg_start + interval\n # Check if segment covers target interval\n if seg_start <= start < seg_end and seg_start < end <= seg_end:\n for inst in seg:\n vote_counter[inst] += 1\n interval_votes[f\"{start}->{end}\"] = vote_counter\n return interval_votes\n\n\ndef vote_for_final_instructions(all_results, param_combinations=None):\n \"\"\"Generate final instructions via voting (interval-based if possible).\"\"\"\n if not all_results:\n return []\n\n conflict_groups = get_mutually_exclusive_groups()\n final_seq = []\n\n # Use interval-based voting if parameters are provided\n if param_combinations and len(param_combinations) == len(all_results):\n interval_votes = collect_interval_based_votes(\n all_results, param_combinations)\n for key in sorted(interval_votes.keys(), key=lambda x: int(x.split('->')[0])):\n votes = interval_votes[key]\n if votes:\n sorted_inst = votes.most_common()\n non_conflict = remove_conflicting_instructions(\n sorted_inst, conflict_groups)\n selected = smart_instruction_selection(non_conflict)\n else:\n selected = [\"Stay\"]\n final_seq.append(selected)\n else:\n # Fallback: frame-wise voting\n max_len = max(len(res) for res in all_results)\n for frame_index in range(max_len):\n votes = Counter()\n for res in all_results:\n if frame_index < len(res):\n for inst in res[frame_index]:\n votes[inst] += 1\n if votes:\n sorted_inst = votes.most_common()\n non_conflict = remove_conflicting_instructions(\n sorted_inst, conflict_groups)\n selected = smart_instruction_selection(non_conflict)\n else:\n selected = [\"Stay\"]\n final_seq.append(selected)\n return final_seq\n\n\n# ------------------------------ Main Workflow ------------------------------\ndef merge_consecutive_instructions(instructions):\n \"\"\"Merge consecutive identical instruction lists (e.g., [A,A,A] → \"0->3\":[A]).\"\"\"\n if not instructions:\n return {}\n merged = {}\n start, prev = 0, instructions[0]\n for i in range(1, len(instructions)):\n if instructions[i] != prev:\n merged[f\"{start}->{i}\"] = prev\n start, prev = i, instructions[i]\n merged[f\"{start}->{len(instructions)}\"] = prev # Add final segment\n return merged\n\n\ndef process_single_row(args, row, param_combinations):\n # Skip if output exists\n out_file = os.path.join(args.dir_path, row['id'], \"instructions.json\")\n if os.path.exists(out_file) and os.path.getsize(out_file) > 0:\n return\n\n # Collect results & vote\n all_results = collect_all_results(args, row, param_combinations)\n if not all_results:\n print(f\"No valid results for {row['id']}\")\n return\n final_inst = vote_for_final_instructions(all_results, param_combinations)\n merged_inst = merge_consecutive_instructions(final_inst)\n\n # Save to JSON\n with open(out_file, \"w\") as f:\n json.dump(merged_inst, f, ensure_ascii=False, indent=2)\n\n\ndef generate_param_combinations(args):\n \"\"\"Generate all (interval, alpha) parameter pairs for grid search.\"\"\"\n intervals = getattr(args, \"intervals\", [1, 3, 5])\n alphas = getattr(args, \"alphas\", [0.03, 0.05, 0.1])\n return list(itertools.product(intervals, alphas))\n\n\ndef worker(task_queue, args, param_combinations, pbar):\n \"\"\"Parallel worker: process tasks from queue.\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n process_single_row(args, row, param_combinations)\n except queue.Empty:\n break\n task_queue.task_done()\n pbar.update(1)\n\n\ndef args_parser():\n parser = argparse.ArgumentParser(\n description=\"Enhanced Camera Pose Instruction Generator\")\n parser.add_argument(\"--csv_path\", type=str, required=True,\n help=\"Input CSV path (The final_results.csv generated by evaluation.py)\")\n parser.add_argument(\"--dir_path\", type=str, required=True,\n help=\"Annotation directory path\")\n parser.add_argument(\"--intervals\", type=int, nargs=\"+\",\n default=[1, 3, 5], help=\"Frame intervals for grid search\")\n parser.add_argument(\"--alphas\", type=float, nargs=\"+\",\n default=[0.03, 0.05, 0.1], help=\"Smoothing factors for grid search\")\n parser.add_argument(\"--f_translation\", type=float,\n default=1.1, help=\"Translation scale factor (>1)\")\n parser.add_argument(\"--min_threshold_translation\", type=float,\n default=0.01, help=\"Min translation threshold\")\n parser.add_argument(\"--rotation_threshold\", type=float,\n default=1.5, help=\"Fixed rotation threshold (degrees)\")\n parser.add_argument(\"--num_workers\", type=int,\n default=8, help=\"Parallel workers count\")\n parser.add_argument(\"--disable_parallel\", action=\"store_true\",\n help=\"Disable parallel processing\")\n return parser.parse_args()\n\n\ndef main():\n args = args_parser()\n csv = pd.read_csv(args.csv_path)\n\n param_combinations = generate_param_combinations(args)\n\n if args.disable_parallel:\n # Serial processing\n for index, row in tqdm(csv.iterrows(), total=len(csv), desc=\"Processing\"):\n process_single_row(args, row, param_combinations)\n else:\n # Parallel processing\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in csv.iterrows():\n task_queue.put((index, row))\n\n with tqdm(total=len(csv), desc=\"Processing\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(max_workers=args.num_workers) as executor:\n for _ in range(args.num_workers):\n executor.submit(worker, task_queue, args,\n param_combinations, pbar)\n task_queue.join()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/merge_tables.py", "content": "\"\"\"\nCSV table merging utility for combining multiple clip information files.\n\"\"\"\n\nimport os\nimport glob\nimport argparse\nimport pandas as pd\n\n\ndef read_csv_file(file_path):\n \"\"\"Read a single CSV file\"\"\"\n return pd.read_csv(file_path)\n\n\ndef merge_tables_from_files(file_list, output_file, merge_on=None):\n \"\"\"\n Merge multiple CSV files using common columns as merge keys.\n\n Args:\n file_list: List of CSV file paths to merge\n output_file: Output path for merged CSV file\n merge_on: List of column names for merging (defaults to first 13 columns)\n \"\"\"\n if not file_list:\n raise ValueError(\"File list is empty!\")\n\n # Read all CSV files\n dfs = [read_csv_file(f) for f in file_list]\n\n # Auto-select merge keys: first 13 columns\n if merge_on is None:\n merge_on = dfs[0].columns[:13].tolist()\n\n # Merge dataframes\n df_merged = dfs[0]\n for df in dfs[1:]:\n # Check if merge keys are consistent\n if merge_on != df.columns[:13].tolist():\n raise ValueError(\n f\"Common columns in one file are inconsistent with previous files!\"\n )\n # Merge based on specified keys\n df_merged = pd.merge(df_merged, df, on=merge_on)\n\n # Save merged result\n df_merged.to_csv(output_file, index=False)\n print(f\"Merge completed. Saved to {output_file}\")\n return df_merged\n\n\ndef main():\n parser = argparse.ArgumentParser(\n description=\"Merge multiple CSV files from a folder\"\n )\n parser.add_argument(\"--csv_dir\", type=str, help=\"Path to folder containing CSV files\")\n parser.add_argument(\n \"--output\", type=str, required=True, help=\"Output path for merged CSV file\"\n )\n\n args = parser.parse_args()\n\n # Match CSV files with 'clips_info_' prefix\n pattern = os.path.join(args.csv_dir, \"clips_info_*.csv\")\n file_list = glob.glob(pattern)\n file_list.sort() # Sort to ensure consistent merge order\n\n if not file_list:\n raise ValueError(f\"No matching CSV files found in folder {args.csv_dir}!\")\n\n print(f\"Found {len(file_list)} CSV files:\")\n for f in file_list:\n print(f\" {f}\")\n\n # Perform merge\n merge_tables_from_files(file_list, args.output)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/normalize_intrinsics.py", "content": "\"\"\"\nCamera intrinsics normalization utility.\n\nThis module provides functionality for:\n- Normalizing camera intrinsics to standard format\n- Converting focal length to normalized coordinates\n- Parallel processing of multiple camera files\n- Support for both threaded and sequential processing\n\"\"\"\n\nimport numpy as np\nimport os\nimport pandas as pd\nimport argparse\nimport concurrent.futures\nimport multiprocessing as mp\nfrom multiprocessing import Manager\nimport queue\nfrom tqdm import tqdm\n\n\ndef possess_single_row(row, args):\n \"\"\"\n Process a single row to normalize camera intrinsics.\n \"\"\"\n id = row[\"id\"]\n dir_path = os.path.join(args.dir_path, id, \"reconstructions\")\n cam_intrinsics_file = os.path.join(dir_path, \"intrinsics.npy\")\n \n # Load and normalize intrinsics\n intrinsics = np.load(cam_intrinsics_file)\n intrinsics[:, 0] /= intrinsics[:, 2] * 2 # Normalize focal length x\n intrinsics[:, 1] /= intrinsics[:, 3] * 2 # Normalize focal length y\n intrinsics[:, 2] = 0.5 # Set principal point x to center\n intrinsics[:, 3] = 0.5 # Set principal point y to center\n \n # Save normalized intrinsics\n np.save(cam_intrinsics_file, intrinsics)\n\n\ndef worker(task_queue, args, pbar):\n \"\"\"\n Worker function for parallel processing of intrinsics normalization.\n \"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n possess_single_row(row, args)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for intrinsics normalization.\"\"\"\n parser = argparse.ArgumentParser(description=\"Normalize camera intrinsics to standard format\")\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\"--dir_path\", type=str, default=\"./outputs\")\n parser.add_argument(\n \"--num_workers\",\n type=int,\n default=8,\n help=\"Number of workers for parallel processing\",\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n\n df = pd.read_csv(args.csv_path)\n\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(df.iterrows(), total=len(df)):\n possess_single_row(row, index, args)\n else:\n # Parallel processing using thread pool\n manager = Manager()\n task_queue = manager.Queue()\n \n # Add all tasks to queue\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n with tqdm(total=len(df), desc=\"Finished tasks\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for _ in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, args, pbar))\n for future in concurrent.futures.as_completed(futures):\n future.result()\n" }, { "path": "utils/pack_clip_assets.py", "content": "\"\"\"\npack_clip_assets.py\n------------------\nThis script unifies depth, RGB frames, intrinsics, extrinsics, etc. of a specified video clip into a single npz file for downstream 3D reconstruction or analysis.\n\nUsage example:\n python pack_clip_assets.py --base_dir /path/to/HQ --clip_id group_xxxx/xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx --height 328 --width 584\n\n\"\"\"\n\nimport argparse\nimport numpy as np\nimport torch\nfrom lietorch import SE3\nimport cv2\nfrom read_depth import read_depth\n\ndef load_video(clip_path, indexes_path, height=720, width=1280):\n \"\"\"\n Read video frames at specified indexes and resize to (height, width).\n Args:\n clip_path (str): Path to video file\n indexes_path (str): Path to frame indexes txt\n height (int): Output frame height\n width (int): Output frame width\n Returns:\n np.ndarray: (N, height, width, 3) RGB frames\n \"\"\"\n indexes = []\n with open(indexes_path, 'r') as f:\n for line in f:\n parts = line.strip().split()\n if len(parts) == 2:\n indexes.append(int(parts[1]))\n print(f\"Frame indexes: {indexes}\")\n cap = cv2.VideoCapture(clip_path)\n frames = []\n for idx in indexes:\n cap.set(cv2.CAP_PROP_POS_FRAMES, idx)\n ret, frame = cap.read()\n if not ret:\n raise ValueError(f\"Frame at index {idx} could not be read.\")\n frame = cv2.resize(frame, (width, height))\n frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n frames.append(frame)\n cap.release()\n return np.array(frames)\n\ndef load_intrinsics(intrinsics_path, tgt_width=1024, tgt_height=576):\n \"\"\"\n Read normalized intrinsics (n,4), convert to 3x3 matrix and scale to target resolution.\n Args:\n intrinsics_path (str): Path to intrinsics npy\n tgt_width (int): Target width\n tgt_height (int): Target height\n Returns:\n np.ndarray: (N, 3, 3) intrinsics matrices\n \"\"\"\n intrinsics = np.load(intrinsics_path)\n intrinsics_3x3 = []\n for intrin in intrinsics:\n fx, fy, cx, cy = intrin\n K = np.array([[fx, 0, cx],\n [0, fy, cy],\n [0, 0, 1]], dtype=np.float32)\n intrinsics_3x3.append(K)\n intrinsics_3x3 = np.array(intrinsics_3x3)\n intrinsics_3x3[:, 0, 0] *= tgt_width\n intrinsics_3x3[:, 1, 1] *= tgt_height\n intrinsics_3x3[:, 0, 2] *= tgt_width\n intrinsics_3x3[:, 1, 2] *= tgt_height\n return intrinsics_3x3\n\ndef main():\n \"\"\"\n Main pipeline: load depth, RGB frames, intrinsics, extrinsics, and save as npz.\n \"\"\"\n parser = argparse.ArgumentParser(description=\"Pack clip assets into a single npz file.\")\n parser.add_argument('--base_dir', type=str, required=True, help='Root directory of HQ data')\n parser.add_argument('--group_id', type=int, required=False, help='Group ID, e.g. group_xxxx')\n parser.add_argument('--clip_id', type=str, required=True, help='Clip ID, e.g. xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx')\n parser.add_argument('--height', type=int, default=328, help='Output image height')\n parser.add_argument('--width', type=int, default=584, help='Output image width')\n parser.add_argument('--output', type=str, default='sgd_cvd_hr.npz', help='Output npz filename')\n args = parser.parse_args()\n\n # Path construction\n annotation_dir = f'{args.base_dir}/annotations/group_{args.group_id:04d}/{args.clip_id}'\n depth_path = f'{args.base_dir}/depths/group_{args.group_id:04d}/{args.clip_id}.zip'\n clip_path = f'{args.base_dir}/videos/group_{args.group_id:04d}/{args.clip_id}.mp4'\n intrinsics_path = f'{annotation_dir}/intrinsics.npy'\n extrinsics_path = f'{annotation_dir}/poses.npy'\n indexes_path = f'{annotation_dir}/indexes.txt'\n\n # Load intrinsics and extrinsics\n intrinsics = load_intrinsics(intrinsics_path, tgt_width=args.width, tgt_height=args.height)\n extrinsics = np.load(extrinsics_path)\n\n # Load and resize depth\n depth = np.clip(read_depth(depth_path), 1e-3, 1e2) # (N, H, W)\n resized_depth = np.zeros((depth.shape[0], args.height, args.width), dtype=depth.dtype)\n for i in range(depth.shape[0]):\n resized_depth[i] = cv2.resize(depth[i], (args.width, args.height), interpolation=cv2.INTER_LINEAR)\n\n # Load RGB frames\n frames = load_video(clip_path, indexes_path, args.height, args.width)\n\n # Compute camera poses\n poses_th = torch.as_tensor(extrinsics, device=\"cpu\").float()\n cam_c2w = SE3(poses_th).inv().matrix()\n K = intrinsics[0]\n K_o = torch.from_numpy(K).float()\n\n # Save as npz\n np.savez(\n args.output,\n images=frames,\n depths=resized_depth,\n intrinsic=K_o.detach().cpu().numpy(),\n cam_c2w=cam_c2w.detach().cpu().numpy(),\n )\n print(f\"Saved to {args.output}\")\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "utils/quat_to_mat.py", "content": "\"\"\"\nCamera pose conversion utility to camera-to-world (c2w) or world-to-camera (w2c) format.\nConverts quaternion representations to rotation matrices and handles pose transformations.\n\nThis module provides utilities for:\n- Converting between quaternion and matrix representations of camera poses\n- Transforming between world-to-camera (w2c) and camera-to-world (c2w) coordinate systems\n- Parallel processing of pose conversion for large datasets\n\"\"\"\n\nimport einops\nimport torch\nimport torch.nn.functional as F\nimport numpy as np\nimport os\nimport pandas as pd\nimport argparse\nimport concurrent.futures\nimport multiprocessing as mp\nfrom multiprocessing import Manager\nimport queue\nfrom tqdm import tqdm\n\n\nclass Pose:\n \"\"\"\n A class of operations on camera poses (numpy arrays with shape [...,3,4]).\n Each [3,4] camera pose takes the form of [R|t].\n \"\"\"\n\n def __call__(self, R=None, t=None):\n \"\"\"\n Construct a camera pose from the given rotation matrix R and/or translation vector t.\n\n Args:\n R: Rotation matrix [...,3,3] or None\n t: Translation vector [...,3] or None\n\n Returns:\n pose: Camera pose matrix [...,3,4]\n \"\"\"\n assert R is not None or t is not None\n if R is None:\n if not isinstance(t, np.ndarray):\n t = np.array(t)\n R = np.eye(3, device=t.device).repeat(*t.shape[:-1], 1, 1)\n elif t is None:\n if not isinstance(R, np.ndarray):\n R = np.array(R)\n t = np.zeros(R.shape[:-1], device=R.device)\n else:\n if not isinstance(R, np.ndarray):\n R = np.array(R)\n if not isinstance(t, np.ndarray):\n t = np.tensor(t)\n assert R.shape[:-1] == t.shape and R.shape[-2:] == (3, 3)\n R = R.astype(np.float32)\n t = t.astype(np.float32)\n pose = np.concatenate([R, t[..., None]], axis=-1) # [...,3,4]\n assert pose.shape[-2:] == (3, 4)\n return pose\n\n def invert(self, pose, use_inverse=False): # c2w <==> w2c\n \"\"\"\n Invert a camera pose transformation matrix.\n Converts between camera-to-world (c2w) and world-to-camera (w2c) representations.\n For a pose [R|t], the inverse is [R^T | -R^T*t].\n\n Args:\n pose: Camera pose matrix [...,3,4] with shape [R|t]\n use_inverse: Whether to use matrix inverse instead of transpose for rotation\n\n Returns:\n pose_inv: Inverted camera pose matrix [...,3,4]\n \"\"\"\n R, t = pose[..., :3], pose[..., 3:]\n R_inv = (\n R.inverse() if use_inverse else R.transpose(0, 2, 1)\n ) # For orthogonal matrices, transpose equals inverse\n t_inv = (-R_inv @ t)[..., 0] # Apply inverse rotation to negative translation\n pose_inv = self(R=R_inv, t=t_inv)\n return pose_inv\n\n def compose(self, pose_list):\n \"\"\"\n Compose a sequence of poses together.\n pose_new(x) = poseN o ... o pose2 o pose1(x)\n\n Args:\n pose_list: List of camera poses to compose\n\n Returns:\n pose_new: Composed camera pose\n \"\"\"\n pose_new = pose_list[0]\n for pose in pose_list[1:]:\n pose_new = self.compose_pair(pose_new, pose)\n return pose_new\n\n def compose_pair(self, pose_a, pose_b):\n \"\"\"\n Compose two poses together.\n pose_new(x) = pose_b o pose_a(x)\n\n Args:\n pose_a: First camera pose\n pose_b: Second camera pose\n\n Returns:\n pose_new: Composed camera pose\n \"\"\"\n R_a, t_a = pose_a[..., :3], pose_a[..., 3:]\n R_b, t_b = pose_b[..., :3], pose_b[..., 3:]\n R_new = R_b @ R_a\n t_new = (R_b @ t_a + t_b)[..., 0]\n pose_new = self(R=R_new, t=t_new)\n return pose_new\n\n def scale_center(self, pose, scale):\n \"\"\"\n Scale the camera center from the origin.\n 0 = R@c+t --> c = -R^T@t (camera center in world coordinates)\n 0 = R@(sc)+t' --> t' = -R@(sc) = -R@(-R^T@st) = st\n\n Args:\n pose: Camera pose to scale\n scale: Scale factor\n\n Returns:\n pose_new: Scaled camera pose\n \"\"\"\n R, t = pose[..., :3], pose[..., 3:]\n pose_new = np.concatenate([R, t * scale], axis=-1)\n return pose_new\n\n\ndef quaternion_to_matrix(quaternions, eps: float = 1e-8):\n \"\"\"\n Convert 4-dimensional quaternions to 3x3 rotation matrices.\n This is adapted from Pytorch3D:\n https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py\n\n Args:\n quaternions: Quaternion tensor [..., 4] (order: i, j, k, r)\n eps: Small value for numerical stability\n\n Returns:\n Rotation matrices [..., 3, 3]\n \"\"\"\n\n # Order changed to match scipy format!\n i, j, k, r = torch.unbind(quaternions, dim=-1)\n two_s = 2 / ((quaternions * quaternions).sum(dim=-1) + eps)\n\n o = torch.stack(\n (\n 1 - two_s * (j * j + k * k),\n two_s * (i * j - k * r),\n two_s * (i * k + j * r),\n two_s * (i * j + k * r),\n 1 - two_s * (i * i + k * k),\n two_s * (j * k - i * r),\n two_s * (i * k - j * r),\n two_s * (j * k + i * r),\n 1 - two_s * (i * i + j * j),\n ),\n -1,\n )\n return einops.rearrange(o, \"... (i j) -> ... i j\", i=3, j=3)\n\n\ndef pose_from_quaternion(pose):\n \"\"\"\n Convert pose from quaternion representation to transformation matrix.\n\n Args:\n pose: Pose tensor [..., 7] where first 3 elements are translation (t)\n and last 4 elements are quaternion rotation (r)\n\n Returns:\n w2c_matrix: World-to-camera transformation matrices [..., 3, 4]\n \"\"\"\n # Input is w2c, pose(n,7) or (n,v,7), output is (N,3,4) w2c matrix\n # Tensor format from https://github.com/pointrix-project/Geomotion/blob/6ab0c364f1b44ab4ea190085dbf068f62b42727c/geomotion/model/cameras.py#L6\n if type(pose) == np.ndarray:\n pose = torch.tensor(pose)\n if len(pose.shape) == 1:\n pose = pose[None]\n quat_t = pose[..., :3] # Translation\n quat_r = pose[..., 3:] # Quaternion rotation\n w2c_matrix = torch.zeros((*list(pose.shape)[:-1], 3, 4), device=pose.device)\n w2c_matrix[..., :3, 3] = quat_t\n w2c_matrix[..., :3, :3] = quaternion_to_matrix(quat_r)\n return w2c_matrix\n\n\ndef possess_single_row(row, index, args):\n \"\"\"\n Process a single row to convert camera poses to c2w/w2c format.\n\n Args:\n row: Data row containing video ID\n index: Row index\n args: Command line arguments\n \"\"\"\n id = row[\"id\"]\n dir_path = os.path.join(args.dir_path, id, \"reconstructions\")\n cam_pos_file = os.path.join(dir_path, \"poses.npy\")\n if not os.path.exists(cam_pos_file):\n return\n output_file = os.path.join(dir_path, \"extrinsics.npy\")\n if os.path.exists(output_file):\n return\n\n # Load quaternion poses\n pose = np.load(cam_pos_file)\n # Convert w2c quaternion format (N,v,7) to w2c matrix format (N,v,3,4)\n poses = pose_from_quaternion(pose)\n poses = poses.cpu().numpy()\n # Convert w2c matrices to c2w matrices (N,v,3,4)\n if args.format == \"c2w\":\n poses = Pose().invert(poses)\n np.save(output_file, poses)\n\n\ndef worker(task_queue, args, pbar):\n \"\"\"Worker function for parallel pose conversion processing.\"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n possess_single_row(row, index, args)\n task_queue.task_done()\n pbar.update(1)\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for camera pose conversion.\"\"\"\n parser = argparse.ArgumentParser(description=\"Convert quaternion to camera pose\")\n parser.add_argument(\"--csv_path\", type=str, help=\"Path to the csv file\")\n parser.add_argument(\"--dir_path\", type=str, default=\"./outputs\")\n parser.add_argument(\"--format\", type=str, default=\"c2w\", choices=[\"c2w\", \"w2c\"])\n parser.add_argument(\n \"--num_workers\",\n type=int,\n default=8,\n help=\"Number of workers for parallel processing\",\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n return parser.parse_args()\n\n\nif __name__ == \"__main__\":\n args = parse_args()\n\n df = pd.read_csv(args.csv_path)\n\n if args.disable_parallel:\n # Sequential processing\n for index, row in tqdm(df.iterrows(), total=len(df)):\n possess_single_row(row, index, args)\n else:\n # Parallel processing with multiple workers\n manager = Manager()\n task_queue = manager.Queue()\n for index, row in df.iterrows():\n task_queue.put((index, row))\n\n with tqdm(total=len(df), desc=\"Finished tasks\") as pbar:\n with concurrent.futures.ThreadPoolExecutor(\n max_workers=args.num_workers\n ) as executor:\n futures = []\n for _ in range(args.num_workers):\n futures.append(executor.submit(worker, task_queue, args, pbar))\n for future in concurrent.futures.as_completed(futures):\n future.result()\n" }, { "path": "utils/read_depth.py", "content": "import zipfile\nimport numpy as np\nimport OpenEXR\n\n\ndef read_depth(zip_file_path):\n \"\"\"\n Read depth from zipped exr files.\n \"\"\"\n valid_width, valid_height = 0, 0\n depth_data_list = []\n with zipfile.ZipFile(zip_file_path, \"r\") as z:\n for file_name in sorted(z.namelist()):\n with z.open(file_name) as f:\n try:\n exr = OpenEXR.InputFile(f)\n except OSError:\n # Sometimes EXR loader might fail, we return all nan maps.\n assert valid_width > 0 and valid_height > 0\n depth_data_list.append(\n np.full((valid_height, valid_width), np.nan, dtype=np.float32))\n continue\n header = exr.header()\n dw = header[\"dataWindow\"]\n valid_width = width = dw.max.x - dw.min.x + 1\n valid_height = height = dw.max.y - dw.min.y + 1\n channels = exr.channels([\"Z\"])\n depth_data = np.frombuffer(\n channels[0], dtype=np.float16).reshape((height, width))\n depth_data_list.append(depth_data.astype(np.float32))\n # Note that the depth with a negative value is an invalid depth.\n # It can be set to the farthest point or other operations.\n depth_array = np.array(depth_data_list)\n depth_array_safe = np.where(depth_array == 0, 1e-12, depth_array)\n return 1.0 / depth_array_safe\n" }, { "path": "utils/read_video.py", "content": "\"\"\"\nVideo reading utilities with memory optimization and multiple backend support.\n\"\"\"\n\nimport gc\nimport math\nimport os\nimport re\nimport warnings\nfrom fractions import Fraction\nfrom typing import Any, Dict, List, Optional, Tuple, Union\nfrom tools.logger import test_lg\nimport av\nimport cv2\nimport numpy as np\nimport torch\nfrom torchvision import get_video_backend\nfrom torchvision.io.video import _check_av_available\n\nMAX_NUM_FRAMES = 2500\n\n\ndef read_video_av(\n filename: str,\n start_pts: Union[float, Fraction] = 0,\n end_pts: Optional[Union[float, Fraction]] = None,\n pts_unit: str = \"pts\",\n output_format: str = \"THWC\",\n) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]:\n \"\"\"\n Read video frames using PyAV backend with memory optimization.\n \n Modified from torchvision.io.video.read_video with improvements:\n - No audio extraction (returns empty aframes)\n - PyAV backend only\n - Added container.close() and gc.collect() to prevent memory leaks\n - Optimized for memory efficiency\n \"\"\"\n # Validate format\n output_format = output_format.upper()\n if output_format not in (\"THWC\", \"TCHW\"):\n raise ValueError(f\"output_format should be either 'THWC' or 'TCHW', got {output_format}.\")\n \n # Check file existence\n if not os.path.exists(filename):\n raise RuntimeError(f\"File not found: {filename}\")\n \n # Validate backend\n assert get_video_backend() == \"pyav\", \"pyav backend is required for read_video_av\"\n _check_av_available()\n \n # Validate time range\n if end_pts is None:\n end_pts = float(\"inf\")\n if end_pts < start_pts:\n raise ValueError(f\"end_pts should be larger than start_pts, got start_pts={start_pts} and end_pts={end_pts}\")\n\n # Extract video metadata\n info = {}\n container = av.open(filename, metadata_errors=\"ignore\")\n video_fps = container.streams.video[0].average_rate\n if video_fps is not None:\n info[\"video_fps\"] = float(video_fps)\n \n # Get frame dimensions\n iter_video = container.decode(**{\"video\": 0})\n frame = next(iter_video).to_rgb().to_ndarray()\n height, width = frame.shape[:2]\n total_frames = container.streams.video[0].frames\n if total_frames == 0:\n total_frames = MAX_NUM_FRAMES\n warnings.warn(f\"total_frames is 0, using {MAX_NUM_FRAMES} as a fallback\")\n container.close()\n del container\n\n # Pre-allocate frame buffer (np.zeros doesn't actually allocate memory)\n video_frames = np.zeros((total_frames, height, width, 3), dtype=np.uint8)\n\n # Read video frames\n try:\n container = av.open(filename, metadata_errors=\"ignore\")\n assert container.streams.video is not None\n video_frames = _read_from_stream(\n video_frames,\n container,\n start_pts,\n end_pts,\n pts_unit,\n container.streams.video[0],\n {\"video\": 0},\n filename=filename,\n )\n except av.AVError as e:\n print(f\"[Warning] Error while reading video {filename}: {e}\")\n\n # Convert to tensor and adjust format\n vframes = torch.from_numpy(video_frames).clone()\n del video_frames\n if output_format == \"TCHW\":\n # Convert [T,H,W,C] to [T,C,H,W]\n vframes = vframes.permute(0, 3, 1, 2)\n\n aframes = torch.empty((1, 0), dtype=torch.float32)\n return vframes, aframes, info\n\n\ndef _read_from_stream(\n video_frames,\n container: \"av.container.Container\",\n start_offset: float,\n end_offset: float,\n pts_unit: str,\n stream: \"av.stream.Stream\",\n stream_name: Dict[str, Optional[Union[int, Tuple[int, ...], List[int]]]],\n filename: Optional[str] = None,\n) -> List[\"av.frame.Frame\"]:\n \"\"\"Read frames from video stream with proper buffering and seeking\"\"\"\n # Convert time units\n if pts_unit == \"sec\":\n start_offset = int(math.floor(start_offset * (1 / stream.time_base)))\n if end_offset != float(\"inf\"):\n end_offset = int(math.ceil(end_offset * (1 / stream.time_base)))\n else:\n warnings.warn(\"The pts_unit 'pts' gives wrong results. Please use pts_unit 'sec'.\")\n\n # Check if buffering is needed for DivX packed B-frames\n should_buffer = True\n max_buffer_size = 5\n if stream.type == \"video\":\n extradata = stream.codec_context.extradata\n if extradata and b\"DivX\" in extradata:\n pos = extradata.find(b\"DivX\")\n d = extradata[pos:]\n o = re.search(rb\"DivX(\\d+)Build(\\d+)(\\w)\", d)\n if o is None:\n o = re.search(rb\"DivX(\\d+)b(\\d+)(\\w)\", d)\n if o is not None:\n should_buffer = o.group(3) == b\"p\"\n \n # Calculate seek offset with safety margin\n seek_offset = start_offset\n seek_offset = max(seek_offset - 1, 0) # Safety margin for seeking\n if should_buffer:\n seek_offset = max(seek_offset - max_buffer_size, 0)\n \n # Seek to start position\n try:\n container.seek(seek_offset, any_frame=False, backward=True, stream=stream)\n except av.AVError as e:\n print(f\"[Warning] Error while seeking video {filename}: {e}\")\n return []\n\n # Read frames from stream\n buffer_count = 0\n frames_pts = []\n cnt = 0\n try:\n for _idx, frame in enumerate(container.decode(**stream_name)):\n frames_pts.append(frame.pts)\n video_frames[cnt] = frame.to_rgb().to_ndarray()\n cnt += 1\n if cnt >= len(video_frames):\n break\n if frame.pts >= end_offset:\n if should_buffer and buffer_count < max_buffer_size:\n buffer_count += 1\n continue\n break\n except av.AVError as e:\n print(f\"[Warning] Error while reading video {filename}: {e}\")\n\n # Clean up resources to prevent memory leaks\n container.close()\n del container\n gc.collect() # Force garbage collection for PyAV threads\n\n # ensure that the results are sorted wrt the pts\n # NOTE: here we assert frames_pts is sorted\n start_ptr = 0\n end_ptr = cnt\n while start_ptr < end_ptr and frames_pts[start_ptr] < start_offset:\n start_ptr += 1\n while start_ptr < end_ptr and frames_pts[end_ptr - 1] > end_offset:\n end_ptr -= 1\n if start_offset > 0 and start_offset not in frames_pts[start_ptr:end_ptr]:\n # if there is no frame that exactly matches the pts of start_offset\n # add the last frame smaller than start_offset, to guarantee that\n # we will have all the necessary data. This is most useful for audio\n if start_ptr > 0:\n start_ptr -= 1\n result = video_frames[start_ptr:end_ptr].copy()\n return result\n\n\ndef read_video_cv2(filename, start_pts=None, end_pts=None, pts_unit=\"pts\"):\n \"\"\"\n Read video using OpenCV backend.\n \"\"\"\n if pts_unit != \"frames\":\n warnings.warn(\"Using pts_unit other than 'frames' is not supported for cv2 backend\")\n \n cap = cv2.VideoCapture(filename)\n \n # Get video metadata\n fps = cap.get(cv2.CAP_PROP_FPS)\n frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n \n # Calculate frame range\n if start_pts is None:\n start_pts = 0\n if end_pts is None:\n end_pts = frame_count\n \n # Limit frame range to video bounds\n start_pts = max(0, start_pts)\n end_pts = min(frame_count, end_pts)\n num_frames = end_pts - start_pts\n \n if num_frames <= 0:\n return torch.zeros(0, 3, 0, 0), None, {\"video_fps\": fps}\n \n # Seek to start frame\n cap.set(cv2.CAP_PROP_POS_FRAMES, start_pts)\n \n # Read frames\n frames = []\n for i in range(num_frames):\n ret, frame = cap.read()\n if not ret:\n break\n # Convert BGR to RGB and change HWC to CHW format\n frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n frame = torch.from_numpy(frame).permute(2, 0, 1).float()\n frames.append(frame)\n \n cap.release()\n \n if frames:\n video_tensor = torch.stack(frames)\n else:\n video_tensor = torch.zeros(0, 3, 0, 0)\n \n metadata = {\"video_fps\": fps}\n return video_tensor, None, metadata\n\n\ndef read_video(video_path, backend=\"av\"):\n \"\"\"\n Read video using specified backend.\n \"\"\"\n if backend == \"cv2\":\n vframes, vinfo = read_video_cv2(video_path)\n elif backend == \"av\":\n vframes, _, vinfo = read_video_av(filename=video_path, pts_unit=\"sec\", output_format=\"TCHW\")\n else:\n raise ValueError(f\"Unsupported backend: {backend}\")\n\n return vframes, vinfo\n" }, { "path": "utils/scene_detect.py", "content": "\"\"\"\nVideo scene detection and timestamp processing utility.\n\nThis module provides functionality for:\n- Scene detection using PySceneDetect library\n- Timestamp processing and filtering\n- Scene duration management\n- Parallel processing of video files\n\"\"\"\n\nimport argparse\nimport os\nimport ast\nimport concurrent.futures\nimport queue\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom scenedetect import (\n AdaptiveDetector,\n detect,\n ContentDetector,\n SceneManager,\n open_video,\n)\nfrom multiprocessing import Manager\n\n\ndef timecode_to_seconds(timecode):\n \"\"\"Convert timecode string to seconds.\"\"\"\n h, m, s = map(float, timecode.split(\":\"))\n return h * 3600 + m * 60 + s\n\n\ndef seconds_to_timecode(seconds):\n \"\"\"Convert seconds to timecode string format.\"\"\"\n h = int(seconds // 3600)\n m = int((seconds % 3600) // 60)\n s = seconds % 60\n return f\"{h:02d}:{m:02d}:{s:06.3f}\"\n\n\ndef process_single_row(\n row,\n frame_skip=0,\n start_remove_sec=0,\n end_remove_sec=0,\n min_seconds=2,\n max_seconds=15,\n backend=\"opencv\",\n):\n \"\"\"\n Process a single video file for scene detection.\n \"\"\"\n video_path = row[\"video_path\"]\n detector1 = ContentDetector(threshold=21, min_scene_len=15)\n detector2 = AdaptiveDetector(\n adaptive_threshold=3.0, min_scene_len=15, luma_only=True\n )\n detector = [detector1, detector2]\n\n try:\n if isinstance(detector, list):\n scene_manager = SceneManager()\n for i in detector:\n scene_manager.add_detector(i)\n if backend == \"opencv\":\n video = open_video(video_path)\n elif backend == \"av\":\n video = open_video(video_path, backend=\"pyav\")\n # Get video frame rate\n fps = video.frame_rate\n scene_manager.detect_scenes(video=video, frame_skip=frame_skip)\n scene_list = scene_manager.get_scene_list()\n else:\n video = open_video(video_path)\n # Get video frame rate\n fps = video.frame_rate\n scene_list = detect(video_path, detector, start_in_scene=True)\n \n if not scene_list:\n # If no scenes are detected, treat the entire video as one scene\n video_duration = video.duration\n timestamp = [(\"00:00:00.000\", seconds_to_timecode(video_duration.get_seconds()))]\n else:\n timestamp = [(s.get_timecode(), t.get_timecode()) for s, t in scene_list]\n\n # Process timestamps: remove specified seconds from start/end, filter by duration\n new_timestamp = []\n total_remove_sec = start_remove_sec + end_remove_sec\n for start_timecode, end_timecode in timestamp:\n start_seconds = timecode_to_seconds(start_timecode)\n end_seconds = timecode_to_seconds(end_timecode)\n duration = end_seconds - start_seconds\n # Only record scenes longer than total removal time\n if duration >= total_remove_sec:\n new_start_seconds = start_seconds + start_remove_sec\n new_end_seconds = end_seconds - end_remove_sec\n new_duration = new_end_seconds - new_start_seconds\n if new_duration <= max_seconds:\n # Duration within max_seconds, check if meets min_seconds\n if min_seconds <= new_duration:\n new_start_timecode = seconds_to_timecode(new_start_seconds)\n new_end_timecode = seconds_to_timecode(new_end_seconds)\n new_timestamp.append((new_start_timecode, new_end_timecode))\n else:\n # Duration exceeds max_seconds, split into segments\n current_start = new_start_seconds\n while current_start + max_seconds <= new_end_seconds:\n new_start_timecode = seconds_to_timecode(current_start)\n new_end_timecode = seconds_to_timecode(\n current_start + max_seconds\n )\n new_timestamp.append((new_start_timecode, new_end_timecode))\n current_start += max_seconds\n\n # Handle remaining segment\n last_duration = new_end_seconds - current_start\n if last_duration >= min_seconds:\n new_start_timecode = seconds_to_timecode(current_start)\n new_end_timecode = seconds_to_timecode(new_end_seconds)\n new_timestamp.append((new_start_timecode, new_end_timecode))\n\n return True, str(new_timestamp), float(fps)\n except Exception as e:\n print(f\"Video '{video_path}' with error {e}\")\n return False, \"\", None\n\n\ndef timecode_to_frames(timecode, fps):\n \"\"\"Convert timecode to frame number using fps.\"\"\"\n h, m, s = map(float, timecode.split(\":\"))\n total_seconds = h * 3600 + m * 60 + s\n return int(total_seconds * fps)\n\n\ndef worker(task_queue, results_queue, args):\n \"\"\"\n Worker function for parallel scene detection processing.\n \"\"\"\n while True:\n try:\n index, row = task_queue.get(timeout=1)\n except queue.Empty:\n break\n result = process_single_row(\n row,\n frame_skip=args.frame_skip,\n start_remove_sec=args.start_remove_sec,\n end_remove_sec=args.end_remove_sec,\n min_seconds=args.min_seconds,\n max_seconds=args.max_seconds,\n backend=args.backend,\n )\n results_queue.put((index, result))\n task_queue.task_done()\n\n\ndef parse_args():\n \"\"\"Parse command line arguments for scene detection.\"\"\"\n parser = argparse.ArgumentParser()\n parser.add_argument(\n \"--csv_path\",\n type=str,\n required=True,\n help=\"Path to the input CSV file containing video paths.\",\n )\n parser.add_argument(\n \"--num_workers\", type=int, default=1, help=\"#workers for concurrent.futures\"\n )\n parser.add_argument(\n \"--frame_skip\", type=int, default=0, help=\"skip frame for detect_scenes\"\n )\n parser.add_argument(\n \"--start_remove_sec\",\n type=float,\n default=0,\n help=\"Seconds to remove from the start of each timestamp\",\n )\n parser.add_argument(\n \"--end_remove_sec\",\n type=float,\n default=0,\n help=\"Seconds to remove from the end of each timestamp\",\n )\n parser.add_argument(\n \"--min_seconds\",\n type=float,\n default=2,\n help=\"Minimum duration of a scene in seconds\",\n )\n parser.add_argument(\n \"--max_seconds\",\n type=float,\n default=15,\n help=\"Maximum duration of a scene in seconds\",\n )\n parser.add_argument(\n \"--backend\",\n type=str,\n default=\"opencv\",\n choices=[\"opencv\", \"av\"],\n help=\"Backend for video reading\",\n )\n parser.add_argument(\n \"--disable_parallel\", action=\"store_true\", help=\"Disable parallel processing\"\n )\n\n args = parser.parse_args()\n return args\n\n\ndef main():\n args = parse_args()\n csv_path = args.csv_path\n if not os.path.exists(csv_path):\n print(f\"csv file '{csv_path}' not found. Exit.\")\n return\n\n csv = pd.read_csv(csv_path)\n ret = []\n\n if args.disable_parallel:\n for index, row in tqdm(csv.iterrows(), total=len(csv)):\n succ, timestamps, fps = process_single_row(\n row,\n frame_skip=args.frame_skip,\n start_remove_sec=args.start_remove_sec,\n end_remove_sec=args.end_remove_sec,\n min_seconds=args.min_seconds,\n max_seconds=args.max_seconds,\n )\n csv.at[index, \"fps\"] = fps\n csv.at[index, \"timestamp\"] = timestamps\n ret.append((index, (succ, timestamps, fps)))\n else:\n manager = Manager()\n task_queue = manager.Queue()\n results_queue = manager.Queue()\n\n # Add all tasks to queue\n for index, row in csv.iterrows():\n task_queue.put((index, row))\n\n # Set number of workers\n if args.num_workers is not None:\n num_workers = args.num_workers\n else:\n num_workers = os.cpu_count() or 1\n\n # Process videos in parallel\n with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:\n futures = []\n for _ in range(num_workers):\n future = executor.submit(worker, task_queue, results_queue, args)\n futures.append(future)\n\n processed = 0\n total_tasks = len(csv)\n with tqdm(total=total_tasks, desc=\"Processing videos\") as pbar:\n while processed < total_tasks:\n try:\n ret.append(results_queue.get(timeout=1))\n processed += 1\n pbar.update(1)\n except queue.Empty:\n if all(f.done() for f in futures) and results_queue.empty():\n break\n\n for future in futures:\n future.result()\n\n # Collect results\n while not results_queue.empty():\n ret.append(results_queue.get())\n\n # Sort results by index\n ret.sort(key=lambda x: x[0])\n succ, timestamps, fps_list = list(zip(*[result for _, result in ret]))\n\n csv[\"fps\"] = fps_list\n csv[\"timestamp\"] = timestamps\n csv = csv[np.array(succ)]\n\n def calculate_frame_numbers(row):\n \"\"\"Calculate frame numbers from timestamps and fps.\"\"\"\n timestamp = ast.literal_eval(row[\"timestamp\"])\n fps = row[\"fps\"]\n frame_numbers = [\n (timecode_to_frames(start, fps), timecode_to_frames(end, fps))\n for start, end in timestamp\n ]\n return str(frame_numbers)\n\n csv[\"frame_numbers\"] = csv.apply(calculate_frame_numbers, axis=1)\n\n # Save results to new CSV file\n wo_ext, ext = os.path.splitext(csv_path)\n out_path = f\"{wo_ext}_timestamp{ext}\"\n csv.to_csv(out_path, index=False)\n print(\n f\"New csv (shape={csv.shape}) with timestamp and frame numbers saved to '{out_path}'.\"\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/.clang-format", "content": "# C++ formatting rules; used for WebAssembly code.\nBasedOnStyle: LLVM\nAlignAfterOpenBracket: BlockIndent\nBinPackArguments: false\nBinPackParameters: false\nIndentWidth: 4\n" }, { "path": "viser/.gitignore", "content": "*.swp\n*.swo\n*.pyc\n*.egg-info\n*.ipynb_checkpoints\n__pycache__\n.coverage\nhtmlcov\n.mypy_cache\n.dmypy.json\n.hypothesis\n.envrc\n.lvimrc\n.DS_Store\n.envrc\n.vite\nbuild\nsrc/viser/client/build\nsrc/viser/client/.nodeenv\nrecord3d_dance" }, { "path": "viser/.pre-commit-config.yaml", "content": "# See https://pre-commit.com for more information\n# See https://pre-commit.com/hooks.html for more hooks\ndefault_language_version:\n python: python3\nrepos:\n - repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v3.2.0\n hooks:\n - id: trailing-whitespace\n - id: end-of-file-fixer\n - repo: https://github.com/astral-sh/ruff-pre-commit\n # Ruff version.\n rev: v0.6.2\n hooks:\n # Run the linter.\n - id: ruff\n args: [--fix]\n # Run the formatter.\n - id: ruff-format\n" }, { "path": "viser/.prettierignore", "content": "*.mjs\nbuild/\n" }, { "path": "viser/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": "viser/README.md", "content": "

\n \"viser\n viser\n \"viser\n

\n\n

\n \"pyright\"\n \"typescript-compile\"\n \n \"codecov\"\n \n

\n\n\n### This repo is a customized version of https://github.com/nerfstudio-project/viser for project MonST3R (https://monst3r-project.github.io/)\n\n`viser` is a library for interactive 3D visualization in Python.\n\nFeatures include:\n\n- API for visualizing 3D primitives\n- GUI building blocks: buttons, checkboxes, text inputs, sliders, etc.\n- Scene interaction tools (clicks, selection, transform gizmos)\n- Programmatic camera control and rendering\n- An entirely web-based client, for easy use over SSH!\n\nFor usage and API reference, see our documentation.\n\n## Installation\n\nYou can install `viser` with `pip`:\n\n```bash\npip install viser\n```\n\nTo include example dependencies:\n\n```bash\npip install viser[examples]\n```\n\nAfter an example script is running, you can connect by navigating to the printed\nURL (default: `http://localhost:8080`).\n\nSee also: our [development docs](https://viser.studio/latest/development/).\n\n## Examples\n\n**Point cloud visualization**\n\nhttps://github.com/nerfstudio-project/viser/assets/6992947/df35c6ee-78a3-43ad-a2c7-1dddf83f7458\n\nSource: `./examples/07_record3d_visualizer.py`\n\n**Gaussian splatting visualization**\n\nhttps://github.com/nerfstudio-project/viser/assets/6992947/c51b4871-6cc8-4987-8751-2bf186bcb1ae\n\nSource:\n[WangFeng18/3d-gaussian-splatting](https://github.com/WangFeng18/3d-gaussian-splatting)\nand\n[heheyas/gaussian_splatting_3d](https://github.com/heheyas/gaussian_splatting_3d).\n\n**SMPLX visualizer**\n\nhttps://github.com/nerfstudio-project/viser/assets/6992947/78ba0e09-612d-4678-abf3-beaeeffddb01\n\nSource: `./example/08_smpl_visualizer.py`\n\n## Acknowledgements\n\n`viser` is heavily inspired by packages like\n[Pangolin](https://github.com/stevenlovegrove/Pangolin),\n[rviz](https://wiki.ros.org/rviz/),\n[meshcat](https://github.com/rdeits/meshcat), and\n[Gradio](https://github.com/gradio-app/gradio).\nIt's made possible by several open-source projects.\n\nThe web client is implemented using [React](https://react.dev/), with:\n\n- [Vite](https://vitejs.dev/) / [Rollup](https://rollupjs.org/) for bundling\n- [three.js](https://threejs.org/) via [react-three-fiber](https://github.com/pmndrs/react-three-fiber) and [drei](https://github.com/pmndrs/drei)\n- [Mantine](https://mantine.dev/) for UI components\n- [zustand](https://github.com/pmndrs/zustand) for state management\n- [vanilla-extract](https://vanilla-extract.style/) for stylesheets\n\nThe Python API communicates via [msgpack](https://msgpack.org/index.html) and [websockets](https://websockets.readthedocs.io/en/stable/index.html).\n" }, { "path": "viser/docs/.gitignore", "content": "build/\n" }, { "path": "viser/docs/Makefile", "content": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line.\nSPHINXOPTS =\nSPHINXBUILD = sphinx-build\nSPHINXPROJ = viser\nSOURCEDIR = source\nBUILDDIR = ./build\n\n# Put it first so that \"make\" without argument is like \"make help\".\nhelp:\n\t@$(SPHINXBUILD) -M help \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n\n.PHONY: help Makefile\n\n# Catch-all target: route all unknown targets to Sphinx using the new\n# \"make mode\" option. $(O) is meant as a shortcut for $(SPHINXOPTS).\n%: Makefile\n\t@$(SPHINXBUILD) -M $@ \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n" }, { "path": "viser/docs/source/_static/css/custom.css", "content": "img.sidebar-logo {\n width: 5em;\n margin: 1em 0 0 0;\n}\n" }, { "path": "viser/docs/source/_templates/sidebar/brand.html", "content": "\n {% block brand_content %} {%- if logo_url %}\n
\n \"Logo\"\n\n
\n {%- endif %} {%- if theme_light_logo and theme_dark_logo %}\n
\n
\n {%- endif %}\n \n\n {% endblock brand_content %}\n\n\n\n
\n \n \n \n Version: {{ version }}\n \n
\n \n \n
\n\n
\n \n \n Github\n \n
\n" }, { "path": "viser/docs/source/camera_handles.md", "content": "# Camera Handles\n\n\n\n.. autoclass:: viser.CameraHandle\n :members:\n :undoc-members:\n :inherited-members:\n\n\n" }, { "path": "viser/docs/source/client_handles.md", "content": "# Client Handles\n\n\n\n.. autoclass:: viser.ClientHandle\n :members:\n :undoc-members:\n :inherited-members:\n\n.. autoclass:: viser.NotificationHandle\n :members:\n :undoc-members:\n :inherited-members:\n\n\n" }, { "path": "viser/docs/source/conf.py", "content": "# -*- coding: utf-8 -*-\n#\n# Configuration file for the Sphinx documentation builder.\n#\n# This file does only contain a selection of the most common options. For a\n# full list see the documentation:\n# http://www.sphinx-doc.org/en/stable/config\n\nfrom pathlib import Path\nfrom typing import Dict, List\n\nimport m2r2\nimport toml\n\n# -- Path setup --------------------------------------------------------------\n\n# If extensions (or modules to document with autodoc) are in another directory,\n# add these directories to sys.path here. If the directory is relative to the\n# documentation root, use os.path.abspath to make it absolute, like shown here.\n#\n\n# -- Project information -----------------------------------------------------\n\nproject = \"viser\"\ncopyright = \"2024\"\nauthor = \"brentyi\"\n\n# The short X.Y version\nversion: str = toml.load(\n Path(__file__).absolute().parent.parent.parent / \"pyproject.toml\"\n)[\"project\"][\"version\"]\n\n# Formatting!\n# 0.1.30 => v0.1.30\n# dev => dev\nif not version.isalpha():\n version = \"v\" + version\n\n# The full version, including alpha/beta/rc tags\nrelease = \"\"\n\n\n# -- General configuration ---------------------------------------------------\n\n# If your documentation needs a minimal Sphinx version, state it here.\n#\n# needs_sphinx = '1.0'\n\n# Add any Sphinx extension module names here, as strings. They can be\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\n# ones.\nextensions = [\n \"sphinx.ext.autodoc\",\n \"sphinx.ext.todo\",\n \"sphinx.ext.coverage\",\n \"sphinx.ext.mathjax\",\n \"sphinx.ext.githubpages\",\n \"sphinx.ext.napoleon\",\n # \"sphinx.ext.inheritance_diagram\",\n \"sphinx.ext.viewcode\",\n \"m2r2\",\n \"sphinxcontrib.programoutput\",\n \"sphinxcontrib.ansi\",\n \"sphinxcontrib.googleanalytics\", # google analytics extension https://github.com/sphinx-contrib/googleanalytics/tree/master\n]\nprogramoutput_use_ansi = True\nhtml_ansi_stylesheet = \"black-on-white.css\"\nhtml_static_path = [\"_static\"]\nhtml_theme_options = {\n \"light_css_variables\": {\n \"color-code-background\": \"#f4f4f4\",\n \"color-code-foreground\": \"#000\",\n },\n \"footer_icons\": [\n {\n \"name\": \"GitHub\",\n \"url\": \"https://github.com/nerfstudio-project/viser\",\n \"html\": \"\"\"\n \n \n \n \"\"\",\n \"class\": \"\",\n },\n ],\n \"light_logo\": \"logo.svg\",\n \"dark_logo\": \"logo.svg\",\n}\n\n# Pull documentation types from hints\nautodoc_typehints = \"both\"\nautodoc_class_signature = \"separated\"\nautodoc_default_options = {\n \"members\": True,\n \"member-order\": \"bysource\",\n \"undoc-members\": True,\n \"inherited-members\": True,\n \"exclude-members\": \"__init__, __post_init__\",\n \"imported-members\": True,\n}\n\n# Add any paths that contain templates here, relative to this directory.\ntemplates_path = [\"_templates\"]\n\n# The suffix(es) of source filenames.\n# You can specify multiple suffix as a list of string:\n#\nsource_suffix = [\".rst\", \".md\"]\n# source_suffix = \".rst\"\n\n# The master toctree document.\nmaster_doc = \"index\"\n\n# The language for content autogenerated by Sphinx. Refer to documentation\n# for a list of supported languages.\n#\n# This is also used if you do content translation via gettext catalogs.\n# Usually you set \"language\" from the command line for these cases.\nlanguage: str = \"en\"\n\n# List of patterns, relative to source directory, that match files and\n# directories to ignore when looking for source files.\n# This pattern also affects html_static_path and html_extra_path .\nexclude_patterns: List[str] = []\n\n# The name of the Pygments (syntax highlighting) style to use.\npygments_style = \"monokai\"\n\n\n# -- Options for HTML output -------------------------------------------------\n\n# The theme to use for HTML and HTML Help pages. See the documentation for\n# a list of builtin themes.\n#\nhtml_theme = \"furo\"\nhtml_title = \"viser\"\n\n\n# Theme options are theme-specific and customize the look and feel of a theme\n# further. For a list of options available for each theme, see the\n# documentation.\n#\n# html_theme_options = {}\n\n# Add any paths that contain custom static files (such as style sheets) here,\n# relative to this directory. They are copied after the builtin static files,\n# so a file named \"default.css\" will overwrite the builtin \"default.css\".\n# html_static_path = [\"_static\"]\n\n# Custom sidebar templates, must be a dictionary that maps document names\n# to template names.\n#\n# The default sidebars (for documents that don't match any pattern) are\n# defined by theme itself. Builtin themes are using these templates by\n# default: ``['localtoc.html', 'relations.html', 'sourcelink.html',\n# 'searchbox.html']``.\n#\n# html_sidebars = {}\n\n\n# -- Options for HTMLHelp output ---------------------------------------------\n\n# Output file base name for HTML help builder.\nhtmlhelp_basename = \"viser_doc\"\n\n\n# -- Options for Github output ------------------------------------------------\n\nsphinx_to_github = True\nsphinx_to_github_verbose = True\nsphinx_to_github_encoding = \"utf-8\"\n\n\n# -- Options for LaTeX output ------------------------------------------------\n\nlatex_elements: Dict[str, str] = {\n # The paper size ('letterpaper' or 'a4paper').\n #\n # 'papersize': 'letterpaper',\n # The font size ('10pt', '11pt' or '12pt').\n #\n # 'pointsize': '10pt',\n # Additional stuff for the LaTeX preamble.\n #\n # 'preamble': '',\n # Latex figure (float) alignment\n #\n # 'figure_align': 'htbp',\n}\n\n# Grouping the document tree into LaTeX files. List of tuples\n# (source start file, target name, title,\n# author, documentclass [howto, manual, or own class]).\nlatex_documents = [\n (\n master_doc,\n \"viser.tex\",\n \"viser\",\n \"brentyi\",\n \"manual\",\n ),\n]\n\n\n# -- Options for manual page output ------------------------------------------\n\n# One entry per manual page. List of tuples\n# (source start file, name, description, authors, manual section).\nman_pages = [(master_doc, \"viser\", \"viser documentation\", [author], 1)]\n\n\n# -- Options for Texinfo output ----------------------------------------------\n\n# Grouping the document tree into Texinfo files. List of tuples\n# (source start file, target name, title, author,\n# dir menu entry, description, category)\ntexinfo_documents = [\n (\n master_doc,\n \"viser\",\n \"viser\",\n author,\n \"viser\",\n \"viser\",\n \"Miscellaneous\",\n ),\n]\n\n\n# -- Extension configuration --------------------------------------------------\n\n# Google Analytics ID\ngoogleanalytics_id = \"G-RRGY51J5ZH\"\n\n# -- Options for todo extension ----------------------------------------------\n\n# If true, `todo` and `todoList` produce output, else they produce nothing.\ntodo_include_todos = True\n\n# -- Enable Markdown -> RST conversion ----------------------------------------\n\n\ndef docstring(app, what, name, obj, options, lines):\n md = \"\\n\".join(lines)\n rst = m2r2.convert(md)\n lines.clear()\n lines += rst.splitlines()\n\n\ndef setup(app):\n app.connect(\"autodoc-process-docstring\", docstring)\n app.add_css_file(\"css/custom.css\")\n\n\n# -- Napoleon settings -------------------------------------------------------\n\n# Settings for parsing non-sphinx style docstrings. We use Google style in this\n# project.\nnapoleon_google_docstring = True\nnapoleon_numpy_docstring = False\nnapoleon_include_init_with_doc = False\nnapoleon_include_private_with_doc = False\nnapoleon_include_special_with_doc = True\nnapoleon_use_admonition_for_examples = False\nnapoleon_use_admonition_for_notes = False\nnapoleon_use_admonition_for_references = False\nnapoleon_use_ivar = False\nnapoleon_use_param = True\nnapoleon_use_rtype = True\nnapoleon_preprocess_types = False\nnapoleon_type_aliases = None\nnapoleon_attr_annotations = True\n" }, { "path": "viser/docs/source/conventions.md", "content": "# Frame Conventions\n\nIn this note, we describe the coordinate frame conventions used in `viser`.\n\n## Scene tree naming\n\nEach object that we add to the scene in viser is instantiated as a node in a\nscene tree. The structure of this tree is determined by the names assigned to\nthe nodes.\n\nIf we add a coordinate frame called `/base_link/shoulder/wrist`, it signifies\nthree nodes: the `wrist` is a child of the `shoulder` which is a child of the\n`base_link`.\n\nIf we set the transformation of a given node like `/base_link/shoulder`, both\nit and its child `/base_link/shoulder/wrist` will move. Its parent,\n`/base_link`, will be unaffected.\n\n## Poses\n\nPoses in `viser` are defined using a pair of fields:\n\n- `wxyz`, a unit quaternion orientation term. This should always be 4D.\n- `position`, a translation term. This should always be 3D.\n\nThese correspond to a transformation from coordinates in the local frame to the\nparent frame:\n\n\n\n.. math::\n\n p_\\mathrm{parent} = \\begin{bmatrix} R & t \\end{bmatrix}\\begin{bmatrix}p_\\mathrm{local} \\\\ 1\\end{bmatrix}\n\n\n\nwhere `wxyz` is the quaternion form of the :math:`\\mathrm{SO}(3)` matrix\n:math:`R` and `position` is the :math:`\\mathbb{R}^3` translation term\n:math:`t`.\n\n## World coordinates\n\nIn the world coordinate space, +Z points upward by default. This can be\noverridden with :func:`viser.SceneApi.set_up_direction()`.\n\n## Cameras\n\nIn `viser`, all camera parameters exposed to the Python API use the\nCOLMAP/OpenCV convention:\n\n- Forward: +Z\n- Up: -Y\n- Right: +X\n\nConfusingly, this is different from Nerfstudio, which adopts the OpenGL/Blender\nconvention:\n\n- Forward: -Z\n- Up: +Y\n- Right: +X\n\nConversion between the two is a simple 180 degree rotation around the local X-axis.\n" }, { "path": "viser/docs/source/development.md", "content": "# Development\n\nIn this note, we outline current practices, tools, and workflows for `viser`\ndevelopment. We assume that the repository is cloned to `~/viser`.\n\n## Python install\n\nWe recommend an editable install for Python development, ideally in a virtual\nenvironment (eg via conda).\n\n```bash\n# Install package.\ncd ~/viser\npip install -e .\n\n# Install example dependencies.\npip install -e .[examples]\n```\n\nAfter installation, any of the example scripts (`~/viser/examples`) should be\nrunnable. A few of them require downloading assets, which can be done via the\nscripts in `~/viser/examples/assets`.\n\n**Linting, formatting, type-checking.**\n\nFirst, install developer tools:\n\n```bash\n# Using pip.\npip install -e .[dev]\npre-commit install\n```\n\nIt would be hard to write unit tests for `viser`. We rely on static typing for\nrobustness. To check your code, you can run the following:\n\n```bash\n# runs linting, formatting, and type-checking\nviser-dev-checks\n```\n\n## Message updates\n\nThe `viser` frontend and backend communicate via a shared set of message\ndefinitions:\n\n- On the server, these are defined as Python dataclasses in\n `~/viser/src/viser/_messages.py`.\n- On the client, these are defined as TypeScript interfaces in\n `~/viser/src/viser/client/src/WebsocketMessages.tsx`.\n\nNote that there is a 1:1 correspondence between the dataclasses message types\nand the TypeScript ones.\n\nThe TypeScript definitions should not be manually modified. Instead, changes\nshould be made in Python and synchronized via the `sync_message_defs.py` script:\n\n```\ncd ~/viser\npython sync_message_defs.py\n```\n\n## Client development\n\nFor client development, we can start by launching a relevant Python script. The\nexamples are a good place to start:\n\n```\ncd ~/viser/examples\npython 05_camera_commands.py\n```\n\nWhen a `viser` script is launched, two URLs will be printed:\n\n- An HTTP URL, like `http://localhost:8080`, which can be used to open a\n _pre-built_ version of the React frontend.\n- A websocket URL, like `ws://localhost:8080`, which client applications can\n connect to.\n\nIf changes to the client source files are detected on startup, `viser` will\nre-build the client automatically. This is okay for quick changes, but for\nfaster iteration we can also launch a development version of the frontend, which\nwill reflect changes we make to the client source files\n(`~/viser/src/viser/client/src`) without a full build. This requires a few more\nsteps.\n\n**Installing dependencies.**\n\n1. [Install nodejs.](https://nodejs.dev/en/download/package-manager)\n2. [Install yarn.](https://yarnpkg.com/getting-started/install)\n3. Install dependencies.\n ```\n cd ~/viser/src/viser/client\n yarn install\n ```\n\n**Launching client.**\n\nTo launch the client, we can run:\n\n```\ncd ~/viser/src/viser/client\nyarn start\n```\n\nfrom the `viser/src/viser/client` directory. After opening the client in a web\nbrowser, the websocket server address typically needs to be updated in the\n\"Server\" tab.\n\n**Formatting.**\n\nWe use [prettier](https://prettier.io/docs/en/install.html). This can be run via\none of:\n\n- `prettier -w .`\n- `npx prettier -w .`\n\nfrom `~/viser/src/viser/client`.\n" }, { "path": "viser/docs/source/events.md", "content": "# Events\n\nWe define a small set of event types, which are passed to callback functions\nwhen events like clicks or GUI updates are triggered.\n\n\n\n.. autoclass:: viser.ScenePointerEvent()\n\n.. autoclass:: viser.SceneNodePointerEvent()\n\n.. autoclass:: viser.GuiEvent()\n\n\n" }, { "path": "viser/docs/source/examples/00_coordinate_frames.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nCoordinate frames\n==========================================\n\n\nIn this basic example, we visualize a set of coordinate frames.\n\nNaming for all scene nodes are hierarchical; /tree/branch, for example, is defined\nrelative to /tree.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import random\n import time\n\n import viser\n\n server = viser.ViserServer()\n\n while True:\n # Add some coordinate frames to the scene. These will be visualized in the viewer.\n server.scene.add_frame(\n \"/tree\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(random.random() * 2.0, 2.0, 0.2),\n )\n server.scene.add_frame(\n \"/tree/branch\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(random.random() * 2.0, 2.0, 0.2),\n )\n leaf = server.scene.add_frame(\n \"/tree/branch/leaf\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(random.random() * 2.0, 2.0, 0.2),\n )\n time.sleep(5.0)\n\n # Remove the leaf node from the scene.\n leaf.remove()\n time.sleep(0.5)\n" }, { "path": "viser/docs/source/examples/01_image.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nImages\n==========================================\n\n\nExample for sending images to the viewer.\n\nWe can send backgrond images to display behind the viewer (useful for visualizing\nNeRFs), or images to render as 3D textures.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n from pathlib import Path\n\n import imageio.v3 as iio\n import numpy as onp\n\n import viser\n\n\n def main() -> None:\n server = viser.ViserServer()\n\n # Add a background image.\n server.scene.set_background_image(\n iio.imread(Path(__file__).parent / \"assets/Cal_logo.png\"),\n format=\"png\",\n )\n\n # Add main image.\n server.scene.add_image(\n \"/img\",\n iio.imread(Path(__file__).parent / \"assets/Cal_logo.png\"),\n 4.0,\n 4.0,\n format=\"png\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(2.0, 2.0, 0.0),\n )\n while True:\n server.scene.add_image(\n \"/noise\",\n onp.random.randint(\n 0,\n 256,\n size=(400, 400, 3),\n dtype=onp.uint8,\n ),\n 4.0,\n 4.0,\n format=\"jpeg\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(2.0, 2.0, -1e-2),\n )\n time.sleep(0.2)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/02_gui.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nGUI basics\n==========================================\n\n\nExamples of basic GUI elements that we can create, read from, and write to.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import numpy as onp\n\n import viser\n\n\n def main() -> None:\n server = viser.ViserServer()\n\n # Add some common GUI elements: number inputs, sliders, vectors, checkboxes.\n with server.gui.add_folder(\"Read-only\"):\n gui_counter = server.gui.add_number(\n \"Counter\",\n initial_value=0,\n disabled=True,\n )\n gui_slider = server.gui.add_slider(\n \"Slider\",\n min=0,\n max=100,\n step=1,\n initial_value=0,\n disabled=True,\n )\n gui_progress = server.gui.add_progress_bar(25, animated=True)\n\n with server.gui.add_folder(\"Editable\"):\n gui_vector2 = server.gui.add_vector2(\n \"Position\",\n initial_value=(0.0, 0.0),\n step=0.1,\n )\n gui_vector3 = server.gui.add_vector3(\n \"Size\",\n initial_value=(1.0, 1.0, 1.0),\n step=0.25,\n )\n with server.gui.add_folder(\"Text toggle\"):\n gui_checkbox_hide = server.gui.add_checkbox(\n \"Hide\",\n initial_value=False,\n )\n gui_text = server.gui.add_text(\n \"Text\",\n initial_value=\"Hello world\",\n )\n gui_button = server.gui.add_button(\"Button\")\n gui_checkbox_disable = server.gui.add_checkbox(\n \"Disable\",\n initial_value=False,\n )\n gui_rgb = server.gui.add_rgb(\n \"Color\",\n initial_value=(255, 255, 0),\n )\n gui_multi_slider = server.gui.add_multi_slider(\n \"Multi slider\",\n min=0,\n max=100,\n step=1,\n initial_value=(0, 30, 100),\n marks=((0, \"0\"), (50, \"5\"), (70, \"7\"), 99),\n )\n gui_slider_positions = server.gui.add_slider(\n \"# sliders\",\n min=0,\n max=10,\n step=1,\n initial_value=3,\n marks=((0, \"0\"), (5, \"5\"), (7, \"7\"), 10),\n )\n gui_upload_button = server.gui.add_upload_button(\n \"Upload\", icon=viser.Icon.UPLOAD\n )\n\n @gui_upload_button.on_upload\n def _(_) -> None:\n \"\"\"Callback for when a file is uploaded.\"\"\"\n file = gui_upload_button.value\n print(file.name, len(file.content), \"bytes\")\n\n # Pre-generate a point cloud to send.\n point_positions = onp.random.uniform(low=-1.0, high=1.0, size=(5000, 3))\n color_coeffs = onp.random.uniform(0.4, 1.0, size=(point_positions.shape[0]))\n\n counter = 0\n while True:\n # We can set the value of an input to a particular value. Changes are\n # automatically reflected in connected clients.\n gui_counter.value = counter\n gui_slider.value = counter % 100\n\n # We can set the position of a scene node with `.position`, and read the value\n # of a gui element with `.value`. Changes are automatically reflected in\n # connected clients.\n server.scene.add_point_cloud(\n \"/point_cloud\",\n points=point_positions * onp.array(gui_vector3.value, dtype=onp.float32),\n colors=(\n onp.tile(gui_rgb.value, point_positions.shape[0]).reshape((-1, 3))\n * color_coeffs[:, None]\n ).astype(onp.uint8),\n position=gui_vector2.value + (0,),\n point_shape=\"circle\",\n )\n\n gui_progress.value = float((counter % 100))\n\n # We can use `.visible` and `.disabled` to toggle GUI elements.\n gui_text.visible = not gui_checkbox_hide.value\n gui_button.visible = not gui_checkbox_hide.value\n gui_rgb.disabled = gui_checkbox_disable.value\n gui_button.disabled = gui_checkbox_disable.value\n gui_upload_button.disabled = gui_checkbox_disable.value\n\n # Update the number of handles in the multi-slider.\n if gui_slider_positions.value != len(gui_multi_slider.value):\n gui_multi_slider.value = onp.linspace(\n 0, 100, gui_slider_positions.value, dtype=onp.int64\n )\n\n counter += 1\n time.sleep(0.01)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/03_gui_callbacks.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nGUI callbacks\n==========================================\n\n\nAsynchronous usage of GUI elements: we can attach callbacks that are called as soon as\nwe get updates.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import numpy as onp\n from typing_extensions import assert_never\n\n import viser\n\n\n def main() -> None:\n server = viser.ViserServer()\n\n gui_reset_scene = server.gui.add_button(\"Reset Scene\")\n\n gui_plane = server.gui.add_dropdown(\n \"Grid plane\", (\"xz\", \"xy\", \"yx\", \"yz\", \"zx\", \"zy\")\n )\n\n def update_plane() -> None:\n server.scene.add_grid(\n \"/grid\",\n width=10.0,\n height=20.0,\n width_segments=10,\n height_segments=20,\n plane=gui_plane.value,\n )\n\n gui_plane.on_update(lambda _: update_plane())\n\n with server.gui.add_folder(\"Control\"):\n gui_show_frame = server.gui.add_checkbox(\"Show Frame\", initial_value=True)\n gui_show_everything = server.gui.add_checkbox(\n \"Show Everything\", initial_value=True\n )\n gui_axis = server.gui.add_dropdown(\"Axis\", (\"x\", \"y\", \"z\"))\n gui_include_z = server.gui.add_checkbox(\"Z in dropdown\", initial_value=True)\n\n @gui_include_z.on_update\n def _(_) -> None:\n gui_axis.options = (\"x\", \"y\", \"z\") if gui_include_z.value else (\"x\", \"y\")\n\n with server.gui.add_folder(\"Sliders\"):\n gui_location = server.gui.add_slider(\n \"Location\", min=-5.0, max=5.0, step=0.05, initial_value=0.0\n )\n gui_num_points = server.gui.add_slider(\n \"# Points\", min=1000, max=200_000, step=1000, initial_value=10_000\n )\n\n def draw_frame() -> None:\n axis = gui_axis.value\n if axis == \"x\":\n pos = (gui_location.value, 0.0, 0.0)\n elif axis == \"y\":\n pos = (0.0, gui_location.value, 0.0)\n elif axis == \"z\":\n pos = (0.0, 0.0, gui_location.value)\n else:\n assert_never(axis)\n\n server.scene.add_frame(\n \"/frame\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=pos,\n show_axes=gui_show_frame.value,\n axes_length=5.0,\n )\n\n def draw_points() -> None:\n num_points = gui_num_points.value\n server.scene.add_point_cloud(\n \"/frame/point_cloud\",\n points=onp.random.normal(size=(num_points, 3)),\n colors=onp.random.randint(0, 256, size=(num_points, 3)),\n )\n\n # We can (optionally) also attach callbacks!\n # Here, we update the point clouds + frames whenever any of the GUI items are updated.\n gui_show_frame.on_update(lambda _: draw_frame())\n gui_show_everything.on_update(\n lambda _: server.scene.set_global_visibility(gui_show_everything.value)\n )\n gui_axis.on_update(lambda _: draw_frame())\n gui_location.on_update(lambda _: draw_frame())\n gui_num_points.on_update(lambda _: draw_points())\n\n @gui_reset_scene.on_click\n def _(_) -> None:\n \"\"\"Reset the scene when the reset button is clicked.\"\"\"\n gui_show_frame.value = True\n gui_location.value = 0.0\n gui_axis.value = \"x\"\n gui_num_points.value = 10_000\n\n draw_frame()\n draw_points()\n\n # Finally, let's add the initial frame + point cloud and just loop infinitely. :)\n update_plane()\n draw_frame()\n draw_points()\n while True:\n time.sleep(1.0)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/04_camera_poses.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nCamera poses\n==========================================\n\n\nExample showing how we can detect new clients and read camera poses from them.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import viser\n\n server = viser.ViserServer()\n server.scene.world_axes.visible = True\n\n\n @server.on_client_connect\n def _(client: viser.ClientHandle) -> None:\n print(\"new client!\")\n\n # This will run whenever we get a new camera!\n @client.camera.on_update\n def _(_: viser.CameraHandle) -> None:\n print(f\"New camera on client {client.client_id}!\")\n\n # Show the client ID in the GUI.\n gui_info = client.gui.add_text(\"Client ID\", initial_value=str(client.client_id))\n gui_info.disabled = True\n\n\n while True:\n # Get all currently connected clients.\n clients = server.get_clients()\n print(\"Connected client IDs\", clients.keys())\n\n for id, client in clients.items():\n print(f\"Camera pose for client {id}\")\n print(f\"\\twxyz: {client.camera.wxyz}\")\n print(f\"\\tposition: {client.camera.position}\")\n print(f\"\\tfov: {client.camera.fov}\")\n print(f\"\\taspect: {client.camera.aspect}\")\n print(f\"\\tlast update: {client.camera.update_timestamp}\")\n\n time.sleep(2.0)\n" }, { "path": "viser/docs/source/examples/05_camera_commands.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nCamera commands\n==========================================\n\n\nIn addition to reads, camera parameters also support writes. These are synced to the\ncorresponding client automatically.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import numpy as onp\n\n import viser\n import viser.transforms as tf\n\n server = viser.ViserServer()\n num_frames = 20\n\n\n @server.on_client_connect\n def _(client: viser.ClientHandle) -> None:\n \"\"\"For each client that connects, we create a set of random frames + a click handler for each frame.\n\n When a frame is clicked, we move the camera to the corresponding frame.\n \"\"\"\n\n rng = onp.random.default_rng(0)\n\n def make_frame(i: int) -> None:\n # Sample a random orientation + position.\n wxyz = rng.normal(size=4)\n wxyz /= onp.linalg.norm(wxyz)\n position = rng.uniform(-3.0, 3.0, size=(3,))\n\n # Create a coordinate frame and label.\n frame = client.scene.add_frame(f\"/frame_{i}\", wxyz=wxyz, position=position)\n client.scene.add_label(f\"/frame_{i}/label\", text=f\"Frame {i}\")\n\n # Move the camera when we click a frame.\n @frame.on_click\n def _(_):\n T_world_current = tf.SE3.from_rotation_and_translation(\n tf.SO3(client.camera.wxyz), client.camera.position\n )\n T_world_target = tf.SE3.from_rotation_and_translation(\n tf.SO3(frame.wxyz), frame.position\n ) @ tf.SE3.from_translation(onp.array([0.0, 0.0, -0.5]))\n\n T_current_target = T_world_current.inverse() @ T_world_target\n\n for j in range(20):\n T_world_set = T_world_current @ tf.SE3.exp(\n T_current_target.log() * j / 19.0\n )\n\n # We can atomically set the orientation and the position of the camera\n # together to prevent jitter that might happen if one was set before the\n # other.\n with client.atomic():\n client.camera.wxyz = T_world_set.rotation().wxyz\n client.camera.position = T_world_set.translation()\n\n client.flush() # Optional!\n time.sleep(1.0 / 60.0)\n\n # Mouse interactions should orbit around the frame origin.\n client.camera.look_at = frame.position\n\n for i in range(num_frames):\n make_frame(i)\n\n\n while True:\n time.sleep(1.0)\n" }, { "path": "viser/docs/source/examples/06_mesh.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nMeshes\n==========================================\n\n\nVisualize a mesh. To get the demo data, see ``./assets/download_dragon_mesh.sh``.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n from pathlib import Path\n\n import numpy as onp\n import trimesh\n\n import viser\n import viser.transforms as tf\n\n mesh = trimesh.load_mesh(str(Path(__file__).parent / \"assets/dragon.obj\"))\n assert isinstance(mesh, trimesh.Trimesh)\n mesh.apply_scale(0.05)\n\n vertices = mesh.vertices\n faces = mesh.faces\n print(f\"Loaded mesh with {vertices.shape} vertices, {faces.shape} faces\")\n\n server = viser.ViserServer()\n server.scene.add_mesh_simple(\n name=\"/simple\",\n vertices=vertices,\n faces=faces,\n wxyz=tf.SO3.from_x_radians(onp.pi / 2).wxyz,\n position=(0.0, 0.0, 0.0),\n )\n server.scene.add_mesh_trimesh(\n name=\"/trimesh\",\n mesh=mesh.smoothed(),\n wxyz=tf.SO3.from_x_radians(onp.pi / 2).wxyz,\n position=(0.0, 5.0, 0.0),\n )\n\n while True:\n time.sleep(10.0)\n" }, { "path": "viser/docs/source/examples/07_record3d_visualizer.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nRecord3D visualizer\n==========================================\n\n\nParse and stream record3d captures. To get the demo data, see ``./assets/download_record3d_dance.sh``.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n from pathlib import Path\n\n import numpy as onp\n import tyro\n from tqdm.auto import tqdm\n\n import viser\n import viser.extras\n import viser.transforms as tf\n\n\n def main(\n data_path: Path = Path(__file__).parent / \"assets/record3d_dance\",\n downsample_factor: int = 4,\n max_frames: int = 100,\n share: bool = False,\n ) -> None:\n server = viser.ViserServer()\n if share:\n server.request_share_url()\n\n print(\"Loading frames!\")\n loader = viser.extras.Record3dLoader(data_path)\n num_frames = min(max_frames, loader.num_frames())\n\n # Add playback UI.\n with server.gui.add_folder(\"Playback\"):\n gui_timestep = server.gui.add_slider(\n \"Timestep\",\n min=0,\n max=num_frames - 1,\n step=1,\n initial_value=0,\n disabled=True,\n )\n gui_next_frame = server.gui.add_button(\"Next Frame\", disabled=True)\n gui_prev_frame = server.gui.add_button(\"Prev Frame\", disabled=True)\n gui_playing = server.gui.add_checkbox(\"Playing\", True)\n gui_framerate = server.gui.add_slider(\n \"FPS\", min=1, max=60, step=0.1, initial_value=loader.fps\n )\n gui_framerate_options = server.gui.add_button_group(\n \"FPS options\", (\"10\", \"20\", \"30\", \"60\")\n )\n\n # Frame step buttons.\n @gui_next_frame.on_click\n def _(_) -> None:\n gui_timestep.value = (gui_timestep.value + 1) % num_frames\n\n @gui_prev_frame.on_click\n def _(_) -> None:\n gui_timestep.value = (gui_timestep.value - 1) % num_frames\n\n # Disable frame controls when we're playing.\n @gui_playing.on_update\n def _(_) -> None:\n gui_timestep.disabled = gui_playing.value\n gui_next_frame.disabled = gui_playing.value\n gui_prev_frame.disabled = gui_playing.value\n\n # Set the framerate when we click one of the options.\n @gui_framerate_options.on_click\n def _(_) -> None:\n gui_framerate.value = int(gui_framerate_options.value)\n\n prev_timestep = gui_timestep.value\n\n # Toggle frame visibility when the timestep slider changes.\n @gui_timestep.on_update\n def _(_) -> None:\n nonlocal prev_timestep\n current_timestep = gui_timestep.value\n with server.atomic():\n frame_nodes[current_timestep].visible = True\n frame_nodes[prev_timestep].visible = False\n prev_timestep = current_timestep\n server.flush() # Optional!\n\n # Load in frames.\n server.scene.add_frame(\n \"/frames\",\n wxyz=tf.SO3.exp(onp.array([onp.pi / 2.0, 0.0, 0.0])).wxyz,\n position=(0, 0, 0),\n show_axes=False,\n )\n frame_nodes: list[viser.FrameHandle] = []\n for i in tqdm(range(num_frames)):\n frame = loader.get_frame(i)\n position, color = frame.get_point_cloud(downsample_factor)\n\n # Add base frame.\n frame_nodes.append(server.scene.add_frame(f\"/frames/t{i}\", show_axes=False))\n\n # Place the point cloud in the frame.\n server.scene.add_point_cloud(\n name=f\"/frames/t{i}/point_cloud\",\n points=position,\n colors=color,\n point_size=0.01,\n point_shape=\"rounded\",\n )\n\n # Place the frustum.\n fov = 2 * onp.arctan2(frame.rgb.shape[0] / 2, frame.K[0, 0])\n aspect = frame.rgb.shape[1] / frame.rgb.shape[0]\n server.scene.add_camera_frustum(\n f\"/frames/t{i}/frustum\",\n fov=fov,\n aspect=aspect,\n scale=0.15,\n image=frame.rgb[::downsample_factor, ::downsample_factor],\n wxyz=tf.SO3.from_matrix(frame.T_world_camera[:3, :3]).wxyz,\n position=frame.T_world_camera[:3, 3],\n )\n\n # Add some axes.\n server.scene.add_frame(\n f\"/frames/t{i}/frustum/axes\",\n axes_length=0.05,\n axes_radius=0.005,\n )\n\n # Hide all but the current frame.\n for i, frame_node in enumerate(frame_nodes):\n frame_node.visible = i == gui_timestep.value\n\n # Playback update loop.\n prev_timestep = gui_timestep.value\n while True:\n if gui_playing.value:\n gui_timestep.value = (gui_timestep.value + 1) % num_frames\n\n time.sleep(1.0 / gui_framerate.value)\n\n\n if __name__ == \"__main__\":\n tyro.cli(main)\n" }, { "path": "viser/docs/source/examples/08_smpl_visualizer.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nSMPL model visualizer\n==========================================\n\n\nVisualizer for SMPL human body models. Requires a .npz model file.\n\nSee here for download instructions:\n https://github.com/vchoutas/smplx?tab=readme-ov-file#downloading-the-model\n\n\n\n.. code-block:: python\n :linenos:\n\n\n from __future__ import annotations\n\n import time\n from dataclasses import dataclass\n from pathlib import Path\n\n import numpy as np\n import numpy as onp\n import tyro\n\n import viser\n import viser.transforms as tf\n\n\n @dataclass(frozen=True)\n class SmplOutputs:\n vertices: np.ndarray\n faces: np.ndarray\n T_world_joint: np.ndarray # (num_joints, 4, 4)\n T_parent_joint: np.ndarray # (num_joints, 4, 4)\n\n\n class SmplHelper:\n \"\"\"Helper for models in the SMPL family, implemented in numpy.\"\"\"\n\n def __init__(self, model_path: Path) -> None:\n assert model_path.suffix.lower() == \".npz\", \"Model should be an .npz file!\"\n body_dict = dict(**onp.load(model_path, allow_pickle=True))\n\n self._J_regressor = body_dict[\"J_regressor\"]\n self._weights = body_dict[\"weights\"]\n self._v_template = body_dict[\"v_template\"]\n self._posedirs = body_dict[\"posedirs\"]\n self._shapedirs = body_dict[\"shapedirs\"]\n self._faces = body_dict[\"f\"]\n\n self.num_joints: int = self._weights.shape[-1]\n self.num_betas: int = self._shapedirs.shape[-1]\n self.parent_idx: np.ndarray = body_dict[\"kintree_table\"][0]\n\n def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutputs:\n # Get shaped vertices + joint positions, when all local poses are identity.\n v_tpose = self._v_template + np.einsum(\"vxb,b->vx\", self._shapedirs, betas)\n j_tpose = np.einsum(\"jv,vx->jx\", self._J_regressor, v_tpose)\n\n # Local SE(3) transforms.\n T_parent_joint = np.zeros((self.num_joints, 4, 4)) + np.eye(4)\n T_parent_joint[:, :3, :3] = joint_rotmats\n T_parent_joint[0, :3, 3] = j_tpose[0]\n T_parent_joint[1:, :3, 3] = j_tpose[1:] - j_tpose[self.parent_idx[1:]]\n\n # Forward kinematics.\n T_world_joint = T_parent_joint.copy()\n for i in range(1, self.num_joints):\n T_world_joint[i] = T_world_joint[self.parent_idx[i]] @ T_parent_joint[i]\n\n # Linear blend skinning.\n pose_delta = (joint_rotmats[1:, ...] - np.eye(3)).flatten()\n v_blend = v_tpose + np.einsum(\"byn,n->by\", self._posedirs, pose_delta)\n v_delta = np.ones((v_blend.shape[0], self.num_joints, 4))\n v_delta[:, :, :3] = v_blend[:, None, :] - j_tpose[None, :, :]\n v_posed = np.einsum(\n \"jxy,vj,vjy->vx\", T_world_joint[:, :3, :], self._weights, v_delta\n )\n return SmplOutputs(v_posed, self._faces, T_world_joint, T_parent_joint)\n\n\n def main(model_path: Path) -> None:\n server = viser.ViserServer()\n server.scene.set_up_direction(\"+y\")\n server.gui.configure_theme(control_layout=\"collapsible\")\n\n # Main loop. We'll read pose/shape from the GUI elements, compute the mesh,\n # and then send the updated mesh in a loop.\n model = SmplHelper(model_path)\n gui_elements = make_gui_elements(\n server,\n num_betas=model.num_betas,\n num_joints=model.num_joints,\n parent_idx=model.parent_idx,\n )\n while True:\n # Do nothing if no change.\n time.sleep(0.02)\n if not gui_elements.changed:\n continue\n\n gui_elements.changed = False\n\n # Compute SMPL outputs.\n smpl_outputs = model.get_outputs(\n betas=np.array([x.value for x in gui_elements.gui_betas]),\n joint_rotmats=tf.SO3.exp(\n # (num_joints, 3)\n np.array([x.value for x in gui_elements.gui_joints])\n ).as_matrix(),\n )\n server.scene.add_mesh_simple(\n \"/human\",\n smpl_outputs.vertices,\n smpl_outputs.faces,\n wireframe=gui_elements.gui_wireframe.value,\n color=gui_elements.gui_rgb.value,\n )\n\n # Match transform control gizmos to joint positions.\n for i, control in enumerate(gui_elements.transform_controls):\n control.position = smpl_outputs.T_parent_joint[i, :3, 3]\n\n\n @dataclass\n class GuiElements:\n \"\"\"Structure containing handles for reading from GUI elements.\"\"\"\n\n gui_rgb: viser.GuiInputHandle[tuple[int, int, int]]\n gui_wireframe: viser.GuiInputHandle[bool]\n gui_betas: list[viser.GuiInputHandle[float]]\n gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]]\n transform_controls: list[viser.TransformControlsHandle]\n\n changed: bool\n \"\"\"This flag will be flipped to True whenever the mesh needs to be re-generated.\"\"\"\n\n\n def make_gui_elements(\n server: viser.ViserServer,\n num_betas: int,\n num_joints: int,\n parent_idx: np.ndarray,\n ) -> GuiElements:\n \"\"\"Make GUI elements for interacting with the model.\"\"\"\n\n tab_group = server.gui.add_tab_group()\n\n def set_changed(_) -> None:\n out.changed = True # out is define later!\n\n # GUI elements: mesh settings + visibility.\n with tab_group.add_tab(\"View\", viser.Icon.VIEWFINDER):\n gui_rgb = server.gui.add_rgb(\"Color\", initial_value=(90, 200, 255))\n gui_wireframe = server.gui.add_checkbox(\"Wireframe\", initial_value=False)\n gui_show_controls = server.gui.add_checkbox(\"Handles\", initial_value=False)\n\n gui_rgb.on_update(set_changed)\n gui_wireframe.on_update(set_changed)\n\n @gui_show_controls.on_update\n def _(_):\n for control in transform_controls:\n control.visible = gui_show_controls.value\n\n # GUI elements: shape parameters.\n with tab_group.add_tab(\"Shape\", viser.Icon.BOX):\n gui_reset_shape = server.gui.add_button(\"Reset Shape\")\n gui_random_shape = server.gui.add_button(\"Random Shape\")\n\n @gui_reset_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = 0.0\n\n @gui_random_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = onp.random.normal(loc=0.0, scale=1.0)\n\n gui_betas = []\n for i in range(num_betas):\n beta = server.gui.add_slider(\n f\"beta{i}\", min=-5.0, max=5.0, step=0.01, initial_value=0.0\n )\n gui_betas.append(beta)\n beta.on_update(set_changed)\n\n # GUI elements: joint angles.\n with tab_group.add_tab(\"Joints\", viser.Icon.ANGLE):\n gui_reset_joints = server.gui.add_button(\"Reset Joints\")\n gui_random_joints = server.gui.add_button(\"Random Joints\")\n\n @gui_reset_joints.on_click\n def _(_):\n for joint in gui_joints:\n joint.value = (0.0, 0.0, 0.0)\n\n @gui_random_joints.on_click\n def _(_):\n for joint in gui_joints:\n # It's hard to uniformly sample orientations directly in so(3), so we\n # first sample on S^3 and then convert.\n quat = onp.random.normal(loc=0.0, scale=1.0, size=(4,))\n quat /= onp.linalg.norm(quat)\n joint.value = tf.SO3(wxyz=quat).log()\n\n gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]] = []\n for i in range(num_joints):\n gui_joint = server.gui.add_vector3(\n label=f\"Joint {i}\",\n initial_value=(0.0, 0.0, 0.0),\n step=0.05,\n )\n gui_joints.append(gui_joint)\n\n def set_callback_in_closure(i: int) -> None:\n @gui_joint.on_update\n def _(_):\n transform_controls[i].wxyz = tf.SO3.exp(\n np.array(gui_joints[i].value)\n ).wxyz\n out.changed = True\n\n set_callback_in_closure(i)\n\n # Transform control gizmos on joints.\n transform_controls: list[viser.TransformControlsHandle] = []\n prefixed_joint_names = [] # Joint names, but prefixed with parents.\n for i in range(num_joints):\n prefixed_joint_name = f\"joint_{i}\"\n if i > 0:\n prefixed_joint_name = (\n prefixed_joint_names[parent_idx[i]] + \"/\" + prefixed_joint_name\n )\n prefixed_joint_names.append(prefixed_joint_name)\n controls = server.scene.add_transform_controls(\n f\"/smpl/{prefixed_joint_name}\",\n depth_test=False,\n scale=0.2 * (0.75 ** prefixed_joint_name.count(\"/\")),\n disable_axes=True,\n disable_sliders=True,\n visible=gui_show_controls.value,\n )\n transform_controls.append(controls)\n\n def set_callback_in_closure(i: int) -> None:\n @controls.on_update\n def _(_) -> None:\n axisangle = tf.SO3(transform_controls[i].wxyz).log()\n gui_joints[i].value = (axisangle[0], axisangle[1], axisangle[2])\n\n set_callback_in_closure(i)\n\n out = GuiElements(\n gui_rgb,\n gui_wireframe,\n gui_betas,\n gui_joints,\n transform_controls=transform_controls,\n changed=True,\n )\n return out\n\n\n if __name__ == \"__main__\":\n tyro.cli(main, description=__doc__)\n" }, { "path": "viser/docs/source/examples/09_urdf_visualizer.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nRobot URDF visualizer\n==========================================\n\n\nRequires yourdfpy and robot_descriptions. Any URDF supported by yourdfpy should work.\n\n\n* https://github.com/robot-descriptions/robot_descriptions.py\n* https://github.com/clemense/yourdfpy\n\nThe :class:`viser.extras.ViserUrdf` is a lightweight interface between yourdfpy\nand viser. It can also take a path to a local URDF file as input.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n from __future__ import annotations\n\n import time\n from typing import Literal\n\n import numpy as onp\n import tyro\n import viser\n from robot_descriptions.loaders.yourdfpy import load_robot_description\n from viser.extras import ViserUrdf\n\n\n def create_robot_control_sliders(\n server: viser.ViserServer, viser_urdf: ViserUrdf\n ) -> tuple[list[viser.GuiInputHandle[float]], list[float]]:\n \"\"\"Create slider for each joint of the robot. We also update robot model\n when slider moves.\"\"\"\n slider_handles: list[viser.GuiInputHandle[float]] = []\n initial_config: list[float] = []\n for joint_name, (\n lower,\n upper,\n ) in viser_urdf.get_actuated_joint_limits().items():\n lower = lower if lower is not None else -onp.pi\n upper = upper if upper is not None else onp.pi\n initial_pos = 0.0 if lower < 0 and upper > 0 else (lower + upper) / 2.0\n slider = server.gui.add_slider(\n label=joint_name,\n min=lower,\n max=upper,\n step=1e-3,\n initial_value=initial_pos,\n )\n slider.on_update( # When sliders move, we update the URDF configuration.\n lambda _: viser_urdf.update_cfg(\n onp.array([slider.value for slider in slider_handles])\n )\n )\n slider_handles.append(slider)\n initial_config.append(initial_pos)\n return slider_handles, initial_config\n\n\n def main(\n robot_type: Literal[\n \"panda\",\n \"ur10\",\n \"cassie\",\n \"allegro_hand\",\n \"barrett_hand\",\n \"robotiq_2f85\",\n \"atlas_drc\",\n \"g1\",\n \"h1\",\n \"anymal_c\",\n \"go2\",\n ] = \"panda\",\n ) -> None:\n # Start viser server.\n server = viser.ViserServer()\n\n # Load URDF.\n #\n # This takes either a yourdfpy.URDF object or a path to a .urdf file.\n viser_urdf = ViserUrdf(\n server,\n urdf_or_path=load_robot_description(robot_type + \"_description\"),\n )\n\n # Create sliders in GUI that help us move the robot joints.\n with server.gui.add_folder(\"Joint position control\"):\n (slider_handles, initial_config) = create_robot_control_sliders(\n server, viser_urdf\n )\n\n # Set initial robot configuration.\n viser_urdf.update_cfg(onp.array(initial_config))\n\n # Create joint reset button.\n reset_button = server.gui.add_button(\"Reset\")\n\n @reset_button.on_click\n def _(_):\n for s, init_q in zip(slider_handles, initial_config):\n s.value = init_q\n\n # Sleep forever.\n while True:\n time.sleep(10.0)\n\n\n if __name__ == \"__main__\":\n tyro.cli(main)\n" }, { "path": "viser/docs/source/examples/10_realsense.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nRealSense visualizer\n==========================================\n\n\nConnect to a RealSense camera, then visualize RGB-D readings as a point clouds. Requires\npyrealsense2.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n from __future__ import annotations\n\n import contextlib\n\n import numpy as np\n import numpy.typing as npt\n import pyrealsense2 as rs # type: ignore\n from tqdm.auto import tqdm\n\n import viser\n\n\n @contextlib.contextmanager\n def realsense_pipeline(fps: int = 30):\n \"\"\"Context manager that yields a RealSense pipeline.\"\"\"\n\n # Configure depth and color streams.\n pipeline = rs.pipeline() # type: ignore\n config = rs.config() # type: ignore\n\n pipeline_wrapper = rs.pipeline_wrapper(pipeline) # type: ignore\n config.resolve(pipeline_wrapper)\n\n config.enable_stream(rs.stream.depth, rs.format.z16, fps) # type: ignore\n config.enable_stream(rs.stream.color, rs.format.rgb8, fps) # type: ignore\n\n # Start streaming.\n pipeline.start(config)\n\n yield pipeline\n\n # Close pipeline when done.\n pipeline.close()\n\n\n def point_cloud_arrays_from_frames(\n depth_frame, color_frame\n ) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.uint8]]:\n \"\"\"Maps realsense frames to two arrays.\n\n Returns:\n - A point position array: (N, 3) float32.\n - A point color array: (N, 3) uint8.\n \"\"\"\n # Processing blocks. Could be tuned.\n point_cloud = rs.pointcloud() # type: ignore\n decimate = rs.decimation_filter() # type: ignore\n decimate.set_option(rs.option.filter_magnitude, 3) # type: ignore\n\n # Downsample depth frame.\n depth_frame = decimate.process(depth_frame)\n\n # Map texture and calculate points from frames. Uses frame intrinsics.\n point_cloud.map_to(color_frame)\n points = point_cloud.calculate(depth_frame)\n\n # Get color coordinates.\n texture_uv = (\n np.asanyarray(points.get_texture_coordinates())\n .view(np.float32)\n .reshape((-1, 2))\n )\n color_image = np.asanyarray(color_frame.get_data())\n color_h, color_w, _ = color_image.shape\n\n # Note: for points that aren't in the view of our RGB camera, we currently clamp to\n # the closes available RGB pixel. We could also just remove these points.\n texture_uv = texture_uv.clip(0.0, 1.0)\n\n # Get positions and colors.\n positions = np.asanyarray(points.get_vertices()).view(np.float32)\n positions = positions.reshape((-1, 3))\n colors = color_image[\n (texture_uv[:, 1] * (color_h - 1.0)).astype(np.int32),\n (texture_uv[:, 0] * (color_w - 1.0)).astype(np.int32),\n :,\n ]\n N = positions.shape[0]\n\n assert positions.shape == (N, 3)\n assert positions.dtype == np.float32\n assert colors.shape == (N, 3)\n assert colors.dtype == np.uint8\n\n return positions, colors\n\n\n def main():\n # Start visualization server.\n server = viser.ViserServer()\n\n with realsense_pipeline() as pipeline:\n for i in tqdm(range(10000000)):\n # Wait for a coherent pair of frames: depth and color\n frames = pipeline.wait_for_frames()\n depth_frame = frames.get_depth_frame()\n color_frame = frames.get_color_frame()\n\n # Compute point cloud from frames.\n positions, colors = point_cloud_arrays_from_frames(depth_frame, color_frame)\n\n R = np.array(\n [\n [1.0, 0.0, 0.0],\n [0.0, 0.0, 1.0],\n [0.0, -1.0, 0.0],\n ],\n dtype=np.float32,\n )\n positions = positions @ R.T\n\n # Visualize.\n server.scene.add_point_cloud(\n \"/realsense\",\n points=positions * 10.0,\n colors=colors,\n point_size=0.1,\n )\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/11_colmap_visualizer.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nCOLMAP visualizer\n==========================================\n\n\nVisualize COLMAP sparse reconstruction outputs. To get demo data, see ``./assets/download_colmap_garden.sh``.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import random\n import time\n from pathlib import Path\n\n import imageio.v3 as iio\n import numpy as onp\n import tyro\n from tqdm.auto import tqdm\n\n import viser\n import viser.transforms as tf\n from viser.extras.colmap import (\n read_cameras_binary,\n read_images_binary,\n read_points3d_binary,\n )\n\n\n def main(\n colmap_path: Path = Path(__file__).parent / \"assets/colmap_garden/sparse/0\",\n images_path: Path = Path(__file__).parent / \"assets/colmap_garden/images_8\",\n downsample_factor: int = 2,\n ) -> None:\n \"\"\"Visualize COLMAP sparse reconstruction outputs.\n\n Args:\n colmap_path: Path to the COLMAP reconstruction directory.\n images_path: Path to the COLMAP images directory.\n downsample_factor: Downsample factor for the images.\n \"\"\"\n server = viser.ViserServer()\n server.gui.configure_theme(titlebar_content=None, control_layout=\"collapsible\")\n\n # Load the colmap info.\n cameras = read_cameras_binary(colmap_path / \"cameras.bin\")\n images = read_images_binary(colmap_path / \"images.bin\")\n points3d = read_points3d_binary(colmap_path / \"points3D.bin\")\n gui_reset_up = server.gui.add_button(\n \"Reset up direction\",\n hint=\"Set the camera control 'up' direction to the current camera's 'up'.\",\n )\n\n @gui_reset_up.on_click\n def _(event: viser.GuiEvent) -> None:\n client = event.client\n assert client is not None\n client.camera.up_direction = tf.SO3(client.camera.wxyz) @ onp.array(\n [0.0, -1.0, 0.0]\n )\n\n gui_points = server.gui.add_slider(\n \"Max points\",\n min=1,\n max=len(points3d),\n step=1,\n initial_value=min(len(points3d), 50_000),\n )\n gui_frames = server.gui.add_slider(\n \"Max frames\",\n min=1,\n max=len(images),\n step=1,\n initial_value=min(len(images), 100),\n )\n gui_point_size = server.gui.add_number(\"Point size\", initial_value=0.05)\n\n def visualize_colmap() -> None:\n \"\"\"Send all COLMAP elements to viser for visualization. This could be optimized\n a ton!\"\"\"\n # Set the point cloud.\n points = onp.array([points3d[p_id].xyz for p_id in points3d])\n colors = onp.array([points3d[p_id].rgb for p_id in points3d])\n points_selection = onp.random.choice(\n points.shape[0], gui_points.value, replace=False\n )\n points = points[points_selection]\n colors = colors[points_selection]\n\n server.scene.add_point_cloud(\n name=\"/colmap/pcd\",\n points=points,\n colors=colors,\n point_size=gui_point_size.value,\n )\n\n # Interpret the images and cameras.\n img_ids = [im.id for im in images.values()]\n random.shuffle(img_ids)\n img_ids = sorted(img_ids[: gui_frames.value])\n\n def attach_callback(\n frustum: viser.CameraFrustumHandle, frame: viser.FrameHandle\n ) -> None:\n @frustum.on_click\n def _(_) -> None:\n for client in server.get_clients().values():\n client.camera.wxyz = frame.wxyz\n client.camera.position = frame.position\n\n for img_id in tqdm(img_ids):\n img = images[img_id]\n cam = cameras[img.camera_id]\n\n # Skip images that don't exist.\n image_filename = images_path / img.name\n if not image_filename.exists():\n continue\n\n T_world_camera = tf.SE3.from_rotation_and_translation(\n tf.SO3(img.qvec), img.tvec\n ).inverse()\n frame = server.scene.add_frame(\n f\"/colmap/frame_{img_id}\",\n wxyz=T_world_camera.rotation().wxyz,\n position=T_world_camera.translation(),\n axes_length=0.1,\n axes_radius=0.005,\n )\n\n # For pinhole cameras, cam.params will be (fx, fy, cx, cy).\n if cam.model != \"PINHOLE\":\n print(f\"Expected pinhole camera, but got {cam.model}\")\n\n H, W = cam.height, cam.width\n fy = cam.params[1]\n image = iio.imread(image_filename)\n image = image[::downsample_factor, ::downsample_factor]\n frustum = server.scene.add_camera_frustum(\n f\"/colmap/frame_{img_id}/frustum\",\n fov=2 * onp.arctan2(H / 2, fy),\n aspect=W / H,\n scale=0.15,\n image=image,\n )\n attach_callback(frustum, frame)\n\n need_update = True\n\n @gui_points.on_update\n def _(_) -> None:\n nonlocal need_update\n need_update = True\n\n @gui_frames.on_update\n def _(_) -> None:\n nonlocal need_update\n need_update = True\n\n @gui_point_size.on_update\n def _(_) -> None:\n nonlocal need_update\n need_update = True\n\n while True:\n if need_update:\n need_update = False\n\n server.scene.reset()\n visualize_colmap()\n\n time.sleep(1e-3)\n\n\n if __name__ == \"__main__\":\n tyro.cli(main)\n" }, { "path": "viser/docs/source/examples/12_click_meshes.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nMesh click events\n==========================================\n\n\nClick on meshes to select them. The index of the last clicked mesh is displayed in the GUI.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import matplotlib\n\n import viser\n\n\n def main() -> None:\n grid_shape = (4, 5)\n server = viser.ViserServer()\n\n with server.gui.add_folder(\"Last clicked\"):\n x_value = server.gui.add_number(\n label=\"x\",\n initial_value=0,\n disabled=True,\n hint=\"x coordinate of the last clicked mesh\",\n )\n y_value = server.gui.add_number(\n label=\"y\",\n initial_value=0,\n disabled=True,\n hint=\"y coordinate of the last clicked mesh\",\n )\n\n def add_swappable_mesh(i: int, j: int) -> None:\n \"\"\"Simple callback that swaps between:\n - a gray box\n - a colored box\n - a colored sphere\n\n Color is chosen based on the position (i, j) of the mesh in the grid.\n \"\"\"\n\n colormap = matplotlib.colormaps[\"tab20\"]\n\n def create_mesh(counter: int) -> None:\n if counter == 0:\n color = (0.8, 0.8, 0.8)\n else:\n index = (i * grid_shape[1] + j) / (grid_shape[0] * grid_shape[1])\n color = colormap(index)[:3]\n\n if counter in (0, 1):\n handle = server.scene.add_box(\n name=f\"/sphere_{i}_{j}\",\n position=(i, j, 0.0),\n color=color,\n dimensions=(0.5, 0.5, 0.5),\n )\n else:\n handle = server.scene.add_icosphere(\n name=f\"/sphere_{i}_{j}\",\n radius=0.4,\n color=color,\n position=(i, j, 0.0),\n )\n\n @handle.on_click\n def _(_) -> None:\n x_value.value = i\n y_value.value = j\n\n # The new mesh will replace the old one because the names\n # /sphere_{i}_{j} are the same.\n create_mesh((counter + 1) % 3)\n\n create_mesh(0)\n\n for i in range(grid_shape[0]):\n for j in range(grid_shape[1]):\n add_swappable_mesh(i, j)\n\n while True:\n time.sleep(10.0)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/13_theming.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nTheming\n==========================================\n\n\nViser includes support for light theming.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import viser\n from viser.theme import TitlebarButton, TitlebarConfig, TitlebarImage\n\n\n def main():\n server = viser.ViserServer(label=\"Viser Theming\")\n\n buttons = (\n TitlebarButton(\n text=\"Getting Started\",\n icon=None,\n href=\"https://nerf.studio\",\n ),\n TitlebarButton(\n text=\"Github\",\n icon=\"GitHub\",\n href=\"https://github.com/nerfstudio-project/nerfstudio\",\n ),\n TitlebarButton(\n text=\"Documentation\",\n icon=\"Description\",\n href=\"https://docs.nerf.studio\",\n ),\n )\n image = TitlebarImage(\n image_url_light=\"https://docs.nerf.studio/_static/imgs/logo.png\",\n image_url_dark=\"https://docs.nerf.studio/_static/imgs/logo-dark.png\",\n image_alt=\"NerfStudio Logo\",\n href=\"https://docs.nerf.studio/\",\n )\n titlebar_theme = TitlebarConfig(buttons=buttons, image=image)\n\n server.gui.add_markdown(\n \"Viser includes support for light theming via the `.configure_theme()` method.\"\n )\n\n gui_theme_code = server.gui.add_markdown(\"no theme applied yet\")\n\n # GUI elements for controllable values.\n titlebar = server.gui.add_checkbox(\"Titlebar\", initial_value=True)\n dark_mode = server.gui.add_checkbox(\"Dark mode\", initial_value=True)\n show_logo = server.gui.add_checkbox(\"Show logo\", initial_value=True)\n show_share_button = server.gui.add_checkbox(\"Show share button\", initial_value=True)\n brand_color = server.gui.add_rgb(\"Brand color\", (230, 180, 30))\n control_layout = server.gui.add_dropdown(\n \"Control layout\", (\"floating\", \"fixed\", \"collapsible\")\n )\n control_width = server.gui.add_dropdown(\n \"Control width\", (\"small\", \"medium\", \"large\"), initial_value=\"medium\"\n )\n synchronize = server.gui.add_button(\"Apply theme\", icon=viser.Icon.CHECK)\n\n def synchronize_theme() -> None:\n server.gui.configure_theme(\n titlebar_content=titlebar_theme if titlebar.value else None,\n control_layout=control_layout.value,\n control_width=control_width.value,\n dark_mode=dark_mode.value,\n show_logo=show_logo.value,\n show_share_button=show_share_button.value,\n brand_color=brand_color.value,\n )\n gui_theme_code.content = f\"\"\"\n ### Current applied theme\n ```\n server.gui.configure_theme(\n titlebar_content={\"titlebar_content\" if titlebar.value else None},\n control_layout=\"{control_layout.value}\",\n control_width=\"{control_width.value}\",\n dark_mode={dark_mode.value},\n show_logo={show_logo.value},\n show_share_button={show_share_button.value},\n brand_color={brand_color.value},\n )\n ```\n \"\"\"\n\n synchronize.on_click(lambda _: synchronize_theme())\n synchronize_theme()\n\n while True:\n time.sleep(10.0)\n\n\n # main()\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/14_markdown.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nMarkdown demonstration\n==========================================\n\n\nViser GUI has MDX 2 support.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n from pathlib import Path\n\n import viser\n\n server = viser.ViserServer()\n server.scene.world_axes.visible = True\n\n markdown_counter = server.gui.add_markdown(\"Counter: 0\")\n\n here = Path(__file__).absolute().parent\n\n button = server.gui.add_button(\"Remove blurb\")\n checkbox = server.gui.add_checkbox(\"Visibility\", initial_value=True)\n\n markdown_source = (here / \"./assets/mdx_example.mdx\").read_text()\n markdown_blurb = server.gui.add_markdown(\n content=markdown_source,\n image_root=here,\n )\n\n\n @button.on_click\n def _(_):\n markdown_blurb.remove()\n\n\n @checkbox.on_update\n def _(_):\n markdown_blurb.visible = checkbox.value\n\n\n counter = 0\n while True:\n markdown_counter.content = f\"Counter: {counter}\"\n counter += 1\n time.sleep(0.1)\n" }, { "path": "viser/docs/source/examples/15_gui_in_scene.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\n3D GUI elements\n==========================================\n\n\n``add_3d_gui_container()`` allows standard GUI elements to be incorporated directly into a\n3D scene. In this example, we click on coordinate frames to show actions that can be\nperformed on them.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n from typing import Optional\n\n import numpy as onp\n\n import viser\n import viser.transforms as tf\n\n server = viser.ViserServer()\n server.gui.configure_theme(dark_mode=True)\n num_frames = 20\n\n\n @server.on_client_connect\n def _(client: viser.ClientHandle) -> None:\n \"\"\"For each client that connects, we create a set of random frames + a click handler for each frame.\n\n When a frame is clicked, we display a 3D gui node.\n \"\"\"\n\n rng = onp.random.default_rng(0)\n\n displayed_3d_container: Optional[viser.Gui3dContainerHandle] = None\n\n def make_frame(i: int) -> None:\n # Sample a random orientation + position.\n wxyz = rng.normal(size=4)\n wxyz /= onp.linalg.norm(wxyz)\n position = rng.uniform(-3.0, 3.0, size=(3,))\n\n # Create a coordinate frame and label.\n frame = client.scene.add_frame(f\"/frame_{i}\", wxyz=wxyz, position=position)\n\n # Move the camera when we click a frame.\n @frame.on_click\n def _(_):\n nonlocal displayed_3d_container\n\n # Close previously opened GUI.\n if displayed_3d_container is not None:\n displayed_3d_container.remove()\n\n displayed_3d_container = client.scene.add_3d_gui_container(\n f\"/frame_{i}/gui\"\n )\n with displayed_3d_container:\n go_to = client.gui.add_button(\"Go to\")\n randomize_orientation = client.gui.add_button(\"Randomize orientation\")\n close = client.gui.add_button(\"Close GUI\")\n\n @go_to.on_click\n def _(_) -> None:\n T_world_current = tf.SE3.from_rotation_and_translation(\n tf.SO3(client.camera.wxyz), client.camera.position\n )\n T_world_target = tf.SE3.from_rotation_and_translation(\n tf.SO3(frame.wxyz), frame.position\n ) @ tf.SE3.from_translation(onp.array([0.0, 0.0, -0.5]))\n\n T_current_target = T_world_current.inverse() @ T_world_target\n\n for j in range(20):\n T_world_set = T_world_current @ tf.SE3.exp(\n T_current_target.log() * j / 19.0\n )\n\n # Important bit: we atomically set both the orientation and the position\n # of the camera.\n with client.atomic():\n client.camera.wxyz = T_world_set.rotation().wxyz\n client.camera.position = T_world_set.translation()\n time.sleep(1.0 / 60.0)\n\n # Mouse interactions should orbit around the frame origin.\n client.camera.look_at = frame.position\n\n @randomize_orientation.on_click\n def _(_) -> None:\n wxyz = rng.normal(size=4)\n wxyz /= onp.linalg.norm(wxyz)\n frame.wxyz = wxyz\n\n @close.on_click\n def _(_) -> None:\n nonlocal displayed_3d_container\n if displayed_3d_container is None:\n return\n displayed_3d_container.remove()\n displayed_3d_container = None\n\n for i in range(num_frames):\n make_frame(i)\n\n\n while True:\n time.sleep(1.0)\n" }, { "path": "viser/docs/source/examples/16_modal.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nModal basics\n==========================================\n\n\nExamples of using modals in Viser.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import viser\n\n\n def main():\n server = viser.ViserServer()\n\n @server.on_client_connect\n def _(client: viser.ClientHandle) -> None:\n with client.gui.add_modal(\"Modal example\"):\n client.gui.add_markdown(\n \"**The input below determines the title of the modal...**\"\n )\n\n gui_title = client.gui.add_text(\n \"Title\",\n initial_value=\"My Modal\",\n )\n\n modal_button = client.gui.add_button(\"Show more modals\")\n\n @modal_button.on_click\n def _(_) -> None:\n with client.gui.add_modal(gui_title.value) as modal:\n client.gui.add_markdown(\"This is content inside the modal!\")\n client.gui.add_button(\"Close\").on_click(lambda _: modal.close())\n\n while True:\n time.sleep(0.15)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/17_background_composite.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nDepth compositing\n==========================================\n\n\nIn this example, we show how to use a background image with depth compositing. This can\nbe useful when we want a 2D image to occlude 3D geometry, such as for NeRF rendering.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import numpy as onp\n import trimesh\n import trimesh.creation\n\n import viser\n\n server = viser.ViserServer()\n\n\n img = onp.random.randint(0, 255, size=(1000, 1000, 3), dtype=onp.uint8)\n depth = onp.ones((1000, 1000, 1), dtype=onp.float32)\n\n # Make a square middle portal.\n depth[250:750, 250:750, :] = 10.0\n img[250:750, 250:750, :] = 255\n\n mesh = trimesh.creation.box((0.5, 0.5, 0.5))\n server.scene.add_mesh_trimesh(\n name=\"/cube\",\n mesh=mesh,\n position=(0, 0, 0.0),\n )\n server.scene.set_background_image(img, depth=depth)\n\n\n while True:\n time.sleep(1.0)\n" }, { "path": "viser/docs/source/examples/18_splines.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nSplines\n==========================================\n\n\nMake a ball with some random splines.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import numpy as onp\n\n import viser\n\n\n def main() -> None:\n server = viser.ViserServer()\n for i in range(10):\n positions = onp.random.normal(size=(30, 3)) * 3.0\n server.scene.add_spline_catmull_rom(\n f\"/catmull_{i}\",\n positions,\n tension=0.5,\n line_width=3.0,\n color=onp.random.uniform(size=3),\n segments=100,\n )\n\n control_points = onp.random.normal(size=(30 * 2 - 2, 3)) * 3.0\n server.scene.add_spline_cubic_bezier(\n f\"/cubic_bezier_{i}\",\n positions,\n control_points,\n line_width=3.0,\n color=onp.random.uniform(size=3),\n segments=100,\n )\n\n while True:\n time.sleep(10.0)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/19_get_renders.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nGet renders\n==========================================\n\n\nExample for getting renders from a client's viewport to the Python API.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import imageio.v3 as iio\n import numpy as onp\n\n import viser\n\n\n def main():\n server = viser.ViserServer()\n\n button = server.gui.add_button(\"Render a GIF\")\n\n @button.on_click\n def _(event: viser.GuiEvent) -> None:\n client = event.client\n assert client is not None\n\n client.scene.reset()\n\n images = []\n\n for i in range(20):\n positions = onp.random.normal(size=(30, 3)) * 3.0\n client.scene.add_spline_catmull_rom(\n f\"/catmull_{i}\",\n positions,\n tension=0.5,\n line_width=3.0,\n color=onp.random.uniform(size=3),\n )\n images.append(client.camera.get_render(height=720, width=1280))\n\n print(\"Generating and sending GIF...\")\n client.send_file_download(\n \"image.gif\", iio.imwrite(\"\", images, extension=\".gif\")\n )\n print(\"Done!\")\n\n while True:\n time.sleep(10.0)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/20_scene_pointer.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nScene pointer events.\n==========================================\n\n\nThis example shows how to use scene pointer events to specify rays, and how they can be\nused to interact with the scene (e.g., ray-mesh intersections).\n\nTo get the demo data, see ``./assets/download_dragon_mesh.sh``.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n from __future__ import annotations\n\n import time\n from pathlib import Path\n from typing import cast\n\n import numpy as onp\n import trimesh\n import trimesh.creation\n import trimesh.ray\n\n import viser\n import viser.transforms as tf\n from viser.theme import TitlebarConfig\n\n server = viser.ViserServer()\n server.gui.configure_theme(\n brand_color=(130, 0, 150),\n titlebar_content=TitlebarConfig(buttons=(), image=None),\n )\n server.scene.set_up_direction(\"+y\")\n\n mesh = cast(\n trimesh.Trimesh, trimesh.load_mesh(str(Path(__file__).parent / \"assets/dragon.obj\"))\n )\n mesh.apply_scale(0.05)\n\n mesh_handle = server.scene.add_mesh_trimesh(\n name=\"/mesh\",\n mesh=mesh,\n position=(0.0, 0.0, 0.0),\n )\n\n hit_pos_handles: list[viser.GlbHandle] = []\n\n\n # Buttons + callbacks will operate on a per-client basis, but will modify the global scene! :)\n @server.on_client_connect\n def _(client: viser.ClientHandle) -> None:\n # Set up the camera -- this gives a nice view of the full mesh.\n client.camera.position = onp.array([0.0, 0.0, -10.0])\n client.camera.wxyz = onp.array([0.0, 0.0, 0.0, 1.0])\n\n # Tests \"click\" scenepointerevent.\n click_button_handle = client.gui.add_button(\"Add sphere\", icon=viser.Icon.POINTER)\n\n @click_button_handle.on_click\n def _(_):\n click_button_handle.disabled = True\n\n @client.scene.on_pointer_event(event_type=\"click\")\n def _(event: viser.ScenePointerEvent) -> None:\n # Check for intersection with the mesh, using trimesh's ray-mesh intersection.\n # Note that mesh is in the mesh frame, so we need to transform the ray.\n R_world_mesh = tf.SO3(mesh_handle.wxyz)\n R_mesh_world = R_world_mesh.inverse()\n origin = (R_mesh_world @ onp.array(event.ray_origin)).reshape(1, 3)\n direction = (R_mesh_world @ onp.array(event.ray_direction)).reshape(1, 3)\n intersector = trimesh.ray.ray_triangle.RayMeshIntersector(mesh)\n hit_pos, _, _ = intersector.intersects_location(origin, direction)\n\n if len(hit_pos) == 0:\n return\n client.scene.remove_pointer_callback()\n\n # Get the first hit position (based on distance from the ray origin).\n hit_pos = hit_pos[onp.argmin(onp.sum((hit_pos - origin) ** 2, axis=-1))]\n\n # Create a sphere at the hit location.\n hit_pos_mesh = trimesh.creation.icosphere(radius=0.1)\n hit_pos_mesh.vertices += R_world_mesh @ hit_pos\n hit_pos_mesh.visual.vertex_colors = (0.5, 0.0, 0.7, 1.0) # type: ignore\n hit_pos_handle = server.scene.add_mesh_trimesh(\n name=f\"/hit_pos_{len(hit_pos_handles)}\", mesh=hit_pos_mesh\n )\n hit_pos_handles.append(hit_pos_handle)\n\n @client.scene.on_pointer_callback_removed\n def _():\n click_button_handle.disabled = False\n\n # Tests \"rect-select\" scenepointerevent.\n paint_button_handle = client.gui.add_button(\"Paint mesh\", icon=viser.Icon.PAINT)\n\n @paint_button_handle.on_click\n def _(_):\n paint_button_handle.disabled = True\n\n @client.scene.on_pointer_event(event_type=\"rect-select\")\n def _(message: viser.ScenePointerEvent) -> None:\n client.scene.remove_pointer_callback()\n\n global mesh_handle\n camera = message.client.camera\n\n # Put the mesh in the camera frame.\n R_world_mesh = tf.SO3(mesh_handle.wxyz)\n R_mesh_world = R_world_mesh.inverse()\n R_camera_world = tf.SE3.from_rotation_and_translation(\n tf.SO3(camera.wxyz), camera.position\n ).inverse()\n vertices = cast(onp.ndarray, mesh.vertices)\n vertices = (R_mesh_world.as_matrix() @ vertices.T).T\n vertices = (\n R_camera_world.as_matrix()\n @ onp.hstack([vertices, onp.ones((vertices.shape[0], 1))]).T\n ).T[:, :3]\n\n # Get the camera intrinsics, and project the vertices onto the image plane.\n fov, aspect = camera.fov, camera.aspect\n vertices_proj = vertices[:, :2] / vertices[:, 2].reshape(-1, 1)\n vertices_proj /= onp.tan(fov / 2)\n vertices_proj[:, 0] /= aspect\n\n # Move the origin to the upper-left corner, and scale to [0, 1].\n # ... make sure to match the OpenCV's image coordinates!\n vertices_proj = (1 + vertices_proj) / 2\n\n # Select the vertices that lie inside the 2D selected box, once projected.\n mask = (\n (vertices_proj > onp.array(message.screen_pos[0]))\n & (vertices_proj < onp.array(message.screen_pos[1]))\n ).all(axis=1)[..., None]\n\n # Update the mesh color based on whether the vertices are inside the box\n mesh.visual.vertex_colors = onp.where( # type: ignore\n mask, (0.5, 0.0, 0.7, 1.0), (0.9, 0.9, 0.9, 1.0)\n )\n mesh_handle = server.scene.add_mesh_trimesh(\n name=\"/mesh\",\n mesh=mesh,\n position=(0.0, 0.0, 0.0),\n )\n\n @client.scene.on_pointer_callback_removed\n def _():\n paint_button_handle.disabled = False\n\n # Button to clear spheres.\n clear_button_handle = client.gui.add_button(\"Clear scene\", icon=viser.Icon.X)\n\n @clear_button_handle.on_click\n def _(_):\n \"\"\"Reset the mesh color and remove all click-generated spheres.\"\"\"\n global mesh_handle\n for handle in hit_pos_handles:\n handle.remove()\n hit_pos_handles.clear()\n mesh.visual.vertex_colors = (0.9, 0.9, 0.9, 1.0) # type: ignore\n mesh_handle = server.scene.add_mesh_trimesh(\n name=\"/mesh\",\n mesh=mesh,\n position=(0.0, 0.0, 0.0),\n )\n\n\n while True:\n time.sleep(10.0)\n" }, { "path": "viser/docs/source/examples/21_set_up_direction.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nSet up direction\n==========================================\n\n\n``.set_up_direction()`` can help us set the global up direction.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import viser\n\n\n def main() -> None:\n server = viser.ViserServer()\n server.scene.world_axes.visible = True\n gui_up = server.gui.add_vector3(\n \"Up Direction\",\n initial_value=(0.0, 0.0, 1.0),\n step=0.01,\n )\n\n @gui_up.on_update\n def _(_) -> None:\n server.scene.set_up_direction(gui_up.value)\n\n while True:\n time.sleep(1.0)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/22_games.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nGames\n==========================================\n\n\nSome two-player games implemented using scene click events.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n from typing import Literal\n\n import numpy as onp\n import trimesh.creation\n from typing_extensions import assert_never\n\n import viser\n import viser.transforms as tf\n\n\n def main() -> None:\n server = viser.ViserServer()\n server.gui.configure_theme(dark_mode=True)\n play_connect_4(server)\n\n server.gui.add_button(\"Tic-Tac-Toe\").on_click(lambda _: play_tic_tac_toe(server))\n server.gui.add_button(\"Connect 4\").on_click(lambda _: play_connect_4(server))\n\n while True:\n time.sleep(10.0)\n\n\n def play_connect_4(server: viser.ViserServer) -> None:\n \"\"\"Play a game of Connect 4.\"\"\"\n server.scene.reset()\n\n num_rows = 6\n num_cols = 7\n\n whose_turn: Literal[\"red\", \"yellow\"] = \"red\"\n pieces_in_col = [0] * num_cols\n\n # Create the board frame.\n for col in range(num_cols):\n for row in range(num_rows):\n server.scene.add_mesh_trimesh(\n f\"/structure/{row}_{col}\",\n trimesh.creation.annulus(0.45, 0.55, 0.125),\n position=(0.0, col, row),\n wxyz=tf.SO3.from_y_radians(onp.pi / 2.0).wxyz,\n )\n\n # Create a sphere to click on for each column.\n def setup_column(col: int) -> None:\n sphere = server.scene.add_icosphere(\n f\"/spheres/{col}\",\n radius=0.25,\n position=(0, col, num_rows - 0.25),\n color=(255, 255, 255),\n )\n\n # Drop piece into the column.\n @sphere.on_click\n def _(_) -> None:\n nonlocal whose_turn\n whose_turn = \"red\" if whose_turn != \"red\" else \"yellow\"\n\n row = pieces_in_col[col]\n if row == num_rows - 1:\n sphere.remove()\n\n pieces_in_col[col] += 1\n cylinder = trimesh.creation.cylinder(radius=0.4, height=0.125)\n piece = server.scene.add_mesh_simple(\n f\"/game_pieces/{row}_{col}\",\n cylinder.vertices,\n cylinder.faces,\n wxyz=tf.SO3.from_y_radians(onp.pi / 2.0).wxyz,\n color={\"red\": (255, 0, 0), \"yellow\": (255, 255, 0)}[whose_turn],\n )\n for row_anim in onp.linspace(num_rows - 1, row, num_rows - row + 1):\n piece.position = (\n 0,\n col,\n row_anim,\n )\n time.sleep(1.0 / 30.0)\n\n for col in range(num_cols):\n setup_column(col)\n\n\n def play_tic_tac_toe(server: viser.ViserServer) -> None:\n \"\"\"Play a game of tic-tac-toe.\"\"\"\n server.scene.reset()\n\n whose_turn: Literal[\"x\", \"o\"] = \"x\"\n\n for i in range(4):\n server.scene.add_spline_catmull_rom(\n f\"/gridlines/{i}\",\n ((-0.5, -1.5, 0), (-0.5, 1.5, 0)),\n color=(127, 127, 127),\n position=(1, 1, 0),\n wxyz=tf.SO3.from_z_radians(onp.pi / 2 * i).wxyz,\n )\n\n def draw_symbol(symbol: Literal[\"x\", \"o\"], i: int, j: int) -> None:\n \"\"\"Draw an X or O in the given cell.\"\"\"\n for scale in onp.linspace(0.01, 1.0, 5):\n if symbol == \"x\":\n for k in range(2):\n server.scene.add_box(\n f\"/symbols/{i}_{j}/{k}\",\n dimensions=(0.7 * scale, 0.125 * scale, 0.125),\n position=(i, j, 0),\n color=(0, 0, 255),\n wxyz=tf.SO3.from_z_radians(\n onp.pi / 2.0 * k + onp.pi / 4.0\n ).wxyz,\n )\n elif symbol == \"o\":\n mesh = trimesh.creation.annulus(0.25 * scale, 0.35 * scale, 0.125)\n server.scene.add_mesh_simple(\n f\"/symbols/{i}_{j}\",\n mesh.vertices,\n mesh.faces,\n position=(i, j, 0),\n color=(255, 0, 0),\n )\n else:\n assert_never(symbol)\n server.flush()\n time.sleep(1.0 / 30.0)\n\n def setup_cell(i: int, j: int) -> None:\n \"\"\"Create a clickable sphere in a given cell.\"\"\"\n sphere = server.scene.add_icosphere(\n f\"/spheres/{i}_{j}\",\n radius=0.25,\n position=(i, j, 0),\n color=(255, 255, 255),\n )\n\n @sphere.on_click\n def _(_) -> None:\n nonlocal whose_turn\n whose_turn = \"x\" if whose_turn != \"x\" else \"o\"\n sphere.remove()\n draw_symbol(whose_turn, i, j)\n\n for i in range(3):\n for j in range(3):\n setup_cell(i, j)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/23_plotly.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nPlotly\n==========================================\n\n\nExamples of visualizing plotly plots in Viser.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import numpy as onp\n import plotly.express as px\n import plotly.graph_objects as go\n from PIL import Image\n\n import viser\n\n\n def create_sinusoidal_wave(t: float) -> go.Figure:\n \"\"\"Create a sinusoidal wave plot, starting at time t.\"\"\"\n x_data = onp.linspace(t, t + 6 * onp.pi, 50)\n y_data = onp.sin(x_data) * 10\n\n fig = px.line(\n x=list(x_data),\n y=list(y_data),\n labels={\"x\": \"x\", \"y\": \"sin(x)\"},\n title=\"Sinusoidal Wave\",\n )\n\n # this sets the margins to be tight around the title.\n fig.layout.title.automargin = True # type: ignore\n fig.update_layout(\n margin=dict(l=20, r=20, t=20, b=20),\n ) # Reduce plot margins.\n\n return fig\n\n\n def main() -> None:\n server = viser.ViserServer()\n\n # Plot type 1: Line plot.\n line_plot_time = 0.0\n line_plot = server.gui.add_plotly(figure=create_sinusoidal_wave(line_plot_time))\n\n # Plot type 2: Image plot.\n fig = px.imshow(Image.open(\"assets/Cal_logo.png\"))\n fig.update_layout(\n margin=dict(l=20, r=20, t=20, b=20),\n )\n server.gui.add_plotly(figure=fig, aspect=1.0)\n\n # Plot type 3: 3D Scatter plot.\n fig = px.scatter_3d(\n px.data.iris(),\n x=\"sepal_length\",\n y=\"sepal_width\",\n z=\"petal_width\",\n color=\"species\",\n )\n fig.update_layout(legend=dict(yanchor=\"top\", y=0.99, xanchor=\"left\", x=0.01))\n fig.update_layout(\n margin=dict(l=20, r=20, t=20, b=20),\n )\n server.gui.add_plotly(figure=fig, aspect=1.0)\n\n while True:\n # Update the line plot.\n line_plot_time += 0.1\n line_plot.figure = create_sinusoidal_wave(line_plot_time)\n\n time.sleep(0.01)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/24_notification.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nNotifications\n==========================================\n\n\nExamples of adding notifications per client in Viser.\n\n\n\n.. code-block:: python\n :linenos:\n\n\n import time\n\n import viser\n\n\n def main() -> None:\n server = viser.ViserServer()\n\n persistent_notif_button = server.gui.add_button(\n \"Show persistent notification (default)\"\n )\n timed_notif_button = server.gui.add_button(\"Show timed notification\")\n controlled_notif_button = server.gui.add_button(\"Show controlled notification\")\n loading_notif_button = server.gui.add_button(\"Show loading notification\")\n\n remove_controlled_notif = server.gui.add_button(\"Remove controlled notification\")\n\n @persistent_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show persistent notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n client.add_notification(\n title=\"Persistent notification\",\n body=\"This can be closed manually and does not disappear on its own!\",\n loading=False,\n with_close_button=True,\n auto_close=False,\n )\n\n @timed_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show timed notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n client.add_notification(\n title=\"Timed notification\",\n body=\"This disappears automatically after 5 seconds!\",\n loading=False,\n with_close_button=True,\n auto_close=5000,\n )\n\n @controlled_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show controlled notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n controlled_notif = client.add_notification(\n title=\"Controlled notification\",\n body=\"This cannot be closed by the user and is controlled in code only!\",\n loading=False,\n with_close_button=False,\n auto_close=False,\n )\n\n @remove_controlled_notif.on_click\n def _(_) -> None:\n \"\"\"Remove controlled notification.\"\"\"\n controlled_notif.remove()\n\n @loading_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show loading notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n loading_notif = client.add_notification(\n title=\"Loading notification\",\n body=\"This indicates that some action is in progress! It will be updated in 3 seconds.\",\n loading=True,\n with_close_button=False,\n auto_close=False,\n )\n\n time.sleep(3.0)\n\n loading_notif.title = \"Updated notification\"\n loading_notif.body = \"This notification has been updated!\"\n loading_notif.loading = False\n loading_notif.with_close_button = True\n loading_notif.auto_close = 5000\n loading_notif.color = \"green\"\n\n while True:\n time.sleep(1.0)\n\n\n if __name__ == \"__main__\":\n main()\n" }, { "path": "viser/docs/source/examples/25_smpl_visualizer_skinned.rst", "content": ".. Comment: this file is automatically generated by `update_example_docs.py`.\n It should not be modified manually.\n\nSMPL visualizer (Skinned Mesh)\n==========================================\n\n\nRequires a .npz model file.\n\nSee here for download instructions:\n https://github.com/vchoutas/smplx?tab=readme-ov-file#downloading-the-model\n\n\n\n.. code-block:: python\n :linenos:\n\n\n from __future__ import annotations\n\n import time\n from dataclasses import dataclass\n from pathlib import Path\n from typing import List, Tuple\n\n import numpy as np\n import numpy as onp\n import tyro\n\n import viser\n import viser.transforms as tf\n\n\n @dataclass(frozen=True)\n class SmplOutputs:\n vertices: np.ndarray\n faces: np.ndarray\n T_world_joint: np.ndarray # (num_joints, 4, 4)\n T_parent_joint: np.ndarray # (num_joints, 4, 4)\n\n\n class SmplHelper:\n \"\"\"Helper for models in the SMPL family, implemented in numpy.\"\"\"\n\n def __init__(self, model_path: Path) -> None:\n assert model_path.suffix.lower() == \".npz\", \"Model should be an .npz file!\"\n body_dict = dict(**onp.load(model_path, allow_pickle=True))\n\n self._J_regressor = body_dict[\"J_regressor\"]\n self._weights = body_dict[\"weights\"]\n self._v_template = body_dict[\"v_template\"]\n self._posedirs = body_dict[\"posedirs\"]\n self._shapedirs = body_dict[\"shapedirs\"]\n self._faces = body_dict[\"f\"]\n\n self.num_joints: int = self._weights.shape[-1]\n self.num_betas: int = self._shapedirs.shape[-1]\n self.parent_idx: np.ndarray = body_dict[\"kintree_table\"][0]\n\n def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutputs:\n # Get shaped vertices + joint positions, when all local poses are identity.\n v_tpose = self._v_template + np.einsum(\"vxb,b->vx\", self._shapedirs, betas)\n j_tpose = np.einsum(\"jv,vx->jx\", self._J_regressor, v_tpose)\n\n # Local SE(3) transforms.\n T_parent_joint = np.zeros((self.num_joints, 4, 4)) + np.eye(4)\n T_parent_joint[:, :3, :3] = joint_rotmats\n T_parent_joint[0, :3, 3] = j_tpose[0]\n T_parent_joint[1:, :3, 3] = j_tpose[1:] - j_tpose[self.parent_idx[1:]]\n\n # Forward kinematics.\n T_world_joint = T_parent_joint.copy()\n for i in range(1, self.num_joints):\n T_world_joint[i] = T_world_joint[self.parent_idx[i]] @ T_parent_joint[i]\n\n # Linear blend skinning.\n pose_delta = (joint_rotmats[1:, ...] - np.eye(3)).flatten()\n v_blend = v_tpose + np.einsum(\"byn,n->by\", self._posedirs, pose_delta)\n v_delta = np.ones((v_blend.shape[0], self.num_joints, 4))\n v_delta[:, :, :3] = v_blend[:, None, :] - j_tpose[None, :, :]\n v_posed = np.einsum(\n \"jxy,vj,vjy->vx\", T_world_joint[:, :3, :], self._weights, v_delta\n )\n return SmplOutputs(v_posed, self._faces, T_world_joint, T_parent_joint)\n\n\n def main(model_path: Path) -> None:\n server = viser.ViserServer()\n server.scene.set_up_direction(\"+y\")\n server.gui.configure_theme(control_layout=\"collapsible\")\n\n # Main loop. We'll read pose/shape from the GUI elements, compute the mesh,\n # and then send the updated mesh in a loop.\n model = SmplHelper(model_path)\n gui_elements = make_gui_elements(\n server,\n num_betas=model.num_betas,\n num_joints=model.num_joints,\n parent_idx=model.parent_idx,\n )\n smpl_outputs = model.get_outputs(\n betas=np.array([x.value for x in gui_elements.gui_betas]),\n joint_rotmats=onp.zeros((model.num_joints, 3, 3)) + onp.eye(3),\n )\n\n bone_wxyzs = np.array(\n [tf.SO3.from_matrix(R).wxyz for R in smpl_outputs.T_world_joint[:, :3, :3]]\n )\n bone_positions = smpl_outputs.T_world_joint[:, :3, 3]\n\n skinned_handle = server.scene.add_mesh_skinned(\n \"/human\",\n smpl_outputs.vertices,\n smpl_outputs.faces,\n bone_wxyzs=bone_wxyzs,\n bone_positions=bone_positions,\n skin_weights=model._weights,\n wireframe=gui_elements.gui_wireframe.value,\n color=gui_elements.gui_rgb.value,\n )\n\n while True:\n # Do nothing if no change.\n time.sleep(0.02)\n if not gui_elements.changed:\n continue\n\n gui_elements.changed = False\n\n # Compute SMPL outputs.\n smpl_outputs = model.get_outputs(\n betas=np.array([x.value for x in gui_elements.gui_betas]),\n joint_rotmats=np.stack(\n [\n tf.SO3.exp(np.array(x.value)).as_matrix()\n for x in gui_elements.gui_joints\n ],\n axis=0,\n ),\n )\n\n # Match transform control gizmos to joint positions.\n for i, control in enumerate(gui_elements.transform_controls):\n control.position = smpl_outputs.T_parent_joint[i, :3, 3]\n skinned_handle.bones[i].wxyz = tf.SO3.from_matrix(\n smpl_outputs.T_world_joint[i, :3, :3]\n ).wxyz\n skinned_handle.bones[i].position = smpl_outputs.T_world_joint[i, :3, 3]\n\n\n @dataclass\n class GuiElements:\n \"\"\"Structure containing handles for reading from GUI elements.\"\"\"\n\n gui_rgb: viser.GuiInputHandle[Tuple[int, int, int]]\n gui_wireframe: viser.GuiInputHandle[bool]\n gui_betas: List[viser.GuiInputHandle[float]]\n gui_joints: List[viser.GuiInputHandle[Tuple[float, float, float]]]\n transform_controls: List[viser.TransformControlsHandle]\n\n changed: bool\n \"\"\"This flag will be flipped to True whenever the mesh needs to be re-generated.\"\"\"\n\n\n def make_gui_elements(\n server: viser.ViserServer,\n num_betas: int,\n num_joints: int,\n parent_idx: np.ndarray,\n ) -> GuiElements:\n \"\"\"Make GUI elements for interacting with the model.\"\"\"\n\n tab_group = server.gui.add_tab_group()\n\n def set_changed(_) -> None:\n out.changed = True # out is define later!\n\n # GUI elements: mesh settings + visibility.\n with tab_group.add_tab(\"View\", viser.Icon.VIEWFINDER):\n gui_rgb = server.gui.add_rgb(\"Color\", initial_value=(90, 200, 255))\n gui_wireframe = server.gui.add_checkbox(\"Wireframe\", initial_value=False)\n gui_show_controls = server.gui.add_checkbox(\"Handles\", initial_value=True)\n\n gui_rgb.on_update(set_changed)\n gui_wireframe.on_update(set_changed)\n\n @gui_show_controls.on_update\n def _(_):\n for control in transform_controls:\n control.visible = gui_show_controls.value\n\n # GUI elements: shape parameters.\n with tab_group.add_tab(\"Shape\", viser.Icon.BOX):\n gui_reset_shape = server.gui.add_button(\"Reset Shape\")\n gui_random_shape = server.gui.add_button(\"Random Shape\")\n\n @gui_reset_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = 0.0\n\n @gui_random_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = onp.random.normal(loc=0.0, scale=1.0)\n\n gui_betas = []\n for i in range(num_betas):\n beta = server.gui.add_slider(\n f\"beta{i}\", min=-5.0, max=5.0, step=0.01, initial_value=0.0\n )\n gui_betas.append(beta)\n beta.on_update(set_changed)\n\n # GUI elements: joint angles.\n with tab_group.add_tab(\"Joints\", viser.Icon.ANGLE):\n gui_reset_joints = server.gui.add_button(\"Reset Joints\")\n gui_random_joints = server.gui.add_button(\"Random Joints\")\n\n @gui_reset_joints.on_click\n def _(_):\n for joint in gui_joints:\n joint.value = (0.0, 0.0, 0.0)\n\n @gui_random_joints.on_click\n def _(_):\n for joint in gui_joints:\n # It's hard to uniformly sample orientations directly in so(3), so we\n # first sample on S^3 and then convert.\n quat = onp.random.normal(loc=0.0, scale=1.0, size=(4,))\n quat /= onp.linalg.norm(quat)\n joint.value = tf.SO3(wxyz=quat).log()\n\n gui_joints: List[viser.GuiInputHandle[Tuple[float, float, float]]] = []\n for i in range(num_joints):\n gui_joint = server.gui.add_vector3(\n label=f\"Joint {i}\",\n initial_value=(0.0, 0.0, 0.0),\n step=0.05,\n )\n gui_joints.append(gui_joint)\n\n def set_callback_in_closure(i: int) -> None:\n @gui_joint.on_update\n def _(_):\n transform_controls[i].wxyz = tf.SO3.exp(\n np.array(gui_joints[i].value)\n ).wxyz\n out.changed = True\n\n set_callback_in_closure(i)\n\n # Transform control gizmos on joints.\n transform_controls: List[viser.TransformControlsHandle] = []\n prefixed_joint_names = [] # Joint names, but prefixed with parents.\n for i in range(num_joints):\n prefixed_joint_name = f\"joint_{i}\"\n if i > 0:\n prefixed_joint_name = (\n prefixed_joint_names[parent_idx[i]] + \"/\" + prefixed_joint_name\n )\n prefixed_joint_names.append(prefixed_joint_name)\n controls = server.scene.add_transform_controls(\n f\"/smpl/{prefixed_joint_name}\",\n depth_test=False,\n scale=0.2 * (0.75 ** prefixed_joint_name.count(\"/\")),\n disable_axes=True,\n disable_sliders=True,\n visible=gui_show_controls.value,\n )\n transform_controls.append(controls)\n\n def set_callback_in_closure(i: int) -> None:\n @controls.on_update\n def _(_) -> None:\n axisangle = tf.SO3(transform_controls[i].wxyz).log()\n gui_joints[i].value = (axisangle[0], axisangle[1], axisangle[2])\n\n set_callback_in_closure(i)\n\n out = GuiElements(\n gui_rgb,\n gui_wireframe,\n gui_betas,\n gui_joints,\n transform_controls=transform_controls,\n changed=True,\n )\n return out\n\n\n if __name__ == \"__main__\":\n tyro.cli(main, description=__doc__)\n" }, { "path": "viser/docs/source/extras.md", "content": "# Record3D + URDF Helpers\n\n\n\n.. automodule:: viser.extras\n\n\n" }, { "path": "viser/docs/source/gui_api.md", "content": "# GUI API\n\n\n\n.. autoclass:: viser.GuiApi\n :members:\n :undoc-members:\n :inherited-members:\n\n\n" }, { "path": "viser/docs/source/gui_handles.md", "content": "# GUI Handles\n\n\n\n.. autoclass:: viser.GuiInputHandle()\n\n.. autoclass:: viser.GuiButtonHandle()\n\n.. autoclass:: viser.GuiButtonGroupHandle()\n\n.. autoclass:: viser.GuiDropdownHandle()\n\n.. autoclass:: viser.GuiFolderHandle()\n\n.. autoclass:: viser.GuiMarkdownHandle()\n\n.. autoclass:: viser.GuiPlotlyHandle()\n\n.. autoclass:: viser.GuiTabGroupHandle()\n\n.. autoclass:: viser.GuiTabHandle()\n\n\n" }, { "path": "viser/docs/source/icons.md", "content": "# Icons\n\nIcons for GUI elements (such as :meth:`GuiApi.add_button()`) can be\nspecified using the :class:`viser.Icon` enum.\n\n\n\n.. autoclass:: viser.IconName\n.. autoclass:: viser.Icon\n\n\n" }, { "path": "viser/docs/source/index.md", "content": "# viser\n\n|mypy| |nbsp| |pyright| |nbsp| |typescript| |nbsp| |versions|\n\n**viser** is a library for interactive 3D visualization in Python.\n\nFeatures include:\n\n- API for visualizing 3D primitives\n- GUI building blocks: buttons, checkboxes, text inputs, sliders, etc.\n- Scene interaction tools (clicks, selection, transform gizmos)\n- Programmatic camera control and rendering\n- An entirely web-based client, for easy use over SSH!\n\n## Installation\n\nYou can install `viser` with `pip`:\n\n```bash\npip install viser\n```\n\nTo include example dependencies:\n\n```bash\npip install viser[examples]\n```\n\nAfter an example script is running, you can connect by navigating to the printed\nURL (default: `http://localhost:8080`).\n\n\n\n.. toctree::\n :caption: Notes\n :hidden:\n :maxdepth: 1\n :titlesonly:\n\n ./conventions.md\n ./development.md\n\n.. toctree::\n :caption: API (Basics)\n :hidden:\n :maxdepth: 1\n :titlesonly:\n\n ./server.md\n ./scene_api.md\n ./gui_api.md\n\n\n.. toctree::\n :caption: API (Advanced)\n :hidden:\n :maxdepth: 1\n :titlesonly:\n\n ./client_handles.md\n ./camera_handles.md\n ./gui_handles.md\n ./scene_handles.md\n ./events.md\n ./icons.md\n\n\n.. toctree::\n :caption: API (Auxiliary)\n :hidden:\n :maxdepth: 1\n :titlesonly:\n\n ./transforms.md\n ./infrastructure.md\n ./extras.md\n\n.. toctree::\n :caption: Examples\n :hidden:\n :maxdepth: 1\n :titlesonly:\n :glob:\n\n examples/*\n\n\n.. |build| image:: https://github.com/nerfstudio-project/viser/workflows/build/badge.svg\n :alt: Build status icon\n :target: https://github.com/nerfstudio-project/viser\n.. |mypy| image:: https://github.com/nerfstudio-project/viser/workflows/mypy/badge.svg?branch=main\n :alt: Mypy status icon\n :target: https://github.com/nerfstudio-project/viser\n.. |pyright| image:: https://github.com/nerfstudio-project/viser/workflows/pyright/badge.svg?branch=main\n :alt: Mypy status icon\n :target: https://github.com/nerfstudio-project/viser\n.. |typescript| image:: https://github.com/nerfstudio-project/viser/workflows/typescript-compile/badge.svg\n :alt: TypeScript status icon\n :target: https://github.com/nerfstudio-project/viser\n.. |versions| image:: https://img.shields.io/pypi/pyversions/viser\n :alt: Version icon\n :target: https://pypi.org/project/viser/\n.. |nbsp| unicode:: 0xA0\n :trim:\n\n\n" }, { "path": "viser/docs/source/infrastructure.md", "content": "# Communication\n\n\n\n.. automodule:: viser.infra\n :show-inheritance:\n\n\n" }, { "path": "viser/docs/source/scene_api.md", "content": "# Scene API\n\n\n\n.. autoclass:: viser.SceneApi\n :members:\n :undoc-members:\n :inherited-members:\n\n\n" }, { "path": "viser/docs/source/scene_handles.md", "content": "# Scene Handles\n\nA handle is created for each object that is added to the scene. These can be\nused to read and set state, as well as detect clicks.\n\nWhen a scene node is added to a server (for example, via\n:func:`viser.ViserServer.add_frame()`), state is synchronized between all\nconnected clients. When a scene node is added to a client (for example, via\n:func:`viser.ClientHandle.add_frame()`), state is local to a specific client.\n\n\n\n.. autoclass:: viser.SceneNodeHandle\n\n.. autoclass:: viser.CameraFrustumHandle\n\n.. autoclass:: viser.FrameHandle\n\n.. autoclass:: viser.BatchedAxesHandle\n\n.. autoclass:: viser.GlbHandle\n\n.. autoclass:: viser.Gui3dContainerHandle\n\n.. autoclass:: viser.ImageHandle\n\n.. autoclass:: viser.LabelHandle\n\n.. autoclass:: viser.MeshHandle\n\n.. autoclass:: viser.MeshSkinnedHandle\n\n.. autoclass:: viser.MeshSkinnedBoneHandle\n\n.. autoclass:: viser.PointCloudHandle\n\n.. autoclass:: viser.TransformControlsHandle\n\n.. autoclass:: viser.GaussianSplatHandle\n\n\n" }, { "path": "viser/docs/source/server.md", "content": "# Viser Server\n\n\n\n.. autoclass:: viser.ViserServer\n\n\n" }, { "path": "viser/docs/source/transforms.md", "content": "# Transforms\n\n\n\n.. automodule:: viser.transforms\n :show-inheritance:\n\n\n" }, { "path": "viser/docs/update_example_docs.py", "content": "\"\"\"Helper script for updating the auto-generated examples pages in the documentation.\"\"\"\n\nfrom __future__ import annotations\n\nimport dataclasses\nimport pathlib\nimport shutil\nfrom typing import Iterable\n\nimport m2r2\nimport tyro\n\n\n@dataclasses.dataclass\nclass ExampleMetadata:\n index: str\n index_with_zero: str\n source: str\n title: str\n description: str\n\n @staticmethod\n def from_path(path: pathlib.Path) -> ExampleMetadata:\n # 01_functions -> 01, _, functions.\n index, _, _ = path.stem.partition(\"_\")\n\n # 01 -> 1.\n index_with_zero = index\n index = str(int(index))\n\n source = path.read_text().strip()\n\n docstring = source.split('\"\"\"')[1].strip()\n\n title, _, description = docstring.partition(\"\\n\")\n\n return ExampleMetadata(\n index=index,\n index_with_zero=index_with_zero,\n source=source.partition('\"\"\"')[2].partition('\"\"\"')[2].strip(),\n title=title,\n description=description.strip(),\n )\n\n\ndef get_example_paths(examples_dir: pathlib.Path) -> Iterable[pathlib.Path]:\n return filter(\n lambda p: not p.name.startswith(\"_\"), sorted(examples_dir.glob(\"*.py\"))\n )\n\n\nREPO_ROOT = pathlib.Path(__file__).absolute().parent.parent\n\n\ndef main(\n examples_dir: pathlib.Path = REPO_ROOT / \"examples\",\n sphinx_source_dir: pathlib.Path = REPO_ROOT / \"docs\" / \"source\",\n) -> None:\n example_doc_dir = sphinx_source_dir / \"examples\"\n shutil.rmtree(example_doc_dir)\n example_doc_dir.mkdir()\n\n for path in get_example_paths(examples_dir):\n ex = ExampleMetadata.from_path(path)\n\n relative_dir = path.parent.relative_to(examples_dir)\n target_dir = example_doc_dir / relative_dir\n target_dir.mkdir(exist_ok=True, parents=True)\n\n (target_dir / f\"{path.stem}.rst\").write_text(\n \"\\n\".join(\n [\n (\n \".. Comment: this file is automatically generated by\"\n \" `update_example_docs.py`.\"\n ),\n \" It should not be modified manually.\",\n \"\",\n f\"{ex.title}\",\n \"==========================================\",\n \"\",\n m2r2.convert(ex.description),\n \"\",\n \"\",\n \".. code-block:: python\",\n \" :linenos:\",\n \"\",\n \"\",\n \"\\n\".join(\n f\" {line}\".rstrip() for line in ex.source.split(\"\\n\")\n ),\n \"\",\n ]\n )\n )\n\n\nif __name__ == \"__main__\":\n tyro.cli(main, description=__doc__)\n" }, { "path": "viser/examples/00_coordinate_frames.py", "content": "\"\"\"Coordinate frames\n\nIn this basic example, we visualize a set of coordinate frames.\n\nNaming for all scene nodes are hierarchical; /tree/branch, for example, is defined\nrelative to /tree.\n\"\"\"\n\nimport random\nimport time\n\nimport viser\n\nserver = viser.ViserServer()\n\nwhile True:\n # Add some coordinate frames to the scene. These will be visualized in the viewer.\n server.scene.add_frame(\n \"/tree\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(random.random() * 2.0, 2.0, 0.2),\n )\n server.scene.add_frame(\n \"/tree/branch\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(random.random() * 2.0, 2.0, 0.2),\n )\n leaf = server.scene.add_frame(\n \"/tree/branch/leaf\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(random.random() * 2.0, 2.0, 0.2),\n )\n time.sleep(5.0)\n\n # Remove the leaf node from the scene.\n leaf.remove()\n time.sleep(0.5)\n" }, { "path": "viser/examples/01_image.py", "content": "\"\"\"Images\n\nExample for sending images to the viewer.\n\nWe can send backgrond images to display behind the viewer (useful for visualizing\nNeRFs), or images to render as 3D textures.\n\"\"\"\n\nimport time\nfrom pathlib import Path\n\nimport imageio.v3 as iio\nimport numpy as onp\n\nimport viser\n\n\ndef main() -> None:\n server = viser.ViserServer()\n\n # Add a background image.\n server.scene.set_background_image(\n iio.imread(Path(__file__).parent / \"assets/Cal_logo.png\"),\n format=\"png\",\n )\n\n # Add main image.\n server.scene.add_image(\n \"/img\",\n iio.imread(Path(__file__).parent / \"assets/Cal_logo.png\"),\n 4.0,\n 4.0,\n format=\"png\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(2.0, 2.0, 0.0),\n )\n while True:\n server.scene.add_image(\n \"/noise\",\n onp.random.randint(\n 0,\n 256,\n size=(400, 400, 3),\n dtype=onp.uint8,\n ),\n 4.0,\n 4.0,\n format=\"jpeg\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=(2.0, 2.0, -1e-2),\n )\n time.sleep(0.2)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/02_gui.py", "content": "\"\"\"GUI basics\n\nExamples of basic GUI elements that we can create, read from, and write to.\"\"\"\n\nimport time\n\nimport numpy as onp\n\nimport viser\n\n\ndef main() -> None:\n server = viser.ViserServer()\n\n # Add some common GUI elements: number inputs, sliders, vectors, checkboxes.\n with server.gui.add_folder(\"Read-only\"):\n gui_counter = server.gui.add_number(\n \"Counter\",\n initial_value=0,\n disabled=True,\n )\n gui_slider = server.gui.add_slider(\n \"Slider\",\n min=0,\n max=100,\n step=1,\n initial_value=0,\n disabled=True,\n )\n gui_progress = server.gui.add_progress_bar(25, animated=True)\n\n with server.gui.add_folder(\"Editable\"):\n gui_vector2 = server.gui.add_vector2(\n \"Position\",\n initial_value=(0.0, 0.0),\n step=0.1,\n )\n gui_vector3 = server.gui.add_vector3(\n \"Size\",\n initial_value=(1.0, 1.0, 1.0),\n step=0.25,\n )\n with server.gui.add_folder(\"Text toggle\"):\n gui_checkbox_hide = server.gui.add_checkbox(\n \"Hide\",\n initial_value=False,\n )\n gui_text = server.gui.add_text(\n \"Text\",\n initial_value=\"Hello world\",\n )\n gui_button = server.gui.add_button(\"Button\")\n gui_checkbox_disable = server.gui.add_checkbox(\n \"Disable\",\n initial_value=False,\n )\n gui_rgb = server.gui.add_rgb(\n \"Color\",\n initial_value=(255, 255, 0),\n )\n gui_multi_slider = server.gui.add_multi_slider(\n \"Multi slider\",\n min=0,\n max=100,\n step=1,\n initial_value=(0, 30, 100),\n marks=((0, \"0\"), (50, \"5\"), (70, \"7\"), 99),\n )\n gui_slider_positions = server.gui.add_slider(\n \"# sliders\",\n min=0,\n max=10,\n step=1,\n initial_value=3,\n marks=((0, \"0\"), (5, \"5\"), (7, \"7\"), 10),\n )\n gui_upload_button = server.gui.add_upload_button(\n \"Upload\", icon=viser.Icon.UPLOAD\n )\n\n @gui_upload_button.on_upload\n def _(_) -> None:\n \"\"\"Callback for when a file is uploaded.\"\"\"\n file = gui_upload_button.value\n print(file.name, len(file.content), \"bytes\")\n\n # Pre-generate a point cloud to send.\n point_positions = onp.random.uniform(low=-1.0, high=1.0, size=(5000, 3))\n color_coeffs = onp.random.uniform(0.4, 1.0, size=(point_positions.shape[0]))\n\n counter = 0\n while True:\n # We can set the value of an input to a particular value. Changes are\n # automatically reflected in connected clients.\n gui_counter.value = counter\n gui_slider.value = counter % 100\n\n # We can set the position of a scene node with `.position`, and read the value\n # of a gui element with `.value`. Changes are automatically reflected in\n # connected clients.\n server.scene.add_point_cloud(\n \"/point_cloud\",\n points=point_positions * onp.array(gui_vector3.value, dtype=onp.float32),\n colors=(\n onp.tile(gui_rgb.value, point_positions.shape[0]).reshape((-1, 3))\n * color_coeffs[:, None]\n ).astype(onp.uint8),\n position=gui_vector2.value + (0,),\n point_shape=\"circle\",\n )\n\n gui_progress.value = float((counter % 100))\n\n # We can use `.visible` and `.disabled` to toggle GUI elements.\n gui_text.visible = not gui_checkbox_hide.value\n gui_button.visible = not gui_checkbox_hide.value\n gui_rgb.disabled = gui_checkbox_disable.value\n gui_button.disabled = gui_checkbox_disable.value\n gui_upload_button.disabled = gui_checkbox_disable.value\n\n # Update the number of handles in the multi-slider.\n if gui_slider_positions.value != len(gui_multi_slider.value):\n gui_multi_slider.value = onp.linspace(\n 0, 100, gui_slider_positions.value, dtype=onp.int64\n )\n\n counter += 1\n time.sleep(0.01)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/03_gui_callbacks.py", "content": "\"\"\"GUI callbacks\n\nAsynchronous usage of GUI elements: we can attach callbacks that are called as soon as\nwe get updates.\"\"\"\n\nimport time\n\nimport numpy as onp\nfrom typing_extensions import assert_never\n\nimport viser\n\n\ndef main() -> None:\n server = viser.ViserServer()\n\n gui_reset_scene = server.gui.add_button(\"Reset Scene\")\n\n gui_plane = server.gui.add_dropdown(\n \"Grid plane\", (\"xz\", \"xy\", \"yx\", \"yz\", \"zx\", \"zy\")\n )\n\n def update_plane() -> None:\n server.scene.add_grid(\n \"/grid\",\n width=10.0,\n height=20.0,\n width_segments=10,\n height_segments=20,\n plane=gui_plane.value,\n )\n\n gui_plane.on_update(lambda _: update_plane())\n\n with server.gui.add_folder(\"Control\"):\n gui_show_frame = server.gui.add_checkbox(\"Show Frame\", initial_value=True)\n gui_show_everything = server.gui.add_checkbox(\n \"Show Everything\", initial_value=True\n )\n gui_axis = server.gui.add_dropdown(\"Axis\", (\"x\", \"y\", \"z\"))\n gui_include_z = server.gui.add_checkbox(\"Z in dropdown\", initial_value=True)\n\n @gui_include_z.on_update\n def _(_) -> None:\n gui_axis.options = (\"x\", \"y\", \"z\") if gui_include_z.value else (\"x\", \"y\")\n\n with server.gui.add_folder(\"Sliders\"):\n gui_location = server.gui.add_slider(\n \"Location\", min=-5.0, max=5.0, step=0.05, initial_value=0.0\n )\n gui_num_points = server.gui.add_slider(\n \"# Points\", min=1000, max=200_000, step=1000, initial_value=10_000\n )\n\n def draw_frame() -> None:\n axis = gui_axis.value\n if axis == \"x\":\n pos = (gui_location.value, 0.0, 0.0)\n elif axis == \"y\":\n pos = (0.0, gui_location.value, 0.0)\n elif axis == \"z\":\n pos = (0.0, 0.0, gui_location.value)\n else:\n assert_never(axis)\n\n server.scene.add_frame(\n \"/frame\",\n wxyz=(1.0, 0.0, 0.0, 0.0),\n position=pos,\n show_axes=gui_show_frame.value,\n axes_length=5.0,\n )\n\n def draw_points() -> None:\n num_points = gui_num_points.value\n server.scene.add_point_cloud(\n \"/frame/point_cloud\",\n points=onp.random.normal(size=(num_points, 3)),\n colors=onp.random.randint(0, 256, size=(num_points, 3)),\n )\n\n # We can (optionally) also attach callbacks!\n # Here, we update the point clouds + frames whenever any of the GUI items are updated.\n gui_show_frame.on_update(lambda _: draw_frame())\n gui_show_everything.on_update(\n lambda _: server.scene.set_global_visibility(gui_show_everything.value)\n )\n gui_axis.on_update(lambda _: draw_frame())\n gui_location.on_update(lambda _: draw_frame())\n gui_num_points.on_update(lambda _: draw_points())\n\n @gui_reset_scene.on_click\n def _(_) -> None:\n \"\"\"Reset the scene when the reset button is clicked.\"\"\"\n gui_show_frame.value = True\n gui_location.value = 0.0\n gui_axis.value = \"x\"\n gui_num_points.value = 10_000\n\n draw_frame()\n draw_points()\n\n # Finally, let's add the initial frame + point cloud and just loop infinitely. :)\n update_plane()\n draw_frame()\n draw_points()\n while True:\n time.sleep(1.0)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/04_camera_poses.py", "content": "\"\"\"Camera poses\n\nExample showing how we can detect new clients and read camera poses from them.\n\"\"\"\n\nimport time\n\nimport viser\n\nserver = viser.ViserServer()\nserver.scene.world_axes.visible = True\n\n\n@server.on_client_connect\ndef _(client: viser.ClientHandle) -> None:\n print(\"new client!\")\n\n # This will run whenever we get a new camera!\n @client.camera.on_update\n def _(_: viser.CameraHandle) -> None:\n print(f\"New camera on client {client.client_id}!\")\n\n # Show the client ID in the GUI.\n gui_info = client.gui.add_text(\"Client ID\", initial_value=str(client.client_id))\n gui_info.disabled = True\n\n\nwhile True:\n # Get all currently connected clients.\n clients = server.get_clients()\n print(\"Connected client IDs\", clients.keys())\n\n for id, client in clients.items():\n print(f\"Camera pose for client {id}\")\n print(f\"\\twxyz: {client.camera.wxyz}\")\n print(f\"\\tposition: {client.camera.position}\")\n print(f\"\\tfov: {client.camera.fov}\")\n print(f\"\\taspect: {client.camera.aspect}\")\n print(f\"\\tlast update: {client.camera.update_timestamp}\")\n\n time.sleep(2.0)\n" }, { "path": "viser/examples/05_camera_commands.py", "content": "\"\"\"Camera commands\n\nIn addition to reads, camera parameters also support writes. These are synced to the\ncorresponding client automatically.\n\"\"\"\n\nimport time\n\nimport numpy as onp\n\nimport viser\nimport viser.transforms as tf\n\nserver = viser.ViserServer()\nnum_frames = 20\n\n\n@server.on_client_connect\ndef _(client: viser.ClientHandle) -> None:\n \"\"\"For each client that connects, we create a set of random frames + a click handler for each frame.\n\n When a frame is clicked, we move the camera to the corresponding frame.\n \"\"\"\n\n rng = onp.random.default_rng(0)\n\n def make_frame(i: int) -> None:\n # Sample a random orientation + position.\n wxyz = rng.normal(size=4)\n wxyz /= onp.linalg.norm(wxyz)\n position = rng.uniform(-3.0, 3.0, size=(3,))\n\n # Create a coordinate frame and label.\n frame = client.scene.add_frame(f\"/frame_{i}\", wxyz=wxyz, position=position)\n client.scene.add_label(f\"/frame_{i}/label\", text=f\"Frame {i}\")\n\n # Move the camera when we click a frame.\n @frame.on_click\n def _(_):\n T_world_current = tf.SE3.from_rotation_and_translation(\n tf.SO3(client.camera.wxyz), client.camera.position\n )\n T_world_target = tf.SE3.from_rotation_and_translation(\n tf.SO3(frame.wxyz), frame.position\n ) @ tf.SE3.from_translation(onp.array([0.0, 0.0, -0.5]))\n\n T_current_target = T_world_current.inverse() @ T_world_target\n\n for j in range(20):\n T_world_set = T_world_current @ tf.SE3.exp(\n T_current_target.log() * j / 19.0\n )\n\n # We can atomically set the orientation and the position of the camera\n # together to prevent jitter that might happen if one was set before the\n # other.\n with client.atomic():\n client.camera.wxyz = T_world_set.rotation().wxyz\n client.camera.position = T_world_set.translation()\n\n client.flush() # Optional!\n time.sleep(1.0 / 60.0)\n\n # Mouse interactions should orbit around the frame origin.\n client.camera.look_at = frame.position\n\n for i in range(num_frames):\n make_frame(i)\n\n\nwhile True:\n time.sleep(1.0)\n" }, { "path": "viser/examples/06_mesh.py", "content": "\"\"\"Meshes\n\nVisualize a mesh. To get the demo data, see `./assets/download_dragon_mesh.sh`.\n\"\"\"\n\nimport time\nfrom pathlib import Path\n\nimport numpy as onp\nimport trimesh\n\nimport viser\nimport viser.transforms as tf\n\nmesh = trimesh.load_mesh(str(Path(__file__).parent / \"assets/dragon.obj\"))\nassert isinstance(mesh, trimesh.Trimesh)\nmesh.apply_scale(0.05)\n\nvertices = mesh.vertices\nfaces = mesh.faces\nprint(f\"Loaded mesh with {vertices.shape} vertices, {faces.shape} faces\")\n\nserver = viser.ViserServer()\nserver.scene.add_mesh_simple(\n name=\"/simple\",\n vertices=vertices,\n faces=faces,\n wxyz=tf.SO3.from_x_radians(onp.pi / 2).wxyz,\n position=(0.0, 0.0, 0.0),\n)\nserver.scene.add_mesh_trimesh(\n name=\"/trimesh\",\n mesh=mesh.smoothed(),\n wxyz=tf.SO3.from_x_radians(onp.pi / 2).wxyz,\n position=(0.0, 5.0, 0.0),\n)\n\nwhile True:\n time.sleep(10.0)\n" }, { "path": "viser/examples/07_record3d_visualizer.py", "content": "\"\"\"Record3D visualizer\n\nParse and stream record3d captures. To get the demo data, see `./assets/download_record3d_dance.sh`.\n\"\"\"\n\nimport time\nfrom pathlib import Path\n\nimport numpy as onp\nimport tyro\nfrom tqdm.auto import tqdm\n\nimport viser\nimport viser.extras\nimport viser.transforms as tf\n\n\ndef main(\n data_path: Path = Path(__file__).parent / \"record3d_dance\",\n downsample_factor: int = 4,\n max_frames: int = 100,\n share: bool = False,\n) -> None:\n server = viser.ViserServer()\n if share:\n server.request_share_url()\n\n print(\"Loading frames!\")\n loader = viser.extras.Record3dLoader(data_path)\n num_frames = min(max_frames, loader.num_frames())\n\n # Add playback UI.\n with server.gui.add_folder(\"Playback\"):\n gui_timestep = server.gui.add_slider(\n \"Timestep\",\n min=0,\n max=num_frames - 1,\n step=1,\n initial_value=0,\n disabled=True,\n )\n gui_next_frame = server.gui.add_button(\"Next Frame\", disabled=True)\n gui_prev_frame = server.gui.add_button(\"Prev Frame\", disabled=True)\n gui_playing = server.gui.add_checkbox(\"Playing\", True)\n gui_framerate = server.gui.add_slider(\n \"FPS\", min=1, max=60, step=0.1, initial_value=loader.fps\n )\n gui_framerate_options = server.gui.add_button_group(\n \"FPS options\", (\"10\", \"20\", \"30\", \"60\")\n )\n\n # Frame step buttons.\n @gui_next_frame.on_click\n def _(_) -> None:\n gui_timestep.value = (gui_timestep.value + 1) % num_frames\n\n @gui_prev_frame.on_click\n def _(_) -> None:\n gui_timestep.value = (gui_timestep.value - 1) % num_frames\n\n # Disable frame controls when we're playing.\n @gui_playing.on_update\n def _(_) -> None:\n gui_timestep.disabled = gui_playing.value\n gui_next_frame.disabled = gui_playing.value\n gui_prev_frame.disabled = gui_playing.value\n\n # Set the framerate when we click one of the options.\n @gui_framerate_options.on_click\n def _(_) -> None:\n gui_framerate.value = int(gui_framerate_options.value)\n\n prev_timestep = gui_timestep.value\n\n # Toggle frame visibility when the timestep slider changes.\n @gui_timestep.on_update\n def _(_) -> None:\n nonlocal prev_timestep\n current_timestep = gui_timestep.value\n with server.atomic():\n frame_nodes[current_timestep].visible = True\n frame_nodes[prev_timestep].visible = False\n prev_timestep = current_timestep\n server.flush() # Optional!\n\n # Load in frames.\n server.scene.add_frame(\n \"/frames\",\n wxyz=tf.SO3.exp(onp.array([onp.pi / 2.0, 0.0, 0.0])).wxyz,\n position=(0, 0, 0),\n show_axes=False,\n )\n frame_nodes: list[viser.FrameHandle] = []\n for i in tqdm(range(num_frames)):\n frame = loader.get_frame(i)\n position, color = frame.get_point_cloud(downsample_factor)\n\n # Add base frame.\n frame_nodes.append(server.scene.add_frame(f\"/frames/t{i}\", show_axes=False))\n\n # Place the point cloud in the frame.\n server.scene.add_point_cloud(\n name=f\"/frames/t{i}/point_cloud\",\n points=position,\n colors=color,\n point_size=0.01,\n point_shape=\"rounded\",\n )\n\n # Place the frustum.\n fov = 2 * onp.arctan2(frame.rgb.shape[0] / 2, frame.K[0, 0])\n aspect = frame.rgb.shape[1] / frame.rgb.shape[0]\n server.scene.add_camera_frustum(\n f\"/frames/t{i}/frustum\",\n fov=fov,\n aspect=aspect,\n scale=0.15,\n image=frame.rgb[::downsample_factor, ::downsample_factor],\n wxyz=tf.SO3.from_matrix(frame.T_world_camera[:3, :3]).wxyz,\n position=frame.T_world_camera[:3, 3],\n )\n\n # Add some axes.\n server.scene.add_frame(\n f\"/frames/t{i}/frustum/axes\",\n axes_length=0.05,\n axes_radius=0.005,\n )\n\n # Hide all but the current frame.\n for i, frame_node in enumerate(frame_nodes):\n frame_node.visible = i == gui_timestep.value\n\n # Playback update loop.\n prev_timestep = gui_timestep.value\n while True:\n if gui_playing.value:\n gui_timestep.value = (gui_timestep.value + 1) % num_frames\n\n time.sleep(1.0 / gui_framerate.value)\n\n\nif __name__ == \"__main__\":\n tyro.cli(main)\n" }, { "path": "viser/examples/08_smpl_visualizer.py", "content": "\"\"\"SMPL model visualizer\n\nVisualizer for SMPL human body models. Requires a .npz model file.\n\nSee here for download instructions:\n https://github.com/vchoutas/smplx?tab=readme-ov-file#downloading-the-model\n\"\"\"\n\nfrom __future__ import annotations\n\nimport time\nfrom dataclasses import dataclass\nfrom pathlib import Path\n\nimport numpy as np\nimport numpy as onp\nimport tyro\n\nimport viser\nimport viser.transforms as tf\n\n\n@dataclass(frozen=True)\nclass SmplOutputs:\n vertices: np.ndarray\n faces: np.ndarray\n T_world_joint: np.ndarray # (num_joints, 4, 4)\n T_parent_joint: np.ndarray # (num_joints, 4, 4)\n\n\nclass SmplHelper:\n \"\"\"Helper for models in the SMPL family, implemented in numpy.\"\"\"\n\n def __init__(self, model_path: Path) -> None:\n assert model_path.suffix.lower() == \".npz\", \"Model should be an .npz file!\"\n body_dict = dict(**onp.load(model_path, allow_pickle=True))\n\n self._J_regressor = body_dict[\"J_regressor\"]\n self._weights = body_dict[\"weights\"]\n self._v_template = body_dict[\"v_template\"]\n self._posedirs = body_dict[\"posedirs\"]\n self._shapedirs = body_dict[\"shapedirs\"]\n self._faces = body_dict[\"f\"]\n\n self.num_joints: int = self._weights.shape[-1]\n self.num_betas: int = self._shapedirs.shape[-1]\n self.parent_idx: np.ndarray = body_dict[\"kintree_table\"][0]\n\n def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutputs:\n # Get shaped vertices + joint positions, when all local poses are identity.\n v_tpose = self._v_template + np.einsum(\"vxb,b->vx\", self._shapedirs, betas)\n j_tpose = np.einsum(\"jv,vx->jx\", self._J_regressor, v_tpose)\n\n # Local SE(3) transforms.\n T_parent_joint = np.zeros((self.num_joints, 4, 4)) + np.eye(4)\n T_parent_joint[:, :3, :3] = joint_rotmats\n T_parent_joint[0, :3, 3] = j_tpose[0]\n T_parent_joint[1:, :3, 3] = j_tpose[1:] - j_tpose[self.parent_idx[1:]]\n\n # Forward kinematics.\n T_world_joint = T_parent_joint.copy()\n for i in range(1, self.num_joints):\n T_world_joint[i] = T_world_joint[self.parent_idx[i]] @ T_parent_joint[i]\n\n # Linear blend skinning.\n pose_delta = (joint_rotmats[1:, ...] - np.eye(3)).flatten()\n v_blend = v_tpose + np.einsum(\"byn,n->by\", self._posedirs, pose_delta)\n v_delta = np.ones((v_blend.shape[0], self.num_joints, 4))\n v_delta[:, :, :3] = v_blend[:, None, :] - j_tpose[None, :, :]\n v_posed = np.einsum(\n \"jxy,vj,vjy->vx\", T_world_joint[:, :3, :], self._weights, v_delta\n )\n return SmplOutputs(v_posed, self._faces, T_world_joint, T_parent_joint)\n\n\ndef main(model_path: Path) -> None:\n server = viser.ViserServer()\n server.scene.set_up_direction(\"+y\")\n server.gui.configure_theme(control_layout=\"collapsible\")\n\n # Main loop. We'll read pose/shape from the GUI elements, compute the mesh,\n # and then send the updated mesh in a loop.\n model = SmplHelper(model_path)\n gui_elements = make_gui_elements(\n server,\n num_betas=model.num_betas,\n num_joints=model.num_joints,\n parent_idx=model.parent_idx,\n )\n while True:\n # Do nothing if no change.\n time.sleep(0.02)\n if not gui_elements.changed:\n continue\n\n gui_elements.changed = False\n\n # Compute SMPL outputs.\n smpl_outputs = model.get_outputs(\n betas=np.array([x.value for x in gui_elements.gui_betas]),\n joint_rotmats=tf.SO3.exp(\n # (num_joints, 3)\n np.array([x.value for x in gui_elements.gui_joints])\n ).as_matrix(),\n )\n server.scene.add_mesh_simple(\n \"/human\",\n smpl_outputs.vertices,\n smpl_outputs.faces,\n wireframe=gui_elements.gui_wireframe.value,\n color=gui_elements.gui_rgb.value,\n )\n\n # Match transform control gizmos to joint positions.\n for i, control in enumerate(gui_elements.transform_controls):\n control.position = smpl_outputs.T_parent_joint[i, :3, 3]\n\n\n@dataclass\nclass GuiElements:\n \"\"\"Structure containing handles for reading from GUI elements.\"\"\"\n\n gui_rgb: viser.GuiInputHandle[tuple[int, int, int]]\n gui_wireframe: viser.GuiInputHandle[bool]\n gui_betas: list[viser.GuiInputHandle[float]]\n gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]]\n transform_controls: list[viser.TransformControlsHandle]\n\n changed: bool\n \"\"\"This flag will be flipped to True whenever the mesh needs to be re-generated.\"\"\"\n\n\ndef make_gui_elements(\n server: viser.ViserServer,\n num_betas: int,\n num_joints: int,\n parent_idx: np.ndarray,\n) -> GuiElements:\n \"\"\"Make GUI elements for interacting with the model.\"\"\"\n\n tab_group = server.gui.add_tab_group()\n\n def set_changed(_) -> None:\n out.changed = True # out is define later!\n\n # GUI elements: mesh settings + visibility.\n with tab_group.add_tab(\"View\", viser.Icon.VIEWFINDER):\n gui_rgb = server.gui.add_rgb(\"Color\", initial_value=(90, 200, 255))\n gui_wireframe = server.gui.add_checkbox(\"Wireframe\", initial_value=False)\n gui_show_controls = server.gui.add_checkbox(\"Handles\", initial_value=False)\n\n gui_rgb.on_update(set_changed)\n gui_wireframe.on_update(set_changed)\n\n @gui_show_controls.on_update\n def _(_):\n for control in transform_controls:\n control.visible = gui_show_controls.value\n\n # GUI elements: shape parameters.\n with tab_group.add_tab(\"Shape\", viser.Icon.BOX):\n gui_reset_shape = server.gui.add_button(\"Reset Shape\")\n gui_random_shape = server.gui.add_button(\"Random Shape\")\n\n @gui_reset_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = 0.0\n\n @gui_random_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = onp.random.normal(loc=0.0, scale=1.0)\n\n gui_betas = []\n for i in range(num_betas):\n beta = server.gui.add_slider(\n f\"beta{i}\", min=-5.0, max=5.0, step=0.01, initial_value=0.0\n )\n gui_betas.append(beta)\n beta.on_update(set_changed)\n\n # GUI elements: joint angles.\n with tab_group.add_tab(\"Joints\", viser.Icon.ANGLE):\n gui_reset_joints = server.gui.add_button(\"Reset Joints\")\n gui_random_joints = server.gui.add_button(\"Random Joints\")\n\n @gui_reset_joints.on_click\n def _(_):\n for joint in gui_joints:\n joint.value = (0.0, 0.0, 0.0)\n\n @gui_random_joints.on_click\n def _(_):\n for joint in gui_joints:\n # It's hard to uniformly sample orientations directly in so(3), so we\n # first sample on S^3 and then convert.\n quat = onp.random.normal(loc=0.0, scale=1.0, size=(4,))\n quat /= onp.linalg.norm(quat)\n joint.value = tf.SO3(wxyz=quat).log()\n\n gui_joints: list[viser.GuiInputHandle[tuple[float, float, float]]] = []\n for i in range(num_joints):\n gui_joint = server.gui.add_vector3(\n label=f\"Joint {i}\",\n initial_value=(0.0, 0.0, 0.0),\n step=0.05,\n )\n gui_joints.append(gui_joint)\n\n def set_callback_in_closure(i: int) -> None:\n @gui_joint.on_update\n def _(_):\n transform_controls[i].wxyz = tf.SO3.exp(\n np.array(gui_joints[i].value)\n ).wxyz\n out.changed = True\n\n set_callback_in_closure(i)\n\n # Transform control gizmos on joints.\n transform_controls: list[viser.TransformControlsHandle] = []\n prefixed_joint_names = [] # Joint names, but prefixed with parents.\n for i in range(num_joints):\n prefixed_joint_name = f\"joint_{i}\"\n if i > 0:\n prefixed_joint_name = (\n prefixed_joint_names[parent_idx[i]] + \"/\" + prefixed_joint_name\n )\n prefixed_joint_names.append(prefixed_joint_name)\n controls = server.scene.add_transform_controls(\n f\"/smpl/{prefixed_joint_name}\",\n depth_test=False,\n scale=0.2 * (0.75 ** prefixed_joint_name.count(\"/\")),\n disable_axes=True,\n disable_sliders=True,\n visible=gui_show_controls.value,\n )\n transform_controls.append(controls)\n\n def set_callback_in_closure(i: int) -> None:\n @controls.on_update\n def _(_) -> None:\n axisangle = tf.SO3(transform_controls[i].wxyz).log()\n gui_joints[i].value = (axisangle[0], axisangle[1], axisangle[2])\n\n set_callback_in_closure(i)\n\n out = GuiElements(\n gui_rgb,\n gui_wireframe,\n gui_betas,\n gui_joints,\n transform_controls=transform_controls,\n changed=True,\n )\n return out\n\n\nif __name__ == \"__main__\":\n tyro.cli(main, description=__doc__)\n" }, { "path": "viser/examples/09_urdf_visualizer.py", "content": "\"\"\"Robot URDF visualizer\n\nRequires yourdfpy and robot_descriptions. Any URDF supported by yourdfpy should work.\n- https://github.com/robot-descriptions/robot_descriptions.py\n- https://github.com/clemense/yourdfpy\n\nThe :class:`viser.extras.ViserUrdf` is a lightweight interface between yourdfpy\nand viser. It can also take a path to a local URDF file as input.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport time\nfrom typing import Literal\n\nimport numpy as onp\nimport tyro\nfrom robot_descriptions.loaders.yourdfpy import load_robot_description\n\nimport viser\nfrom viser.extras import ViserUrdf\n\n\ndef create_robot_control_sliders(\n server: viser.ViserServer, viser_urdf: ViserUrdf\n) -> tuple[list[viser.GuiInputHandle[float]], list[float]]:\n \"\"\"Create slider for each joint of the robot. We also update robot model\n when slider moves.\"\"\"\n slider_handles: list[viser.GuiInputHandle[float]] = []\n initial_config: list[float] = []\n for joint_name, (\n lower,\n upper,\n ) in viser_urdf.get_actuated_joint_limits().items():\n lower = lower if lower is not None else -onp.pi\n upper = upper if upper is not None else onp.pi\n initial_pos = 0.0 if lower < 0 and upper > 0 else (lower + upper) / 2.0\n slider = server.gui.add_slider(\n label=joint_name,\n min=lower,\n max=upper,\n step=1e-3,\n initial_value=initial_pos,\n )\n slider.on_update( # When sliders move, we update the URDF configuration.\n lambda _: viser_urdf.update_cfg(\n onp.array([slider.value for slider in slider_handles])\n )\n )\n slider_handles.append(slider)\n initial_config.append(initial_pos)\n return slider_handles, initial_config\n\n\ndef main(\n robot_type: Literal[\n \"panda\",\n \"ur10\",\n \"cassie\",\n \"allegro_hand\",\n \"barrett_hand\",\n \"robotiq_2f85\",\n \"atlas_drc\",\n \"g1\",\n \"h1\",\n \"anymal_c\",\n \"go2\",\n ] = \"panda\",\n) -> None:\n # Start viser server.\n server = viser.ViserServer()\n\n # Load URDF.\n #\n # This takes either a yourdfpy.URDF object or a path to a .urdf file.\n viser_urdf = ViserUrdf(\n server,\n urdf_or_path=load_robot_description(robot_type + \"_description\"),\n )\n\n # Create sliders in GUI that help us move the robot joints.\n with server.gui.add_folder(\"Joint position control\"):\n (slider_handles, initial_config) = create_robot_control_sliders(\n server, viser_urdf\n )\n\n # Set initial robot configuration.\n viser_urdf.update_cfg(onp.array(initial_config))\n\n # Create joint reset button.\n reset_button = server.gui.add_button(\"Reset\")\n\n @reset_button.on_click\n def _(_):\n for s, init_q in zip(slider_handles, initial_config):\n s.value = init_q\n\n # Sleep forever.\n while True:\n time.sleep(10.0)\n\n\nif __name__ == \"__main__\":\n tyro.cli(main)\n" }, { "path": "viser/examples/10_realsense.py", "content": "\"\"\"RealSense visualizer\n\nConnect to a RealSense camera, then visualize RGB-D readings as a point clouds. Requires\npyrealsense2.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport contextlib\n\nimport numpy as np\nimport numpy.typing as npt\nimport pyrealsense2 as rs # type: ignore\nfrom tqdm.auto import tqdm\n\nimport viser\n\n\n@contextlib.contextmanager\ndef realsense_pipeline(fps: int = 30):\n \"\"\"Context manager that yields a RealSense pipeline.\"\"\"\n\n # Configure depth and color streams.\n pipeline = rs.pipeline() # type: ignore\n config = rs.config() # type: ignore\n\n pipeline_wrapper = rs.pipeline_wrapper(pipeline) # type: ignore\n config.resolve(pipeline_wrapper)\n\n config.enable_stream(rs.stream.depth, rs.format.z16, fps) # type: ignore\n config.enable_stream(rs.stream.color, rs.format.rgb8, fps) # type: ignore\n\n # Start streaming.\n pipeline.start(config)\n\n yield pipeline\n\n # Close pipeline when done.\n pipeline.close()\n\n\ndef point_cloud_arrays_from_frames(\n depth_frame, color_frame\n) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.uint8]]:\n \"\"\"Maps realsense frames to two arrays.\n\n Returns:\n - A point position array: (N, 3) float32.\n - A point color array: (N, 3) uint8.\n \"\"\"\n # Processing blocks. Could be tuned.\n point_cloud = rs.pointcloud() # type: ignore\n decimate = rs.decimation_filter() # type: ignore\n decimate.set_option(rs.option.filter_magnitude, 3) # type: ignore\n\n # Downsample depth frame.\n depth_frame = decimate.process(depth_frame)\n\n # Map texture and calculate points from frames. Uses frame intrinsics.\n point_cloud.map_to(color_frame)\n points = point_cloud.calculate(depth_frame)\n\n # Get color coordinates.\n texture_uv = (\n np.asanyarray(points.get_texture_coordinates())\n .view(np.float32)\n .reshape((-1, 2))\n )\n color_image = np.asanyarray(color_frame.get_data())\n color_h, color_w, _ = color_image.shape\n\n # Note: for points that aren't in the view of our RGB camera, we currently clamp to\n # the closes available RGB pixel. We could also just remove these points.\n texture_uv = texture_uv.clip(0.0, 1.0)\n\n # Get positions and colors.\n positions = np.asanyarray(points.get_vertices()).view(np.float32)\n positions = positions.reshape((-1, 3))\n colors = color_image[\n (texture_uv[:, 1] * (color_h - 1.0)).astype(np.int32),\n (texture_uv[:, 0] * (color_w - 1.0)).astype(np.int32),\n :,\n ]\n N = positions.shape[0]\n\n assert positions.shape == (N, 3)\n assert positions.dtype == np.float32\n assert colors.shape == (N, 3)\n assert colors.dtype == np.uint8\n\n return positions, colors\n\n\ndef main():\n # Start visualization server.\n server = viser.ViserServer()\n\n with realsense_pipeline() as pipeline:\n for i in tqdm(range(10000000)):\n # Wait for a coherent pair of frames: depth and color\n frames = pipeline.wait_for_frames()\n depth_frame = frames.get_depth_frame()\n color_frame = frames.get_color_frame()\n\n # Compute point cloud from frames.\n positions, colors = point_cloud_arrays_from_frames(depth_frame, color_frame)\n\n R = np.array(\n [\n [1.0, 0.0, 0.0],\n [0.0, 0.0, 1.0],\n [0.0, -1.0, 0.0],\n ],\n dtype=np.float32,\n )\n positions = positions @ R.T\n\n # Visualize.\n server.scene.add_point_cloud(\n \"/realsense\",\n points=positions * 10.0,\n colors=colors,\n point_size=0.1,\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/11_colmap_visualizer.py", "content": "\"\"\"COLMAP visualizer\n\nVisualize COLMAP sparse reconstruction outputs. To get demo data, see `./assets/download_colmap_garden.sh`.\n\"\"\"\n\nimport random\nimport time\nfrom pathlib import Path\n\nimport imageio.v3 as iio\nimport numpy as onp\nimport tyro\nfrom tqdm.auto import tqdm\n\nimport viser\nimport viser.transforms as tf\nfrom viser.extras.colmap import (\n read_cameras_binary,\n read_images_binary,\n read_points3d_binary,\n)\n\n\ndef main(\n colmap_path: Path = Path(__file__).parent / \"assets/colmap_garden/sparse/0\",\n images_path: Path = Path(__file__).parent / \"assets/colmap_garden/images_8\",\n downsample_factor: int = 2,\n) -> None:\n \"\"\"Visualize COLMAP sparse reconstruction outputs.\n\n Args:\n colmap_path: Path to the COLMAP reconstruction directory.\n images_path: Path to the COLMAP images directory.\n downsample_factor: Downsample factor for the images.\n \"\"\"\n server = viser.ViserServer()\n server.gui.configure_theme(titlebar_content=None, control_layout=\"collapsible\")\n\n # Load the colmap info.\n cameras = read_cameras_binary(colmap_path / \"cameras.bin\")\n images = read_images_binary(colmap_path / \"images.bin\")\n points3d = read_points3d_binary(colmap_path / \"points3D.bin\")\n gui_reset_up = server.gui.add_button(\n \"Reset up direction\",\n hint=\"Set the camera control 'up' direction to the current camera's 'up'.\",\n )\n\n @gui_reset_up.on_click\n def _(event: viser.GuiEvent) -> None:\n client = event.client\n assert client is not None\n client.camera.up_direction = tf.SO3(client.camera.wxyz) @ onp.array(\n [0.0, -1.0, 0.0]\n )\n\n gui_points = server.gui.add_slider(\n \"Max points\",\n min=1,\n max=len(points3d),\n step=1,\n initial_value=min(len(points3d), 50_000),\n )\n gui_frames = server.gui.add_slider(\n \"Max frames\",\n min=1,\n max=len(images),\n step=1,\n initial_value=min(len(images), 100),\n )\n gui_point_size = server.gui.add_number(\"Point size\", initial_value=0.05)\n\n def visualize_colmap() -> None:\n \"\"\"Send all COLMAP elements to viser for visualization. This could be optimized\n a ton!\"\"\"\n # Set the point cloud.\n points = onp.array([points3d[p_id].xyz for p_id in points3d])\n colors = onp.array([points3d[p_id].rgb for p_id in points3d])\n points_selection = onp.random.choice(\n points.shape[0], gui_points.value, replace=False\n )\n points = points[points_selection]\n colors = colors[points_selection]\n\n server.scene.add_point_cloud(\n name=\"/colmap/pcd\",\n points=points,\n colors=colors,\n point_size=gui_point_size.value,\n )\n\n # Interpret the images and cameras.\n img_ids = [im.id for im in images.values()]\n random.shuffle(img_ids)\n img_ids = sorted(img_ids[: gui_frames.value])\n\n def attach_callback(\n frustum: viser.CameraFrustumHandle, frame: viser.FrameHandle\n ) -> None:\n @frustum.on_click\n def _(_) -> None:\n for client in server.get_clients().values():\n client.camera.wxyz = frame.wxyz\n client.camera.position = frame.position\n\n for img_id in tqdm(img_ids):\n img = images[img_id]\n cam = cameras[img.camera_id]\n\n # Skip images that don't exist.\n image_filename = images_path / img.name\n if not image_filename.exists():\n continue\n\n T_world_camera = tf.SE3.from_rotation_and_translation(\n tf.SO3(img.qvec), img.tvec\n ).inverse()\n frame = server.scene.add_frame(\n f\"/colmap/frame_{img_id}\",\n wxyz=T_world_camera.rotation().wxyz,\n position=T_world_camera.translation(),\n axes_length=0.1,\n axes_radius=0.005,\n )\n\n # For pinhole cameras, cam.params will be (fx, fy, cx, cy).\n if cam.model != \"PINHOLE\":\n print(f\"Expected pinhole camera, but got {cam.model}\")\n\n H, W = cam.height, cam.width\n fy = cam.params[1]\n image = iio.imread(image_filename)\n image = image[::downsample_factor, ::downsample_factor]\n frustum = server.scene.add_camera_frustum(\n f\"/colmap/frame_{img_id}/frustum\",\n fov=2 * onp.arctan2(H / 2, fy),\n aspect=W / H,\n scale=0.15,\n image=image,\n )\n attach_callback(frustum, frame)\n\n need_update = True\n\n @gui_points.on_update\n def _(_) -> None:\n nonlocal need_update\n need_update = True\n\n @gui_frames.on_update\n def _(_) -> None:\n nonlocal need_update\n need_update = True\n\n @gui_point_size.on_update\n def _(_) -> None:\n nonlocal need_update\n need_update = True\n\n while True:\n if need_update:\n need_update = False\n\n server.scene.reset()\n visualize_colmap()\n\n time.sleep(1e-3)\n\n\nif __name__ == \"__main__\":\n tyro.cli(main)\n" }, { "path": "viser/examples/12_click_meshes.py", "content": "\"\"\"Mesh click events\n\nClick on meshes to select them. The index of the last clicked mesh is displayed in the GUI.\n\"\"\"\n\nimport time\n\nimport matplotlib\n\nimport viser\n\n\ndef main() -> None:\n grid_shape = (4, 5)\n server = viser.ViserServer()\n\n with server.gui.add_folder(\"Last clicked\"):\n x_value = server.gui.add_number(\n label=\"x\",\n initial_value=0,\n disabled=True,\n hint=\"x coordinate of the last clicked mesh\",\n )\n y_value = server.gui.add_number(\n label=\"y\",\n initial_value=0,\n disabled=True,\n hint=\"y coordinate of the last clicked mesh\",\n )\n\n def add_swappable_mesh(i: int, j: int) -> None:\n \"\"\"Simple callback that swaps between:\n - a gray box\n - a colored box\n - a colored sphere\n\n Color is chosen based on the position (i, j) of the mesh in the grid.\n \"\"\"\n\n colormap = matplotlib.colormaps[\"tab20\"]\n\n def create_mesh(counter: int) -> None:\n if counter == 0:\n color = (0.8, 0.8, 0.8)\n else:\n index = (i * grid_shape[1] + j) / (grid_shape[0] * grid_shape[1])\n color = colormap(index)[:3]\n\n if counter in (0, 1):\n handle = server.scene.add_box(\n name=f\"/sphere_{i}_{j}\",\n position=(i, j, 0.0),\n color=color,\n dimensions=(0.5, 0.5, 0.5),\n )\n else:\n handle = server.scene.add_icosphere(\n name=f\"/sphere_{i}_{j}\",\n radius=0.4,\n color=color,\n position=(i, j, 0.0),\n )\n\n @handle.on_click\n def _(_) -> None:\n x_value.value = i\n y_value.value = j\n\n # The new mesh will replace the old one because the names\n # /sphere_{i}_{j} are the same.\n create_mesh((counter + 1) % 3)\n\n create_mesh(0)\n\n for i in range(grid_shape[0]):\n for j in range(grid_shape[1]):\n add_swappable_mesh(i, j)\n\n while True:\n time.sleep(10.0)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/13_theming.py", "content": "\"\"\"Theming\n\nViser includes support for light theming.\n\"\"\"\n\nimport time\n\nimport viser\nfrom viser.theme import TitlebarButton, TitlebarConfig, TitlebarImage\n\n\ndef main():\n server = viser.ViserServer(label=\"Viser Theming\")\n\n buttons = (\n TitlebarButton(\n text=\"Getting Started\",\n icon=None,\n href=\"https://nerf.studio\",\n ),\n TitlebarButton(\n text=\"Github\",\n icon=\"GitHub\",\n href=\"https://github.com/nerfstudio-project/nerfstudio\",\n ),\n TitlebarButton(\n text=\"Documentation\",\n icon=\"Description\",\n href=\"https://docs.nerf.studio\",\n ),\n )\n image = TitlebarImage(\n image_url_light=\"https://docs.nerf.studio/_static/imgs/logo.png\",\n image_url_dark=\"https://docs.nerf.studio/_static/imgs/logo-dark.png\",\n image_alt=\"NerfStudio Logo\",\n href=\"https://docs.nerf.studio/\",\n )\n titlebar_theme = TitlebarConfig(buttons=buttons, image=image)\n\n server.gui.add_markdown(\n \"Viser includes support for light theming via the `.configure_theme()` method.\"\n )\n\n gui_theme_code = server.gui.add_markdown(\"no theme applied yet\")\n\n # GUI elements for controllable values.\n titlebar = server.gui.add_checkbox(\"Titlebar\", initial_value=True)\n dark_mode = server.gui.add_checkbox(\"Dark mode\", initial_value=True)\n show_logo = server.gui.add_checkbox(\"Show logo\", initial_value=True)\n show_share_button = server.gui.add_checkbox(\"Show share button\", initial_value=True)\n brand_color = server.gui.add_rgb(\"Brand color\", (230, 180, 30))\n control_layout = server.gui.add_dropdown(\n \"Control layout\", (\"floating\", \"fixed\", \"collapsible\")\n )\n control_width = server.gui.add_dropdown(\n \"Control width\", (\"small\", \"medium\", \"large\"), initial_value=\"medium\"\n )\n synchronize = server.gui.add_button(\"Apply theme\", icon=viser.Icon.CHECK)\n\n def synchronize_theme() -> None:\n server.gui.configure_theme(\n titlebar_content=titlebar_theme if titlebar.value else None,\n control_layout=control_layout.value,\n control_width=control_width.value,\n dark_mode=dark_mode.value,\n show_logo=show_logo.value,\n show_share_button=show_share_button.value,\n brand_color=brand_color.value,\n )\n gui_theme_code.content = f\"\"\"\n ### Current applied theme\n ```\n server.gui.configure_theme(\n titlebar_content={\"titlebar_content\" if titlebar.value else None},\n control_layout=\"{control_layout.value}\",\n control_width=\"{control_width.value}\",\n dark_mode={dark_mode.value},\n show_logo={show_logo.value},\n show_share_button={show_share_button.value},\n brand_color={brand_color.value},\n )\n ```\n \"\"\"\n\n synchronize.on_click(lambda _: synchronize_theme())\n synchronize_theme()\n\n while True:\n time.sleep(10.0)\n\n\n# main()\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/14_markdown.py", "content": "\"\"\"Markdown demonstration\n\nViser GUI has MDX 2 support.\n\"\"\"\n\nimport time\nfrom pathlib import Path\n\nimport viser\n\nserver = viser.ViserServer()\nserver.scene.world_axes.visible = True\n\nmarkdown_counter = server.gui.add_markdown(\"Counter: 0\")\n\nhere = Path(__file__).absolute().parent\n\nbutton = server.gui.add_button(\"Remove blurb\")\ncheckbox = server.gui.add_checkbox(\"Visibility\", initial_value=True)\n\nmarkdown_source = (here / \"./assets/mdx_example.mdx\").read_text()\nmarkdown_blurb = server.gui.add_markdown(\n content=markdown_source,\n image_root=here,\n)\n\n\n@button.on_click\ndef _(_):\n markdown_blurb.remove()\n\n\n@checkbox.on_update\ndef _(_):\n markdown_blurb.visible = checkbox.value\n\n\ncounter = 0\nwhile True:\n markdown_counter.content = f\"Counter: {counter}\"\n counter += 1\n time.sleep(0.1)\n" }, { "path": "viser/examples/15_gui_in_scene.py", "content": "\"\"\"3D GUI elements\n\n`add_3d_gui_container()` allows standard GUI elements to be incorporated directly into a\n3D scene. In this example, we click on coordinate frames to show actions that can be\nperformed on them.\n\"\"\"\n\nimport time\nfrom typing import Optional\n\nimport numpy as onp\n\nimport viser\nimport viser.transforms as tf\n\nserver = viser.ViserServer()\nserver.gui.configure_theme(dark_mode=True)\nnum_frames = 20\n\n\n@server.on_client_connect\ndef _(client: viser.ClientHandle) -> None:\n \"\"\"For each client that connects, we create a set of random frames + a click handler for each frame.\n\n When a frame is clicked, we display a 3D gui node.\n \"\"\"\n\n rng = onp.random.default_rng(0)\n\n displayed_3d_container: Optional[viser.Gui3dContainerHandle] = None\n\n def make_frame(i: int) -> None:\n # Sample a random orientation + position.\n wxyz = rng.normal(size=4)\n wxyz /= onp.linalg.norm(wxyz)\n position = rng.uniform(-3.0, 3.0, size=(3,))\n\n # Create a coordinate frame and label.\n frame = client.scene.add_frame(f\"/frame_{i}\", wxyz=wxyz, position=position)\n\n # Move the camera when we click a frame.\n @frame.on_click\n def _(_):\n nonlocal displayed_3d_container\n\n # Close previously opened GUI.\n if displayed_3d_container is not None:\n displayed_3d_container.remove()\n\n displayed_3d_container = client.scene.add_3d_gui_container(\n f\"/frame_{i}/gui\"\n )\n with displayed_3d_container:\n go_to = client.gui.add_button(\"Go to\")\n randomize_orientation = client.gui.add_button(\"Randomize orientation\")\n close = client.gui.add_button(\"Close GUI\")\n\n @go_to.on_click\n def _(_) -> None:\n T_world_current = tf.SE3.from_rotation_and_translation(\n tf.SO3(client.camera.wxyz), client.camera.position\n )\n T_world_target = tf.SE3.from_rotation_and_translation(\n tf.SO3(frame.wxyz), frame.position\n ) @ tf.SE3.from_translation(onp.array([0.0, 0.0, -0.5]))\n\n T_current_target = T_world_current.inverse() @ T_world_target\n\n for j in range(20):\n T_world_set = T_world_current @ tf.SE3.exp(\n T_current_target.log() * j / 19.0\n )\n\n # Important bit: we atomically set both the orientation and the position\n # of the camera.\n with client.atomic():\n client.camera.wxyz = T_world_set.rotation().wxyz\n client.camera.position = T_world_set.translation()\n time.sleep(1.0 / 60.0)\n\n # Mouse interactions should orbit around the frame origin.\n client.camera.look_at = frame.position\n\n @randomize_orientation.on_click\n def _(_) -> None:\n wxyz = rng.normal(size=4)\n wxyz /= onp.linalg.norm(wxyz)\n frame.wxyz = wxyz\n\n @close.on_click\n def _(_) -> None:\n nonlocal displayed_3d_container\n if displayed_3d_container is None:\n return\n displayed_3d_container.remove()\n displayed_3d_container = None\n\n for i in range(num_frames):\n make_frame(i)\n\n\nwhile True:\n time.sleep(1.0)\n" }, { "path": "viser/examples/16_modal.py", "content": "\"\"\"Modal basics\n\nExamples of using modals in Viser.\"\"\"\n\nimport time\n\nimport viser\n\n\ndef main():\n server = viser.ViserServer()\n\n @server.on_client_connect\n def _(client: viser.ClientHandle) -> None:\n with client.gui.add_modal(\"Modal example\"):\n client.gui.add_markdown(\n \"**The input below determines the title of the modal...**\"\n )\n\n gui_title = client.gui.add_text(\n \"Title\",\n initial_value=\"My Modal\",\n )\n\n modal_button = client.gui.add_button(\"Show more modals\")\n\n @modal_button.on_click\n def _(_) -> None:\n with client.gui.add_modal(gui_title.value) as modal:\n client.gui.add_markdown(\"This is content inside the modal!\")\n client.gui.add_button(\"Close\").on_click(lambda _: modal.close())\n\n while True:\n time.sleep(0.15)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/17_background_composite.py", "content": "\"\"\"Depth compositing\n\nIn this example, we show how to use a background image with depth compositing. This can\nbe useful when we want a 2D image to occlude 3D geometry, such as for NeRF rendering.\n\"\"\"\n\nimport time\n\nimport numpy as onp\nimport trimesh\nimport trimesh.creation\n\nimport viser\n\nserver = viser.ViserServer()\n\n\nimg = onp.random.randint(0, 255, size=(1000, 1000, 3), dtype=onp.uint8)\ndepth = onp.ones((1000, 1000, 1), dtype=onp.float32)\n\n# Make a square middle portal.\ndepth[250:750, 250:750, :] = 10.0\nimg[250:750, 250:750, :] = 255\n\nmesh = trimesh.creation.box((0.5, 0.5, 0.5))\nserver.scene.add_mesh_trimesh(\n name=\"/cube\",\n mesh=mesh,\n position=(0, 0, 0.0),\n)\nserver.scene.set_background_image(img, depth=depth)\n\n\nwhile True:\n time.sleep(1.0)\n" }, { "path": "viser/examples/18_splines.py", "content": "\"\"\"Splines\n\nMake a ball with some random splines.\n\"\"\"\n\nimport time\n\nimport numpy as onp\n\nimport viser\n\n\ndef main() -> None:\n server = viser.ViserServer()\n for i in range(10):\n positions = onp.random.normal(size=(30, 3)) * 3.0\n server.scene.add_spline_catmull_rom(\n f\"/catmull_{i}\",\n positions,\n tension=0.5,\n line_width=3.0,\n color=onp.random.uniform(size=3),\n segments=100,\n )\n\n control_points = onp.random.normal(size=(30 * 2 - 2, 3)) * 3.0\n server.scene.add_spline_cubic_bezier(\n f\"/cubic_bezier_{i}\",\n positions,\n control_points,\n line_width=3.0,\n color=onp.random.uniform(size=3),\n segments=100,\n )\n\n while True:\n time.sleep(10.0)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/19_get_renders.py", "content": "\"\"\"Get renders\n\nExample for getting renders from a client's viewport to the Python API.\"\"\"\n\nimport time\n\nimport imageio.v3 as iio\nimport numpy as onp\n\nimport viser\n\n\ndef main():\n server = viser.ViserServer()\n\n button = server.gui.add_button(\"Render a GIF\")\n\n @button.on_click\n def _(event: viser.GuiEvent) -> None:\n client = event.client\n assert client is not None\n\n client.scene.reset()\n\n images = []\n\n for i in range(20):\n positions = onp.random.normal(size=(30, 3)) * 3.0\n client.scene.add_spline_catmull_rom(\n f\"/catmull_{i}\",\n positions,\n tension=0.5,\n line_width=3.0,\n color=onp.random.uniform(size=3),\n )\n images.append(client.camera.get_render(height=720, width=1280))\n\n print(\"Generating and sending GIF...\")\n client.send_file_download(\n \"image.gif\", iio.imwrite(\"\", images, extension=\".gif\")\n )\n print(\"Done!\")\n\n while True:\n time.sleep(10.0)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/20_scene_pointer.py", "content": "\"\"\"Scene pointer events.\n\nThis example shows how to use scene pointer events to specify rays, and how they can be\nused to interact with the scene (e.g., ray-mesh intersections).\n\nTo get the demo data, see `./assets/download_dragon_mesh.sh`.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport time\nfrom pathlib import Path\nfrom typing import cast\n\nimport numpy as onp\nimport trimesh\nimport trimesh.creation\nimport trimesh.ray\n\nimport viser\nimport viser.transforms as tf\nfrom viser.theme import TitlebarConfig\n\nserver = viser.ViserServer()\nserver.gui.configure_theme(\n brand_color=(130, 0, 150),\n titlebar_content=TitlebarConfig(buttons=(), image=None),\n)\nserver.scene.set_up_direction(\"+y\")\n\nmesh = cast(\n trimesh.Trimesh, trimesh.load_mesh(str(Path(__file__).parent / \"assets/dragon.obj\"))\n)\nmesh.apply_scale(0.05)\n\nmesh_handle = server.scene.add_mesh_trimesh(\n name=\"/mesh\",\n mesh=mesh,\n position=(0.0, 0.0, 0.0),\n)\n\nhit_pos_handles: list[viser.GlbHandle] = []\n\n\n# Buttons + callbacks will operate on a per-client basis, but will modify the global scene! :)\n@server.on_client_connect\ndef _(client: viser.ClientHandle) -> None:\n # Set up the camera -- this gives a nice view of the full mesh.\n client.camera.position = onp.array([0.0, 0.0, -10.0])\n client.camera.wxyz = onp.array([0.0, 0.0, 0.0, 1.0])\n\n # Tests \"click\" scenepointerevent.\n click_button_handle = client.gui.add_button(\"Add sphere\", icon=viser.Icon.POINTER)\n\n @click_button_handle.on_click\n def _(_):\n click_button_handle.disabled = True\n\n @client.scene.on_pointer_event(event_type=\"click\")\n def _(event: viser.ScenePointerEvent) -> None:\n # Check for intersection with the mesh, using trimesh's ray-mesh intersection.\n # Note that mesh is in the mesh frame, so we need to transform the ray.\n R_world_mesh = tf.SO3(mesh_handle.wxyz)\n R_mesh_world = R_world_mesh.inverse()\n origin = (R_mesh_world @ onp.array(event.ray_origin)).reshape(1, 3)\n direction = (R_mesh_world @ onp.array(event.ray_direction)).reshape(1, 3)\n intersector = trimesh.ray.ray_triangle.RayMeshIntersector(mesh)\n hit_pos, _, _ = intersector.intersects_location(origin, direction)\n\n if len(hit_pos) == 0:\n return\n client.scene.remove_pointer_callback()\n\n # Get the first hit position (based on distance from the ray origin).\n hit_pos = hit_pos[onp.argmin(onp.sum((hit_pos - origin) ** 2, axis=-1))]\n\n # Create a sphere at the hit location.\n hit_pos_mesh = trimesh.creation.icosphere(radius=0.1)\n hit_pos_mesh.vertices += R_world_mesh @ hit_pos\n hit_pos_mesh.visual.vertex_colors = (0.5, 0.0, 0.7, 1.0) # type: ignore\n hit_pos_handle = server.scene.add_mesh_trimesh(\n name=f\"/hit_pos_{len(hit_pos_handles)}\", mesh=hit_pos_mesh\n )\n hit_pos_handles.append(hit_pos_handle)\n\n @client.scene.on_pointer_callback_removed\n def _():\n click_button_handle.disabled = False\n\n # Tests \"rect-select\" scenepointerevent.\n paint_button_handle = client.gui.add_button(\"Paint mesh\", icon=viser.Icon.PAINT)\n\n @paint_button_handle.on_click\n def _(_):\n paint_button_handle.disabled = True\n\n @client.scene.on_pointer_event(event_type=\"rect-select\")\n def _(message: viser.ScenePointerEvent) -> None:\n client.scene.remove_pointer_callback()\n\n global mesh_handle\n camera = message.client.camera\n\n # Put the mesh in the camera frame.\n R_world_mesh = tf.SO3(mesh_handle.wxyz)\n R_mesh_world = R_world_mesh.inverse()\n R_camera_world = tf.SE3.from_rotation_and_translation(\n tf.SO3(camera.wxyz), camera.position\n ).inverse()\n vertices = cast(onp.ndarray, mesh.vertices)\n vertices = (R_mesh_world.as_matrix() @ vertices.T).T\n vertices = (\n R_camera_world.as_matrix()\n @ onp.hstack([vertices, onp.ones((vertices.shape[0], 1))]).T\n ).T[:, :3]\n\n # Get the camera intrinsics, and project the vertices onto the image plane.\n fov, aspect = camera.fov, camera.aspect\n vertices_proj = vertices[:, :2] / vertices[:, 2].reshape(-1, 1)\n vertices_proj /= onp.tan(fov / 2)\n vertices_proj[:, 0] /= aspect\n\n # Move the origin to the upper-left corner, and scale to [0, 1].\n # ... make sure to match the OpenCV's image coordinates!\n vertices_proj = (1 + vertices_proj) / 2\n\n # Select the vertices that lie inside the 2D selected box, once projected.\n mask = (\n (vertices_proj > onp.array(message.screen_pos[0]))\n & (vertices_proj < onp.array(message.screen_pos[1]))\n ).all(axis=1)[..., None]\n\n # Update the mesh color based on whether the vertices are inside the box\n mesh.visual.vertex_colors = onp.where( # type: ignore\n mask, (0.5, 0.0, 0.7, 1.0), (0.9, 0.9, 0.9, 1.0)\n )\n mesh_handle = server.scene.add_mesh_trimesh(\n name=\"/mesh\",\n mesh=mesh,\n position=(0.0, 0.0, 0.0),\n )\n\n @client.scene.on_pointer_callback_removed\n def _():\n paint_button_handle.disabled = False\n\n # Button to clear spheres.\n clear_button_handle = client.gui.add_button(\"Clear scene\", icon=viser.Icon.X)\n\n @clear_button_handle.on_click\n def _(_):\n \"\"\"Reset the mesh color and remove all click-generated spheres.\"\"\"\n global mesh_handle\n for handle in hit_pos_handles:\n handle.remove()\n hit_pos_handles.clear()\n mesh.visual.vertex_colors = (0.9, 0.9, 0.9, 1.0) # type: ignore\n mesh_handle = server.scene.add_mesh_trimesh(\n name=\"/mesh\",\n mesh=mesh,\n position=(0.0, 0.0, 0.0),\n )\n\n\nwhile True:\n time.sleep(10.0)\n" }, { "path": "viser/examples/21_set_up_direction.py", "content": "\"\"\"Set up direction\n\n`.set_up_direction()` can help us set the global up direction.\"\"\"\n\nimport time\n\nimport viser\n\n\ndef main() -> None:\n server = viser.ViserServer()\n server.scene.world_axes.visible = True\n gui_up = server.gui.add_vector3(\n \"Up Direction\",\n initial_value=(0.0, 0.0, 1.0),\n step=0.01,\n )\n\n @gui_up.on_update\n def _(_) -> None:\n server.scene.set_up_direction(gui_up.value)\n\n while True:\n time.sleep(1.0)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/22_games.py", "content": "\"\"\"Games\n\nSome two-player games implemented using scene click events.\"\"\"\n\nimport time\nfrom typing import Literal\n\nimport numpy as onp\nimport trimesh.creation\nfrom typing_extensions import assert_never\n\nimport viser\nimport viser.transforms as tf\n\n\ndef main() -> None:\n server = viser.ViserServer()\n server.gui.configure_theme(dark_mode=True)\n play_connect_4(server)\n\n server.gui.add_button(\"Tic-Tac-Toe\").on_click(lambda _: play_tic_tac_toe(server))\n server.gui.add_button(\"Connect 4\").on_click(lambda _: play_connect_4(server))\n\n while True:\n time.sleep(10.0)\n\n\ndef play_connect_4(server: viser.ViserServer) -> None:\n \"\"\"Play a game of Connect 4.\"\"\"\n server.scene.reset()\n\n num_rows = 6\n num_cols = 7\n\n whose_turn: Literal[\"red\", \"yellow\"] = \"red\"\n pieces_in_col = [0] * num_cols\n\n # Create the board frame.\n for col in range(num_cols):\n for row in range(num_rows):\n server.scene.add_mesh_trimesh(\n f\"/structure/{row}_{col}\",\n trimesh.creation.annulus(0.45, 0.55, 0.125),\n position=(0.0, col, row),\n wxyz=tf.SO3.from_y_radians(onp.pi / 2.0).wxyz,\n )\n\n # Create a sphere to click on for each column.\n def setup_column(col: int) -> None:\n sphere = server.scene.add_icosphere(\n f\"/spheres/{col}\",\n radius=0.25,\n position=(0, col, num_rows - 0.25),\n color=(255, 255, 255),\n )\n\n # Drop piece into the column.\n @sphere.on_click\n def _(_) -> None:\n nonlocal whose_turn\n whose_turn = \"red\" if whose_turn != \"red\" else \"yellow\"\n\n row = pieces_in_col[col]\n if row == num_rows - 1:\n sphere.remove()\n\n pieces_in_col[col] += 1\n cylinder = trimesh.creation.cylinder(radius=0.4, height=0.125)\n piece = server.scene.add_mesh_simple(\n f\"/game_pieces/{row}_{col}\",\n cylinder.vertices,\n cylinder.faces,\n wxyz=tf.SO3.from_y_radians(onp.pi / 2.0).wxyz,\n color={\"red\": (255, 0, 0), \"yellow\": (255, 255, 0)}[whose_turn],\n )\n for row_anim in onp.linspace(num_rows - 1, row, num_rows - row + 1):\n piece.position = (\n 0,\n col,\n row_anim,\n )\n time.sleep(1.0 / 30.0)\n\n for col in range(num_cols):\n setup_column(col)\n\n\ndef play_tic_tac_toe(server: viser.ViserServer) -> None:\n \"\"\"Play a game of tic-tac-toe.\"\"\"\n server.scene.reset()\n\n whose_turn: Literal[\"x\", \"o\"] = \"x\"\n\n for i in range(4):\n server.scene.add_spline_catmull_rom(\n f\"/gridlines/{i}\",\n ((-0.5, -1.5, 0), (-0.5, 1.5, 0)),\n color=(127, 127, 127),\n position=(1, 1, 0),\n wxyz=tf.SO3.from_z_radians(onp.pi / 2 * i).wxyz,\n )\n\n def draw_symbol(symbol: Literal[\"x\", \"o\"], i: int, j: int) -> None:\n \"\"\"Draw an X or O in the given cell.\"\"\"\n for scale in onp.linspace(0.01, 1.0, 5):\n if symbol == \"x\":\n for k in range(2):\n server.scene.add_box(\n f\"/symbols/{i}_{j}/{k}\",\n dimensions=(0.7 * scale, 0.125 * scale, 0.125),\n position=(i, j, 0),\n color=(0, 0, 255),\n wxyz=tf.SO3.from_z_radians(\n onp.pi / 2.0 * k + onp.pi / 4.0\n ).wxyz,\n )\n elif symbol == \"o\":\n mesh = trimesh.creation.annulus(0.25 * scale, 0.35 * scale, 0.125)\n server.scene.add_mesh_simple(\n f\"/symbols/{i}_{j}\",\n mesh.vertices,\n mesh.faces,\n position=(i, j, 0),\n color=(255, 0, 0),\n )\n else:\n assert_never(symbol)\n server.flush()\n time.sleep(1.0 / 30.0)\n\n def setup_cell(i: int, j: int) -> None:\n \"\"\"Create a clickable sphere in a given cell.\"\"\"\n sphere = server.scene.add_icosphere(\n f\"/spheres/{i}_{j}\",\n radius=0.25,\n position=(i, j, 0),\n color=(255, 255, 255),\n )\n\n @sphere.on_click\n def _(_) -> None:\n nonlocal whose_turn\n whose_turn = \"x\" if whose_turn != \"x\" else \"o\"\n sphere.remove()\n draw_symbol(whose_turn, i, j)\n\n for i in range(3):\n for j in range(3):\n setup_cell(i, j)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/23_plotly.py", "content": "\"\"\"Plotly\n\nExamples of visualizing plotly plots in Viser.\"\"\"\n\nimport time\n\nimport numpy as onp\nimport plotly.express as px\nimport plotly.graph_objects as go\nfrom PIL import Image\n\nimport viser\n\n\ndef create_sinusoidal_wave(t: float) -> go.Figure:\n \"\"\"Create a sinusoidal wave plot, starting at time t.\"\"\"\n x_data = onp.linspace(t, t + 6 * onp.pi, 50)\n y_data = onp.sin(x_data) * 10\n\n fig = px.line(\n x=list(x_data),\n y=list(y_data),\n labels={\"x\": \"x\", \"y\": \"sin(x)\"},\n title=\"Sinusoidal Wave\",\n )\n\n # this sets the margins to be tight around the title.\n fig.layout.title.automargin = True # type: ignore\n fig.update_layout(\n margin=dict(l=20, r=20, t=20, b=20),\n ) # Reduce plot margins.\n\n return fig\n\n\ndef main() -> None:\n server = viser.ViserServer()\n\n # Plot type 1: Line plot.\n line_plot_time = 0.0\n line_plot = server.gui.add_plotly(figure=create_sinusoidal_wave(line_plot_time))\n\n # Plot type 2: Image plot.\n fig = px.imshow(Image.open(\"assets/Cal_logo.png\"))\n fig.update_layout(\n margin=dict(l=20, r=20, t=20, b=20),\n )\n server.gui.add_plotly(figure=fig, aspect=1.0)\n\n # Plot type 3: 3D Scatter plot.\n fig = px.scatter_3d(\n px.data.iris(),\n x=\"sepal_length\",\n y=\"sepal_width\",\n z=\"petal_width\",\n color=\"species\",\n )\n fig.update_layout(legend=dict(yanchor=\"top\", y=0.99, xanchor=\"left\", x=0.01))\n fig.update_layout(\n margin=dict(l=20, r=20, t=20, b=20),\n )\n server.gui.add_plotly(figure=fig, aspect=1.0)\n\n while True:\n # Update the line plot.\n line_plot_time += 0.1\n line_plot.figure = create_sinusoidal_wave(line_plot_time)\n\n time.sleep(0.01)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/24_notification.py", "content": "\"\"\"Notifications\n\nExamples of adding notifications per client in Viser.\"\"\"\n\nimport time\n\nimport viser\n\n\ndef main() -> None:\n server = viser.ViserServer()\n\n persistent_notif_button = server.gui.add_button(\n \"Show persistent notification (default)\"\n )\n timed_notif_button = server.gui.add_button(\"Show timed notification\")\n controlled_notif_button = server.gui.add_button(\"Show controlled notification\")\n loading_notif_button = server.gui.add_button(\"Show loading notification\")\n\n remove_controlled_notif = server.gui.add_button(\"Remove controlled notification\")\n\n @persistent_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show persistent notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n client.add_notification(\n title=\"Persistent notification\",\n body=\"This can be closed manually and does not disappear on its own!\",\n loading=False,\n with_close_button=True,\n auto_close=False,\n )\n\n @timed_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show timed notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n client.add_notification(\n title=\"Timed notification\",\n body=\"This disappears automatically after 5 seconds!\",\n loading=False,\n with_close_button=True,\n auto_close=5000,\n )\n\n @controlled_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show controlled notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n controlled_notif = client.add_notification(\n title=\"Controlled notification\",\n body=\"This cannot be closed by the user and is controlled in code only!\",\n loading=False,\n with_close_button=False,\n auto_close=False,\n )\n\n @remove_controlled_notif.on_click\n def _(_) -> None:\n \"\"\"Remove controlled notification.\"\"\"\n controlled_notif.remove()\n\n @loading_notif_button.on_click\n def _(event: viser.GuiEvent) -> None:\n \"\"\"Show loading notification when the button is clicked.\"\"\"\n client = event.client\n assert client is not None\n\n loading_notif = client.add_notification(\n title=\"Loading notification\",\n body=\"This indicates that some action is in progress! It will be updated in 3 seconds.\",\n loading=True,\n with_close_button=False,\n auto_close=False,\n )\n\n time.sleep(3.0)\n\n loading_notif.title = \"Updated notification\"\n loading_notif.body = \"This notification has been updated!\"\n loading_notif.loading = False\n loading_notif.with_close_button = True\n loading_notif.auto_close = 5000\n loading_notif.color = \"green\"\n\n while True:\n time.sleep(1.0)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "viser/examples/25_smpl_visualizer_skinned.py", "content": "# mypy: disable-error-code=\"assignment\"\n#\n# Asymmetric properties are supported in Pyright, but not yet in mypy.\n# - https://github.com/python/mypy/issues/3004\n# - https://github.com/python/mypy/pull/11643\n\"\"\"SMPL visualizer (Skinned Mesh)\n\nRequires a .npz model file.\n\nSee here for download instructions:\n https://github.com/vchoutas/smplx?tab=readme-ov-file#downloading-the-model\n\"\"\"\n\nfrom __future__ import annotations\n\nimport time\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import List, Tuple\n\nimport numpy as np\nimport numpy as onp\nimport tyro\n\nimport viser\nimport viser.transforms as tf\n\n\n@dataclass(frozen=True)\nclass SmplOutputs:\n vertices: np.ndarray\n faces: np.ndarray\n T_world_joint: np.ndarray # (num_joints, 4, 4)\n T_parent_joint: np.ndarray # (num_joints, 4, 4)\n\n\nclass SmplHelper:\n \"\"\"Helper for models in the SMPL family, implemented in numpy.\"\"\"\n\n def __init__(self, model_path: Path) -> None:\n assert model_path.suffix.lower() == \".npz\", \"Model should be an .npz file!\"\n body_dict = dict(**onp.load(model_path, allow_pickle=True))\n\n self._J_regressor = body_dict[\"J_regressor\"]\n self._weights = body_dict[\"weights\"]\n self._v_template = body_dict[\"v_template\"]\n self._posedirs = body_dict[\"posedirs\"]\n self._shapedirs = body_dict[\"shapedirs\"]\n self._faces = body_dict[\"f\"]\n\n self.num_joints: int = self._weights.shape[-1]\n self.num_betas: int = self._shapedirs.shape[-1]\n self.parent_idx: np.ndarray = body_dict[\"kintree_table\"][0]\n\n def get_outputs(self, betas: np.ndarray, joint_rotmats: np.ndarray) -> SmplOutputs:\n # Get shaped vertices + joint positions, when all local poses are identity.\n v_tpose = self._v_template + np.einsum(\"vxb,b->vx\", self._shapedirs, betas)\n j_tpose = np.einsum(\"jv,vx->jx\", self._J_regressor, v_tpose)\n\n # Local SE(3) transforms.\n T_parent_joint = np.zeros((self.num_joints, 4, 4)) + np.eye(4)\n T_parent_joint[:, :3, :3] = joint_rotmats\n T_parent_joint[0, :3, 3] = j_tpose[0]\n T_parent_joint[1:, :3, 3] = j_tpose[1:] - j_tpose[self.parent_idx[1:]]\n\n # Forward kinematics.\n T_world_joint = T_parent_joint.copy()\n for i in range(1, self.num_joints):\n T_world_joint[i] = T_world_joint[self.parent_idx[i]] @ T_parent_joint[i]\n\n # Linear blend skinning.\n pose_delta = (joint_rotmats[1:, ...] - np.eye(3)).flatten()\n v_blend = v_tpose + np.einsum(\"byn,n->by\", self._posedirs, pose_delta)\n v_delta = np.ones((v_blend.shape[0], self.num_joints, 4))\n v_delta[:, :, :3] = v_blend[:, None, :] - j_tpose[None, :, :]\n v_posed = np.einsum(\n \"jxy,vj,vjy->vx\", T_world_joint[:, :3, :], self._weights, v_delta\n )\n return SmplOutputs(v_posed, self._faces, T_world_joint, T_parent_joint)\n\n\ndef main(model_path: Path) -> None:\n server = viser.ViserServer()\n server.scene.set_up_direction(\"+y\")\n server.gui.configure_theme(control_layout=\"collapsible\")\n\n # Main loop. We'll read pose/shape from the GUI elements, compute the mesh,\n # and then send the updated mesh in a loop.\n model = SmplHelper(model_path)\n gui_elements = make_gui_elements(\n server,\n num_betas=model.num_betas,\n num_joints=model.num_joints,\n parent_idx=model.parent_idx,\n )\n smpl_outputs = model.get_outputs(\n betas=np.array([x.value for x in gui_elements.gui_betas]),\n joint_rotmats=onp.zeros((model.num_joints, 3, 3)) + onp.eye(3),\n )\n\n bone_wxyzs = np.array(\n [tf.SO3.from_matrix(R).wxyz for R in smpl_outputs.T_world_joint[:, :3, :3]]\n )\n bone_positions = smpl_outputs.T_world_joint[:, :3, 3]\n\n skinned_handle = server.scene.add_mesh_skinned(\n \"/human\",\n smpl_outputs.vertices,\n smpl_outputs.faces,\n bone_wxyzs=bone_wxyzs,\n bone_positions=bone_positions,\n skin_weights=model._weights,\n wireframe=gui_elements.gui_wireframe.value,\n color=gui_elements.gui_rgb.value,\n )\n\n while True:\n # Do nothing if no change.\n time.sleep(0.02)\n if not gui_elements.changed:\n continue\n\n gui_elements.changed = False\n\n # Compute SMPL outputs.\n smpl_outputs = model.get_outputs(\n betas=np.array([x.value for x in gui_elements.gui_betas]),\n joint_rotmats=np.stack(\n [\n tf.SO3.exp(np.array(x.value)).as_matrix()\n for x in gui_elements.gui_joints\n ],\n axis=0,\n ),\n )\n\n # Match transform control gizmos to joint positions.\n for i, control in enumerate(gui_elements.transform_controls):\n control.position = smpl_outputs.T_parent_joint[i, :3, 3]\n skinned_handle.bones[i].wxyz = tf.SO3.from_matrix(\n smpl_outputs.T_world_joint[i, :3, :3]\n ).wxyz\n skinned_handle.bones[i].position = smpl_outputs.T_world_joint[i, :3, 3]\n\n\n@dataclass\nclass GuiElements:\n \"\"\"Structure containing handles for reading from GUI elements.\"\"\"\n\n gui_rgb: viser.GuiInputHandle[Tuple[int, int, int]]\n gui_wireframe: viser.GuiInputHandle[bool]\n gui_betas: List[viser.GuiInputHandle[float]]\n gui_joints: List[viser.GuiInputHandle[Tuple[float, float, float]]]\n transform_controls: List[viser.TransformControlsHandle]\n\n changed: bool\n \"\"\"This flag will be flipped to True whenever the mesh needs to be re-generated.\"\"\"\n\n\ndef make_gui_elements(\n server: viser.ViserServer,\n num_betas: int,\n num_joints: int,\n parent_idx: np.ndarray,\n) -> GuiElements:\n \"\"\"Make GUI elements for interacting with the model.\"\"\"\n\n tab_group = server.gui.add_tab_group()\n\n def set_changed(_) -> None:\n out.changed = True # out is define later!\n\n # GUI elements: mesh settings + visibility.\n with tab_group.add_tab(\"View\", viser.Icon.VIEWFINDER):\n gui_rgb = server.gui.add_rgb(\"Color\", initial_value=(90, 200, 255))\n gui_wireframe = server.gui.add_checkbox(\"Wireframe\", initial_value=False)\n gui_show_controls = server.gui.add_checkbox(\"Handles\", initial_value=True)\n\n gui_rgb.on_update(set_changed)\n gui_wireframe.on_update(set_changed)\n\n @gui_show_controls.on_update\n def _(_):\n for control in transform_controls:\n control.visible = gui_show_controls.value\n\n # GUI elements: shape parameters.\n with tab_group.add_tab(\"Shape\", viser.Icon.BOX):\n gui_reset_shape = server.gui.add_button(\"Reset Shape\")\n gui_random_shape = server.gui.add_button(\"Random Shape\")\n\n @gui_reset_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = 0.0\n\n @gui_random_shape.on_click\n def _(_):\n for beta in gui_betas:\n beta.value = onp.random.normal(loc=0.0, scale=1.0)\n\n gui_betas = []\n for i in range(num_betas):\n beta = server.gui.add_slider(\n f\"beta{i}\", min=-5.0, max=5.0, step=0.01, initial_value=0.0\n )\n gui_betas.append(beta)\n beta.on_update(set_changed)\n\n # GUI elements: joint angles.\n with tab_group.add_tab(\"Joints\", viser.Icon.ANGLE):\n gui_reset_joints = server.gui.add_button(\"Reset Joints\")\n gui_random_joints = server.gui.add_button(\"Random Joints\")\n\n @gui_reset_joints.on_click\n def _(_):\n for joint in gui_joints:\n joint.value = (0.0, 0.0, 0.0)\n\n @gui_random_joints.on_click\n def _(_):\n for joint in gui_joints:\n # It's hard to uniformly sample orientations directly in so(3), so we\n # first sample on S^3 and then convert.\n quat = onp.random.normal(loc=0.0, scale=1.0, size=(4,))\n quat /= onp.linalg.norm(quat)\n joint.value = tf.SO3(wxyz=quat).log()\n\n gui_joints: List[viser.GuiInputHandle[Tuple[float, float, float]]] = []\n for i in range(num_joints):\n gui_joint = server.gui.add_vector3(\n label=f\"Joint {i}\",\n initial_value=(0.0, 0.0, 0.0),\n step=0.05,\n )\n gui_joints.append(gui_joint)\n\n def set_callback_in_closure(i: int) -> None:\n @gui_joint.on_update\n def _(_):\n transform_controls[i].wxyz = tf.SO3.exp(\n np.array(gui_joints[i].value)\n ).wxyz\n out.changed = True\n\n set_callback_in_closure(i)\n\n # Transform control gizmos on joints.\n transform_controls: List[viser.TransformControlsHandle] = []\n prefixed_joint_names = [] # Joint names, but prefixed with parents.\n for i in range(num_joints):\n prefixed_joint_name = f\"joint_{i}\"\n if i > 0:\n prefixed_joint_name = (\n prefixed_joint_names[parent_idx[i]] + \"/\" + prefixed_joint_name\n )\n prefixed_joint_names.append(prefixed_joint_name)\n controls = server.scene.add_transform_controls(\n f\"/smpl/{prefixed_joint_name}\",\n depth_test=False,\n scale=0.2 * (0.75 ** prefixed_joint_name.count(\"/\")),\n disable_axes=True,\n disable_sliders=True,\n visible=gui_show_controls.value,\n )\n transform_controls.append(controls)\n\n def set_callback_in_closure(i: int) -> None:\n @controls.on_update\n def _(_) -> None:\n axisangle = tf.SO3(transform_controls[i].wxyz).log()\n gui_joints[i].value = (axisangle[0], axisangle[1], axisangle[2])\n\n set_callback_in_closure(i)\n\n out = GuiElements(\n gui_rgb,\n gui_wireframe,\n gui_betas,\n gui_joints,\n transform_controls=transform_controls,\n changed=True,\n )\n return out\n\n\nif __name__ == \"__main__\":\n tyro.cli(main, description=__doc__)\n" }, { "path": "viser/examples/assets/.gitignore", "content": "dragon.obj\n/record3d_dance/\n/colmap_garden/\n" }, { "path": "viser/examples/assets/download_colmap_garden.sh", "content": "# This downloads the COLMAP model for the MIP-NeRF garden dataset\n# with the images that are downscaled by a factor of 8.\n# The full dataset is available at https://jonbarron.info/mipnerf360/.\n\nset -e -x\n\ngdown \"https://drive.google.com/uc?id=1wYHdrgwXPHtREdCjItvt4gqRQGISMade\"\n\nmkdir -p colmap_garden\n# shellcheck disable=SC2035\nunzip *.zip && rm *.zip" }, { "path": "viser/examples/assets/download_dragon_mesh.sh", "content": "set -e -x\n\ngdown \"https://drive.google.com/uc?id=1uRDvoS_l2Or8g8YDDPYV79K6_RfFYBeF\"\n" }, { "path": "viser/examples/assets/download_record3d_dance.sh", "content": "set -e -x\n\ngdown \"https://drive.google.com/uc?id=1_vd5bK_MhtlfisA6BkK1IgiJNfDbIntq\"\n\nmkdir -p record3d_dance\n# shellcheck disable=SC2035\nunzip *.r3d -d record3d_dance && rm *.r3d\n" }, { "path": "viser/examples/assets/mdx_example.mdx", "content": "## Markdown in Viser\n\n---\n\nViser has full support for the GFM markdown spec, including **bold**, _italics_,\n~~strikethrough~~, and many other features.\n\nHere's a [masked link](https://github.com/nerfstudio-project/viser). Not a fan?\nHere's a normal one: https://pypi.org/project/viser/\n\nAnywhere where you can insert GUI elements, you can also insert `images`,\n`blockquotes`, `lists`, `tables`, `task lists`, and `(unstyled) code blocks`.\n\nIn inline code blocks, you can show off colors with color chips: `#FED363`\n`hsl(0, 0%, 82%)` `rgb(255, 255, 255)`\n\nAdding images from a remote origin is simple.\n\n![Viser Logo](https://viser.studio/latest/_static/logo.svg)\n\nFor local images with relative paths, you can either directly use a\n[data URL](https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URLs)\nor set the `image_root` argument to the `Path` object that you'd like your paths\nrelative to. If no such `image_root` is provided, the file system will be scoped\nto the directory that Viser is installed in.\n\n![Cal Logo](../examples/assets/Cal_logo.png)\n\nTables follow the standard markdown spec:\n\n| Application | Description |\n| ---------------------------------------------------- | -------------------------------------------------- |\n| [NS](https://nerf.studio) | A collaboration friendly studio for NeRFs |\n| [Viser](https://nerfstudio-project.github.io/viser/) | An interactive 3D visualization toolbox for Python |\n\nCode blocks, while being not nearly as exciting as some of the things presented,\nwork as expected. Currently, while you can specify a language and metadata in\nyour blocks, they will remain unused by the Markdown renderer.\n\n```python\n\"\"\"Markdown Demonstration\n\nViser GUI has MDX 2 support.\n\"\"\"\n\nimport time\nfrom pathlib import Path\n\nimport viser\n\nserver = viser.ViserServer()\nserver.world_axes.visible = True\n\n\n@server.on_client_connect\ndef _(client: viser.ClientHandle) -> None:\n with open(\"./assets/mdx_example.mdx\", \"r\") as mkdn:\n markdown = client.gui.add_markdown(\n markdown=mkdn.read(), image_root=Path(__file__).parent\n )\n\n button = client.gui.add_button(\"Remove Markdown\")\n\n @button.on_click\n def _(_):\n markdown.remove()\n\n\nwhile True:\n time.sleep(10.0)\n```\n\nAs a bonus, MDX is extensible and JS capable. This means that you have the\nfreedom to do things like:\n\nThis page loaded on {(new Date()).toString()}\n\nOr:\n\n> Oh yes, mdx PR would be exciting\n>\n> — Brent Yi \n\n**Note**: Be careful when playing with JSX, it's very easy to break markdown.\n\nSo that's MDX in Viser. It has support for:\n\n- [x] CommonMark and GFM standards\n - bold, italics, strikethrough, images, blockquotes, tables, task lists, code\n blocks, inline code\n- [x] Color chips\n- [x] JSX enhanced components\n" }, { "path": "viser/examples/experimental/gaussian_splats.py", "content": "\"\"\"WebGL-based Gaussian splat rendering. This is still under developmentt.\"\"\"\n\nfrom __future__ import annotations\n\nimport time\nfrom pathlib import Path\nfrom typing import TypedDict\n\nimport numpy as onp\nimport numpy.typing as onpt\nimport tyro\nfrom plyfile import PlyData\n\nimport viser\nfrom viser import transforms as tf\n\n\nclass SplatFile(TypedDict):\n \"\"\"Data loaded from an antimatter15-style splat file.\"\"\"\n\n centers: onpt.NDArray[onp.floating]\n \"\"\"(N, 3).\"\"\"\n rgbs: onpt.NDArray[onp.floating]\n \"\"\"(N, 3). Range [0, 1].\"\"\"\n opacities: onpt.NDArray[onp.floating]\n \"\"\"(N, 1). Range [0, 1].\"\"\"\n covariances: onpt.NDArray[onp.floating]\n \"\"\"(N, 3, 3).\"\"\"\n\n\ndef load_splat_file(splat_path: Path, center: bool = False) -> SplatFile:\n \"\"\"Load an antimatter15-style splat file.\"\"\"\n start_time = time.time()\n splat_buffer = splat_path.read_bytes()\n bytes_per_gaussian = (\n # Each Gaussian is serialized as:\n # - position (vec3, float32)\n 3 * 4\n # - xyz (vec3, float32)\n + 3 * 4\n # - rgba (vec4, uint8)\n + 4\n # - ijkl (vec4, uint8), where 0 => -1, 255 => 1.\n + 4\n )\n assert len(splat_buffer) % bytes_per_gaussian == 0\n num_gaussians = len(splat_buffer) // bytes_per_gaussian\n\n # Reinterpret cast to dtypes that we want to extract.\n splat_uint8 = onp.frombuffer(splat_buffer, dtype=onp.uint8).reshape(\n (num_gaussians, bytes_per_gaussian)\n )\n scales = splat_uint8[:, 12:24].copy().view(onp.float32)\n wxyzs = splat_uint8[:, 28:32] / 255.0 * 2.0 - 1.0\n Rs = tf.SO3(wxyzs).as_matrix()\n covariances = onp.einsum(\n \"nij,njk,nlk->nil\", Rs, onp.eye(3)[None, :, :] * scales[:, None, :] ** 2, Rs\n )\n centers = splat_uint8[:, 0:12].copy().view(onp.float32)\n if center:\n centers -= onp.mean(centers, axis=0, keepdims=True)\n print(\n f\"Splat file with {num_gaussians=} loaded in {time.time() - start_time} seconds\"\n )\n return {\n \"centers\": centers,\n # Colors should have shape (N, 3).\n \"rgbs\": splat_uint8[:, 24:27] / 255.0,\n \"opacities\": splat_uint8[:, 27:28] / 255.0,\n # Covariances should have shape (N, 3, 3).\n \"covariances\": covariances,\n }\n\n\ndef load_ply_file(ply_file_path: Path, center: bool = False) -> SplatFile:\n \"\"\"Load Gaussians stored in a PLY file.\"\"\"\n start_time = time.time()\n\n SH_C0 = 0.28209479177387814\n\n plydata = PlyData.read(ply_file_path)\n v = plydata[\"vertex\"]\n positions = onp.stack([v[\"x\"], v[\"y\"], v[\"z\"]], axis=-1)\n scales = onp.exp(onp.stack([v[\"scale_0\"], v[\"scale_1\"], v[\"scale_2\"]], axis=-1))\n wxyzs = onp.stack([v[\"rot_0\"], v[\"rot_1\"], v[\"rot_2\"], v[\"rot_3\"]], axis=1)\n colors = 0.5 + SH_C0 * onp.stack([v[\"f_dc_0\"], v[\"f_dc_1\"], v[\"f_dc_2\"]], axis=1)\n opacities = 1.0 / (1.0 + onp.exp(-v[\"opacity\"][:, None]))\n\n Rs = tf.SO3(wxyzs).as_matrix()\n covariances = onp.einsum(\n \"nij,njk,nlk->nil\", Rs, onp.eye(3)[None, :, :] * scales[:, None, :] ** 2, Rs\n )\n if center:\n positions -= onp.mean(positions, axis=0, keepdims=True)\n\n num_gaussians = len(v)\n print(\n f\"PLY file with {num_gaussians=} loaded in {time.time() - start_time} seconds\"\n )\n return {\n \"centers\": positions,\n \"rgbs\": colors,\n \"opacities\": opacities,\n \"covariances\": covariances,\n }\n\n\ndef main(splat_paths: tuple[Path, ...]) -> None:\n server = viser.ViserServer()\n server.gui.configure_theme(dark_mode=True)\n gui_reset_up = server.gui.add_button(\n \"Reset up direction\",\n hint=\"Set the camera control 'up' direction to the current camera's 'up'.\",\n )\n\n @gui_reset_up.on_click\n def _(event: viser.GuiEvent) -> None:\n client = event.client\n assert client is not None\n client.camera.up_direction = tf.SO3(client.camera.wxyz) @ onp.array(\n [0.0, -1.0, 0.0]\n )\n\n for i, splat_path in enumerate(splat_paths):\n if splat_path.suffix == \".splat\":\n splat_data = load_splat_file(splat_path, center=True)\n elif splat_path.suffix == \".ply\":\n splat_data = load_ply_file(splat_path, center=True)\n else:\n raise SystemExit(\"Please provide a filepath to a .splat or .ply file.\")\n\n server.scene.add_transform_controls(f\"/{i}\")\n gs_handle = server.scene._add_gaussian_splats(\n f\"/{i}/gaussian_splats\",\n centers=splat_data[\"centers\"],\n rgbs=splat_data[\"rgbs\"],\n opacities=splat_data[\"opacities\"],\n covariances=splat_data[\"covariances\"],\n )\n\n remove_button = server.gui.add_button(f\"Remove splat object {i}\")\n\n @remove_button.on_click\n def _(_, gs_handle=gs_handle, remove_button=remove_button) -> None:\n gs_handle.remove()\n remove_button.remove()\n\n while True:\n time.sleep(10.0)\n\n\nif __name__ == \"__main__\":\n tyro.cli(main)\n" }, { "path": "viser/examples/quick_save.py", "content": "\"\"\"Record3D visualizer\n\nBatch process Record3D captures to generate recordings for multiple data folders.\n\"\"\"\n\nimport time\nimport sys\nimport os\nimport argparse\nfrom pathlib import Path\n\nimport numpy as onp\nfrom tqdm.auto import tqdm\n\nimport viser\nimport viser.extras\nimport viser.transforms as tf\nimport matplotlib.cm as cm # For colormap\n\ndef process_folder(\n data_path: Path,\n downsample_factor: int,\n max_frames: int,\n conf_threshold: float,\n foreground_conf_threshold: float,\n point_size: float,\n camera_frustum_scale: float,\n no_mask: bool,\n xyzw: bool,\n axes_scale: float,\n bg_downsample_factor: int,\n output_dir: Path,\n) -> None:\n print(f\"Processing folder: {data_path}\")\n\n server = viser.ViserServer()\n server.scene.set_up_direction('-z')\n\n loader = viser.extras.Record3dLoader_Customized(\n data_path,\n conf_threshold=conf_threshold,\n foreground_conf_threshold=foreground_conf_threshold,\n no_mask=no_mask,\n xyzw=xyzw,\n init_conf=True,\n )\n num_frames = min(max_frames, loader.num_frames())\n\n # Load frames\n server.scene.add_frame(\n \"/frames\",\n wxyz=tf.SO3.exp(onp.array([onp.pi / 2.0, 0.0, 0.0])).wxyz,\n position=(0, 0, 0),\n show_axes=False,\n )\n frame_nodes: list[viser.FrameHandle] = []\n bg_positions = []\n bg_colors = []\n for i in tqdm(range(num_frames)):\n frame = loader.get_frame(i)\n position, color, bg_position, bg_color = frame.get_point_cloud(downsample_factor, bg_downsample_factor)\n\n bg_positions.append(bg_position)\n bg_colors.append(bg_color)\n\n # Add base frame\n frame_nodes.append(server.scene.add_frame(f\"/frames/t{i}\", show_axes=False))\n\n # Place the point cloud in the frame\n server.scene.add_point_cloud(\n name=f\"/frames/t{i}/point_cloud\",\n points=position,\n colors=color,\n point_size=point_size,\n point_shape=\"rounded\",\n )\n\n # Compute color for frustum based on frame index\n norm_i = i / (num_frames - 1) if num_frames > 1 else 0 # Normalize index to [0, 1]\n color_rgba = cm.viridis(norm_i) # Get RGBA color from colormap\n color_rgb = color_rgba[:3] # Use RGB components\n\n # Place the frustum with the computed color\n fov = 2 * onp.arctan2(frame.rgb.shape[0] / 2, frame.K[0, 0])\n aspect = frame.rgb.shape[1] / frame.rgb.shape[0]\n server.scene.add_camera_frustum(\n f\"/frames/t{i}/frustum\",\n fov=fov,\n aspect=aspect,\n scale=camera_frustum_scale,\n image=frame.rgb[::downsample_factor, ::downsample_factor],\n wxyz=tf.SO3.from_matrix(frame.T_world_camera[:3, :3]).wxyz,\n position=frame.T_world_camera[:3, 3],\n color=color_rgb, # Set the color for the frustum\n )\n\n # Add axes\n server.scene.add_frame(\n f\"/frames/t{i}/frustum/axes\",\n axes_length=camera_frustum_scale * axes_scale * 10,\n axes_radius=camera_frustum_scale * axes_scale,\n )\n\n # Add background frame\n bg_positions = onp.concatenate(bg_positions, axis=0)\n bg_colors = onp.concatenate(bg_colors, axis=0)\n server.scene.add_point_cloud(\n name=f\"/frames/background\",\n points=bg_positions,\n colors=bg_colors,\n point_size=point_size,\n point_shape=\"rounded\",\n )\n\n # Automatically play through frames and record the scene\n rec = server._start_scene_recording()\n rec.set_loop_start()\n\n sleep_duration = 1.0 / loader.fps if loader.fps > 0 else 0.033 # Default to ~30 FPS\n\n for t in range(num_frames):\n # Update the scene to show frame t\n with server.atomic():\n for i, frame_node in enumerate(frame_nodes):\n frame_node.visible = (i == t)\n server.flush()\n rec.insert_sleep(sleep_duration)\n\n # Set all frames invisible\n with server.atomic():\n for frame_node in frame_nodes:\n frame_node.visible = False\n server.flush()\n\n # Finish recording\n bs = rec.end_and_serialize()\n\n # Save the recording to a file\n output_path = output_dir / f\"recording_{data_path.name}.viser\"\n # Ensure the output directory exists\n output_path.parent.mkdir(parents=True, exist_ok=True)\n output_path.write_bytes(bs)\n print(f\"Recording saved to {output_path.resolve()}\")\n\ndef main(\n data_paths: list[Path],\n output_dir: Path = Path(\"./viser_result\"),\n downsample_factor: int = 1,\n max_frames: int = 100,\n conf_threshold: float = 1.0,\n foreground_conf_threshold: float = 0.1,\n point_size: float = 0.001,\n camera_frustum_scale: float = 0.02,\n no_mask: bool = False,\n xyzw: bool = True,\n axes_scale: float = 0.25,\n bg_downsample_factor: int = 1,\n) -> None:\n # if data_path[0] has subfolders, process each subfolder\n if data_paths[0].is_dir():\n new_data_paths = sorted([subfolder for subfolder in data_paths[0].iterdir() if subfolder.is_dir()])\n \n if len(new_data_paths) > 0:\n data_paths = new_data_paths\n\n for data_path in data_paths:\n process_folder(\n data_path,\n downsample_factor,\n max_frames,\n conf_threshold,\n foreground_conf_threshold,\n point_size,\n camera_frustum_scale,\n no_mask,\n xyzw,\n axes_scale,\n bg_downsample_factor,\n output_dir,\n )\n\nif __name__ == \"__main__\":\n # Initialize parser\n parser = argparse.ArgumentParser(description=\"Process input arguments.\")\n\n # Define arguments\n parser.add_argument(\n \"--data\",\n type=Path,\n nargs=\"+\",\n required=True,\n help=\"Paths to the data folders (can specify multiple paths)\",\n )\n parser.add_argument(\n \"--output_dir\",\n type=Path,\n default=Path(\"./viser_result\"),\n help=\"Output directory for recordings\",\n )\n parser.add_argument(\n \"--conf_thre\",\n type=float,\n default=0.1,\n help=\"Confidence threshold, default is 0.1\",\n )\n parser.add_argument(\n \"--fg_conf_thre\",\n type=float,\n default=0.5,\n help=\"Foreground confidence threshold, default is 0.0\",\n )\n parser.add_argument(\n \"--point_size\",\n type=float,\n default=0.001,\n help=\"Point size, default is 0.001\",\n )\n parser.add_argument(\n \"--camera_size\",\n type=float,\n default=0.015,\n help=\"Camera frustum scale, default is 0.015\",\n )\n parser.add_argument(\n \"--no_mask\",\n action=\"store_true\",\n help=\"Don't use mask to filter out points\",\n )\n parser.add_argument(\n \"--wxyz\",\n action=\"store_true\",\n help=\"Use wxyz for SO3 representation\",\n )\n parser.add_argument(\n \"--axes_scale\",\n type=float,\n default=0.1,\n help=\"Scale of axes\",\n )\n parser.add_argument(\n \"--bg_downsample\",\n type=int,\n default=1,\n help=\"Background downsample factor\",\n )\n parser.add_argument(\n \"--downsample\",\n type=int,\n default=2,\n help=\"Downsample factor\",\n )\n parser.add_argument(\n \"--max_frames\",\n type=int,\n default=100,\n help=\"Maximum number of frames to process\",\n )\n\n # Parse arguments\n args = parser.parse_args()\n\n # Call the main function with the parsed arguments\n main(\n data_paths=args.data,\n output_dir=args.output_dir,\n conf_threshold=args.conf_thre,\n foreground_conf_threshold=args.fg_conf_thre,\n point_size=args.point_size,\n camera_frustum_scale=args.camera_size,\n no_mask=args.no_mask,\n xyzw=not args.wxyz,\n axes_scale=args.axes_scale,\n bg_downsample_factor=args.bg_downsample,\n downsample_factor=args.downsample,\n max_frames=args.max_frames,\n )\n" }, { "path": "viser/pyproject.toml", "content": "[build-system]\nrequires = [\"setuptools>=61.0\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"viser\"\nversion = \"0.2.7\"\ndescription = \"3D visualization + Python\"\nreadme = \"README.md\"\nlicense = { text=\"MIT\" }\nrequires-python = \">=3.8\"\nclassifiers = [\n \"Programming Language :: Python :: 3\",\n \"Programming Language :: Python :: 3.8\",\n \"Programming Language :: Python :: 3.9\",\n \"Programming Language :: Python :: 3.10\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"License :: OSI Approved :: MIT License\",\n \"Operating System :: OS Independent\"\n]\ndependencies = [\n \"websockets>=10.4\",\n \"numpy>=1.0.0\",\n \"msgspec>=0.18.6\",\n \"imageio>=2.0.0\",\n \"pyliblzfse>=0.4.1; platform_system!='Windows'\",\n \"scikit-image>=0.18.0\",\n \"scipy>=1.7.3\",\n \"tqdm>=4.0.0\",\n \"tyro>=0.2.0\",\n \"rich>=13.3.3\",\n \"trimesh>=3.21.7\",\n \"nodeenv>=1.8.0\",\n \"psutil>=5.9.5\",\n \"yourdfpy>=0.0.53\",\n \"plyfile>=1.0.2\"\n]\n\n[project.optional-dependencies]\ndev = [\n \"pyright>=1.1.308\",\n \"ruff==0.6.2\",\n \"pre-commit==3.3.2\",\n]\nexamples = [\n \"torch>=1.13.1\",\n \"matplotlib>=3.7.1\",\n \"plotly>=5.21.0\",\n \"robot_descriptions>=1.10.0\",\n \"gdown>=4.6.6\",\n \"plyfile\",\n]\n\n[project.urls]\n\"GitHub\" = \"https://github.com/nerfstudio-project/viser\"\n\n# <>\n# Important: in the ./.github/workflows/publish.yml action, we have sed\n# commands that assume the `viser = ...` line below directly follows\n# `[tool.setuptools.package-data]`. We use this to remove the client source\n# from PyPI builds.\n#\n# We should make sure that any modifications to the package-data list remain\n# compatible with the sed commands!\n#\n# We keep the client source in by default to support things like pip\n# installation via the Git URL, because build artifacts aren't\n# version-controlled.\n[tool.setuptools.package-data]\nviser = [\"py.typed\", \"*.pyi\", \"_icons/tabler-icons.tar\", \"client/**/*\", \"client/**/.*\"]\n# \n\n[tool.setuptools.exclude-package-data]\n# We exclude node_modules to prevent long build times for wheels when\n# installing from source, eg via `pip install .`.\n#\n# https://github.com/nerfstudio-project/viser/issues/271\nviser = [\"**/node_modules/**\"]\n\n[project.scripts]\nviser-dev-checks = \"viser.scripts.dev_checks:entrypoint\"\n\n[tool.pyright]\nexclude = [\"./docs/**/*\", \"./examples/assets/**/*\", \"./src/viser/client/.nodeenv\", \"./build\"]\n\n[tool.ruff]\nlint.select = [\n \"E\", # pycodestyle errors.\n \"F\", # Pyflakes rules.\n \"PLC\", # Pylint convention warnings.\n \"PLE\", # Pylint errors.\n \"PLR\", # Pylint refactor recommendations.\n \"PLW\", # Pylint warnings.\n \"I\", # Import sorting.\n]\nlint.ignore = [\n \"E741\", # Ambiguous variable name. (l, O, or I)\n \"E501\", # Line too long.\n \"E721\", # Do not compare types, use `isinstance()`.\n \"F722\", # Forward annotation false positive from jaxtyping. Should be caught by pyright.\n \"F821\", # Forward annotation false positive from jaxtyping. Should be caught by pyright.\n \"PLR2004\", # Magic value used in comparison.\n \"PLR0915\", # Too many statements.\n \"PLR0913\", # Too many arguments.\n \"PLC0414\", # Import alias does not rename variable. (this is used for exporting names)\n \"PLC1901\", # Use falsey strings.\n \"PLR5501\", # Use `elif` instead of `else if`.\n \"PLR0911\", # Too many return statements.\n \"PLR0912\", # Too many branches.\n \"PLW0603\", # Globa statement updates are discouraged.\n \"PLW2901\", # For loop variable overwritten.\n \"PLW0642\", # Reassigned self in instance method.\n]\nexclude = [ \".nodeenv\" ]\n" }, { "path": "viser/src/viser/__init__.py", "content": "from ._gui_api import GuiApi as GuiApi\nfrom ._gui_handles import GuiButtonGroupHandle as GuiButtonGroupHandle\nfrom ._gui_handles import GuiButtonHandle as GuiButtonHandle\nfrom ._gui_handles import GuiDropdownHandle as GuiDropdownHandle\nfrom ._gui_handles import GuiEvent as GuiEvent\nfrom ._gui_handles import GuiFolderHandle as GuiFolderHandle\nfrom ._gui_handles import GuiInputHandle as GuiInputHandle\nfrom ._gui_handles import GuiMarkdownHandle as GuiMarkdownHandle\nfrom ._gui_handles import GuiPlotlyHandle as GuiPlotlyHandle\nfrom ._gui_handles import GuiTabGroupHandle as GuiTabGroupHandle\nfrom ._gui_handles import GuiTabHandle as GuiTabHandle\nfrom ._icons_enum import Icon as Icon\nfrom ._icons_enum import IconName as IconName\nfrom ._notification_handle import NotificationHandle as NotificationHandle\nfrom ._scene_api import SceneApi as SceneApi\nfrom ._scene_handles import BatchedAxesHandle as BatchedAxesHandle\nfrom ._scene_handles import CameraFrustumHandle as CameraFrustumHandle\nfrom ._scene_handles import FrameHandle as FrameHandle\nfrom ._scene_handles import GaussianSplatHandle as GaussianSplatHandle\nfrom ._scene_handles import GlbHandle as GlbHandle\nfrom ._scene_handles import Gui3dContainerHandle as Gui3dContainerHandle\nfrom ._scene_handles import ImageHandle as ImageHandle\nfrom ._scene_handles import LabelHandle as LabelHandle\nfrom ._scene_handles import MeshHandle as MeshHandle\nfrom ._scene_handles import MeshSkinnedBoneHandle as MeshSkinnedBoneHandle\nfrom ._scene_handles import MeshSkinnedHandle as MeshSkinnedHandle\nfrom ._scene_handles import PointCloudHandle as PointCloudHandle\nfrom ._scene_handles import SceneNodeHandle as SceneNodeHandle\nfrom ._scene_handles import SceneNodePointerEvent as SceneNodePointerEvent\nfrom ._scene_handles import ScenePointerEvent as ScenePointerEvent\nfrom ._scene_handles import TransformControlsHandle as TransformControlsHandle\nfrom ._viser import CameraHandle as CameraHandle\nfrom ._viser import ClientHandle as ClientHandle\nfrom ._viser import ViserServer as ViserServer\n" }, { "path": "viser/src/viser/_client_autobuild.py", "content": "import os\nimport subprocess\nimport sys\nfrom pathlib import Path\n\nimport rich\n\nclient_dir = Path(__file__).absolute().parent / \"client\"\nbuild_dir = client_dir / \"build\"\n\n\ndef _check_viser_yarn_running() -> bool:\n \"\"\"Returns True if the viewer client has been launched via `yarn start`.\"\"\"\n import psutil\n\n for process in psutil.process_iter():\n try:\n if Path(process.cwd()).as_posix().endswith(\"viser/client\") and any(\n [part.endswith(\"yarn\") for part in process.cmdline()]\n + [part.endswith(\"yarn.js\") for part in process.cmdline()]\n ):\n return True\n except (psutil.AccessDenied, psutil.ZombieProcess):\n pass\n return False\n\n\ndef ensure_client_is_built() -> None:\n \"\"\"Ensure that the client is built or already running.\"\"\"\n\n if not (client_dir / \"src\").exists():\n # Can't build client.\n assert (build_dir / \"index.html\").exists(), (\n \"Something went wrong! At least one of the client source or build\"\n \" directories should be present.\"\n )\n return\n\n # Do we need to re-trigger a build?\n build = False\n if _check_viser_yarn_running():\n # Don't run `yarn build` if `yarn start` is already running.\n rich.print(\n \"[bold](viser)[/bold] The Viser viewer looks like it has been launched via\"\n \" `yarn start`. Skipping build check...\"\n )\n build = False\n elif not (build_dir / \"index.html\").exists():\n rich.print(\"[bold](viser)[/bold] No client build found. Building now...\")\n build = True\n elif (\n # We should be at least 10 seconds newer than the last build.\n # This buffer is important when we install from pip, and the src/ +\n # build/ directories have very similar timestamps.\n _modified_time_recursive(client_dir / \"src\")\n > _modified_time_recursive(build_dir) + 10.0\n ):\n rich.print(\n \"[bold](viser)[/bold] Client build looks out of date. Building now...\"\n )\n build = True\n\n # Install nodejs and build if necessary. We assume bash is installed.\n if build:\n node_bin_dir = _install_sandboxed_node()\n npx_path = node_bin_dir / \"npx\"\n\n subprocess_env = os.environ.copy()\n subprocess_env[\"NODE_VIRTUAL_ENV\"] = str(node_bin_dir.parent)\n subprocess_env[\"PATH\"] = (\n str(node_bin_dir)\n + (\";\" if sys.platform == \"win32\" else \":\")\n + subprocess_env[\"PATH\"]\n )\n subprocess.run(\n args=f\"{npx_path} --yes yarn install\",\n env=subprocess_env,\n cwd=client_dir,\n shell=True,\n check=False,\n )\n subprocess.run(\n args=f\"{npx_path} --yes yarn run build\",\n env=subprocess_env,\n cwd=client_dir,\n shell=True,\n check=False,\n )\n\n\ndef _install_sandboxed_node() -> Path:\n \"\"\"Install a sandboxed copy of nodejs using nodeenv, and return a path to the\n environment's bin directory (`.nodeenv/bin` or `.nodeenv/Scripts`).\n\n On Windows, the `.nodeenv/bin` does not exist. Instead, executables are\n installed to `.nodeenv/Scripts`.\"\"\"\n\n def get_node_bin_dir() -> Path:\n env_dir = client_dir / \".nodeenv\"\n node_bin_dir = env_dir / \"bin\"\n if not node_bin_dir.exists():\n node_bin_dir = env_dir / \"Scripts\"\n return node_bin_dir\n\n node_bin_dir = get_node_bin_dir()\n if (node_bin_dir / \"npx\").exists():\n rich.print(\"[bold](viser)[/bold] nodejs is set up!\")\n return node_bin_dir\n\n env_dir = client_dir / \".nodeenv\"\n subprocess.run(\n [sys.executable, \"-m\", \"nodeenv\", \"--node=20.4.0\", env_dir], check=False\n )\n\n node_bin_dir = get_node_bin_dir()\n assert (node_bin_dir / \"npx\").exists()\n return node_bin_dir\n\n\ndef _modified_time_recursive(dir: Path) -> float:\n \"\"\"Recursively get the last time a file was modified in a directory.\"\"\"\n return max([f.stat().st_mtime for f in dir.glob(\"**/*\")])\n" }, { "path": "viser/src/viser/_gui_api.py", "content": "from __future__ import annotations\n\nimport builtins\nimport colorsys\nimport dataclasses\nimport functools\nimport threading\nimport time\nfrom concurrent.futures import ThreadPoolExecutor\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Any, Sequence, Tuple, TypeVar, cast, overload\n\nimport numpy as onp\nfrom typing_extensions import (\n Literal,\n LiteralString,\n TypeAlias,\n TypedDict,\n get_type_hints,\n)\n\nfrom viser import theme\n\nfrom . import _messages\nfrom ._gui_handles import (\n GuiButtonGroupHandle,\n GuiButtonHandle,\n GuiContainerProtocol,\n GuiDropdownHandle,\n GuiEvent,\n GuiFolderHandle,\n GuiInputHandle,\n GuiMarkdownHandle,\n GuiModalHandle,\n GuiPlotlyHandle,\n GuiProgressBarHandle,\n GuiTabGroupHandle,\n GuiUploadButtonHandle,\n SupportsRemoveProtocol,\n UploadedFile,\n _GuiHandleState,\n _GuiInputHandle,\n _make_unique_id,\n)\nfrom ._icons import svg_from_icon\nfrom ._icons_enum import IconName\nfrom ._messages import FileTransferPartAck\nfrom ._scene_api import cast_vector\n\nif TYPE_CHECKING:\n import plotly.graph_objects as go\n\n from ._viser import ClientHandle, ViserServer\n from .infra import ClientId\n\n\nIntOrFloat = TypeVar(\"IntOrFloat\", int, float)\nTString = TypeVar(\"TString\", bound=str)\nTLiteralString = TypeVar(\"TLiteralString\", bound=LiteralString)\nT = TypeVar(\"T\")\nLengthTenStrTuple: TypeAlias = Tuple[str, str, str, str, str, str, str, str, str, str]\nColor: TypeAlias = Literal[\n \"dark\",\n \"gray\",\n \"red\",\n \"pink\",\n \"grape\",\n \"violet\",\n \"indigo\",\n \"blue\",\n \"cyan\",\n \"green\",\n \"lime\",\n \"yellow\",\n \"orange\",\n \"teal\",\n]\n\n\ndef _hex_from_hls(h: float, l: float, s: float) -> str:\n \"\"\"Converts HLS values in [0.0, 1.0] to a hex-formatted string, eg 0xffffff.\"\"\"\n return \"#\" + \"\".join(\n [\n int(min(255, max(0, channel * 255.0)) + 0.5).to_bytes(1, \"little\").hex()\n for channel in colorsys.hls_to_rgb(h, l, s)\n ]\n )\n\n\ndef _compute_step(x: float | None) -> float: # type: ignore\n \"\"\"For number inputs: compute an increment size from some number.\n\n Example inputs/outputs:\n 100 => 1\n 12 => 1\n 12.1 => 0.1\n 12.02 => 0.01\n 0.004 => 0.001\n \"\"\"\n return 1 if x is None else 10 ** (-_compute_precision_digits(x))\n\n\ndef _compute_precision_digits(x: float) -> int:\n \"\"\"For number inputs: compute digits of precision from some number.\n\n Example inputs/outputs:\n 100 => 0\n 12 => 0\n 12.1 => 1\n 10.2 => 1\n 0.007 => 3\n \"\"\"\n digits = 0\n while x != round(x, ndigits=digits) and digits < 7:\n digits += 1\n return digits\n\n\n@dataclasses.dataclass\nclass _RootGuiContainer:\n _children: dict[str, SupportsRemoveProtocol]\n\n\n_global_order_counter = 0\n\n\ndef _apply_default_order(order: float | None) -> float:\n \"\"\"Apply default ordering logic for GUI elements.\n\n If `order` is set to a float, this function is a no-op and returns it back.\n Otherwise, we increment and return the value of a global counter.\n \"\"\"\n if order is not None:\n return order\n\n global _global_order_counter\n _global_order_counter += 1\n return _global_order_counter\n\n\n@functools.lru_cache(maxsize=None)\ndef get_type_hints_cached(cls: type[Any]) -> dict[str, Any]:\n return get_type_hints(cls) # type: ignore\n\n\nclass _FileUploadState(TypedDict):\n filename: str\n mime_type: str\n part_count: int\n parts: dict[int, bytes]\n total_bytes: int\n transferred_bytes: int\n lock: threading.Lock\n\n\nclass GuiApi:\n \"\"\"Interface for working with the 2D GUI in viser.\n\n Used by both our global server object, for sharing the same GUI elements\n with all clients, and by individual client handles.\"\"\"\n\n _target_container_from_thread_id: dict[int, str] = {}\n \"\"\"ID of container to put GUI elements into.\"\"\"\n\n def __init__(\n self,\n owner: ViserServer | ClientHandle, # Who do I belong to?\n thread_executor: ThreadPoolExecutor,\n ) -> None:\n from ._viser import ViserServer\n\n self._owner = owner\n \"\"\"Entity that owns this API.\"\"\"\n self._thread_executor = thread_executor\n\n self._websock_interface = (\n owner._websock_server\n if isinstance(owner, ViserServer)\n else owner._websock_connection\n )\n \"\"\"Interface for sending and listening to messages.\"\"\"\n\n self._gui_input_handle_from_id: dict[str, _GuiInputHandle[Any]] = {}\n self._container_handle_from_id: dict[str, GuiContainerProtocol] = {\n \"root\": _RootGuiContainer({})\n }\n self._current_file_upload_states: dict[str, _FileUploadState] = {}\n\n # Set to True when plotly.min.js has been sent to client.\n self._setup_plotly_js: bool = False\n\n self._websock_interface.register_handler(\n _messages.GuiUpdateMessage, self._handle_gui_updates\n )\n self._websock_interface.register_handler(\n _messages.FileTransferStart, self._handle_file_transfer_start\n )\n self._websock_interface.register_handler(\n _messages.FileTransferPart,\n self._handle_file_transfer_part,\n )\n\n def _handle_gui_updates(\n self, client_id: ClientId, message: _messages.GuiUpdateMessage\n ) -> None:\n \"\"\"Callback for handling GUI messages.\"\"\"\n handle = self._gui_input_handle_from_id.get(message.id, None)\n if handle is None:\n return\n handle_state = handle._impl\n\n has_changed = False\n updates_cast = {}\n for prop_name, prop_value in message.updates.items():\n assert hasattr(handle_state, prop_name)\n current_value = getattr(handle_state, prop_name)\n\n # Do some type casting. This is brittle, but necessary (1) when we\n # expect floats but the Javascript side gives us integers or (2)\n # when we expect tuples but the Javascript side gives us lists.\n if prop_name == \"value\":\n if isinstance(handle_state.value, tuple):\n # We currently assume all tuple types have length >0, and\n # contents are all the same type.\n assert len(handle_state.value) > 0\n typ = type(handle_state.value[0])\n assert all([type(x) == typ for x in handle_state.value])\n prop_value = tuple([typ(new) for new in prop_value])\n else:\n prop_value = type(handle_state.value)(prop_value)\n\n # Update handle property.\n if current_value != prop_value:\n has_changed = True\n setattr(handle_state, prop_name, prop_value)\n\n # Save value, which might have been cast.\n updates_cast[prop_name] = prop_value\n\n # Only call update when value has actually changed.\n if not handle_state.is_button and not has_changed:\n return\n\n # GUI element has been updated!\n handle_state.update_timestamp = time.time()\n for cb in handle_state.update_cb:\n from ._viser import ClientHandle, ViserServer\n\n # Get the handle of the client that triggered this event.\n if isinstance(self._owner, ClientHandle):\n client = self._owner\n elif isinstance(self._owner, ViserServer):\n client = self._owner.get_clients()[client_id]\n else:\n assert False\n\n cb(GuiEvent(client, client_id, handle))\n\n if handle_state.sync_cb is not None:\n handle_state.sync_cb(client_id, updates_cast)\n\n def _handle_file_transfer_start(\n self, client_id: ClientId, message: _messages.FileTransferStart\n ) -> None:\n if message.source_component_id not in self._gui_input_handle_from_id:\n return\n self._current_file_upload_states[message.transfer_uuid] = {\n \"filename\": message.filename,\n \"mime_type\": message.mime_type,\n \"part_count\": message.part_count,\n \"parts\": {},\n \"total_bytes\": message.size_bytes,\n \"transferred_bytes\": 0,\n \"lock\": threading.Lock(),\n }\n\n def _handle_file_transfer_part(\n self, client_id: ClientId, message: _messages.FileTransferPart\n ) -> None:\n if message.transfer_uuid not in self._current_file_upload_states:\n return\n assert message.source_component_id in self._gui_input_handle_from_id\n\n state = self._current_file_upload_states[message.transfer_uuid]\n state[\"parts\"][message.part] = message.content\n total_bytes = state[\"total_bytes\"]\n\n with state[\"lock\"]:\n state[\"transferred_bytes\"] += len(message.content)\n\n # Send ack to the server.\n self._websock_interface.queue_message(\n FileTransferPartAck(\n source_component_id=message.source_component_id,\n transfer_uuid=message.transfer_uuid,\n transferred_bytes=state[\"transferred_bytes\"],\n total_bytes=total_bytes,\n )\n )\n\n if state[\"transferred_bytes\"] < total_bytes:\n return\n\n # Finish the upload.\n assert state[\"transferred_bytes\"] == total_bytes\n state = self._current_file_upload_states.pop(message.transfer_uuid)\n\n handle = self._gui_input_handle_from_id.get(message.source_component_id, None)\n if handle is None:\n return\n\n handle_state = handle._impl\n\n value = UploadedFile(\n name=state[\"filename\"],\n content=b\"\".join(state[\"parts\"][i] for i in range(state[\"part_count\"])),\n )\n\n # Update state.\n with self._owner.atomic():\n handle_state.value = value\n handle_state.update_timestamp = time.time()\n\n # Trigger callbacks.\n for cb in handle_state.update_cb:\n from ._viser import ClientHandle, ViserServer\n\n # Get the handle of the client that triggered this event.\n if isinstance(self._owner, ClientHandle):\n client = self._owner\n elif isinstance(self._owner, ViserServer):\n client = self._owner.get_clients()[client_id]\n else:\n assert False\n\n cb(GuiEvent(client, client_id, handle))\n\n def _get_container_id(self) -> str:\n \"\"\"Get container ID associated with the current thread.\"\"\"\n return self._target_container_from_thread_id.get(threading.get_ident(), \"root\")\n\n def _set_container_id(self, container_id: str) -> None:\n \"\"\"Set container ID associated with the current thread.\"\"\"\n self._target_container_from_thread_id[threading.get_ident()] = container_id\n\n def reset(self) -> None:\n \"\"\"Reset the GUI.\"\"\"\n self._websock_interface.queue_message(_messages.ResetGuiMessage())\n\n def set_panel_label(self, label: str | None) -> None:\n \"\"\"Set the main label that appears in the GUI panel.\n\n Args:\n label: The new label.\n \"\"\"\n self._websock_interface.queue_message(_messages.SetGuiPanelLabelMessage(label))\n\n def configure_theme(\n self,\n *,\n titlebar_content: theme.TitlebarConfig | None = None,\n control_layout: Literal[\"floating\", \"collapsible\", \"fixed\"] = \"floating\",\n control_width: Literal[\"small\", \"medium\", \"large\"] = \"medium\",\n dark_mode: bool = False,\n show_logo: bool = True,\n show_share_button: bool = True,\n brand_color: tuple[int, int, int] | None = None,\n ) -> None:\n \"\"\"Configures the visual appearance of the viser front-end.\n\n Args:\n titlebar_content: Optional configuration for the title bar.\n control_layout: The layout of control elements, options are \"floating\",\n \"collapsible\", or \"fixed\".\n control_width: The width of control elements, options are \"small\",\n \"medium\", or \"large\".\n dark_mode: A boolean indicating if dark mode should be enabled.\n show_logo: A boolean indicating if the logo should be displayed.\n show_share_button: A boolean indicating if the share button should be displayed.\n brand_color: An optional tuple of integers (RGB) representing the brand color.\n \"\"\"\n\n colors_cast: LengthTenStrTuple | None = None\n\n if brand_color is not None:\n assert len(brand_color) in (3, 10)\n if len(brand_color) == 3:\n assert all(\n map(lambda val: isinstance(val, int), brand_color)\n ), \"All channels should be integers.\"\n\n # RGB => HLS.\n h, l, s = colorsys.rgb_to_hls(\n brand_color[0] / 255.0,\n brand_color[1] / 255.0,\n brand_color[2] / 255.0,\n )\n\n # Automatically generate a 10-color palette.\n min_l = max(l - 0.08, 0.0)\n max_l = min(0.8 + 0.5, 0.9)\n l = max(min_l, min(max_l, l))\n\n primary_index = 8\n ls = tuple(\n onp.interp(\n x=onp.arange(10),\n xp=onp.array([0, primary_index, 9]),\n fp=onp.array([max_l, l, min_l]),\n )\n )\n colors_cast = cast(\n LengthTenStrTuple,\n tuple(_hex_from_hls(h, ls[i], s) for i in range(10)),\n )\n\n assert colors_cast is None or all(\n [isinstance(val, str) and val.startswith(\"#\") for val in colors_cast]\n ), \"All string colors should be in hexadecimal + prefixed with #, eg #ffffff.\"\n\n self._websock_interface.queue_message(\n _messages.ThemeConfigurationMessage(\n titlebar_content=titlebar_content,\n control_layout=control_layout,\n control_width=control_width,\n dark_mode=dark_mode,\n show_logo=show_logo,\n show_share_button=show_share_button,\n colors=colors_cast,\n ),\n )\n\n def add_folder(\n self,\n label: str,\n order: float | None = None,\n expand_by_default: bool = True,\n visible: bool = True,\n ) -> GuiFolderHandle:\n \"\"\"Add a folder, and return a handle that can be used to populate it.\n\n Args:\n label: Label to display on the folder.\n order: Optional ordering, smallest values will be displayed first.\n expand_by_default: Open the folder by default. Set to False to collapse it by\n default.\n visible: Whether the component is visible.\n\n Returns:\n A handle that can be used as a context to populate the folder.\n \"\"\"\n folder_container_id = _make_unique_id()\n order = _apply_default_order(order)\n self._websock_interface.queue_message(\n _messages.GuiAddFolderMessage(\n order=order,\n id=folder_container_id,\n label=label,\n container_id=self._get_container_id(),\n expand_by_default=expand_by_default,\n visible=visible,\n )\n )\n return GuiFolderHandle(\n _gui_api=self,\n _id=folder_container_id,\n _parent_container_id=self._get_container_id(),\n _order=order,\n )\n\n def add_modal(\n self,\n title: str,\n order: float | None = None,\n ) -> GuiModalHandle:\n \"\"\"Show a modal window, which can be useful for popups and messages, then return\n a handle that can be used to populate it.\n\n Args:\n title: Title to display on the modal.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used as a context to populate the modal.\n \"\"\"\n modal_container_id = _make_unique_id()\n order = _apply_default_order(order)\n self._websock_interface.queue_message(\n _messages.GuiModalMessage(\n order=order,\n id=modal_container_id,\n title=title,\n )\n )\n return GuiModalHandle(\n _gui_api=self,\n _id=modal_container_id,\n )\n\n def add_tab_group(\n self,\n order: float | None = None,\n visible: bool = True,\n ) -> GuiTabGroupHandle:\n \"\"\"Add a tab group.\n\n Args:\n order: Optional ordering, smallest values will be displayed first.\n visible: Whether the component is visible.\n\n Returns:\n A handle that can be used as a context to populate the tab group.\n \"\"\"\n tab_group_id = _make_unique_id()\n order = _apply_default_order(order)\n\n self._websock_interface.queue_message(\n _messages.GuiAddTabGroupMessage(\n order=order,\n id=tab_group_id,\n container_id=self._get_container_id(),\n tab_labels=(),\n visible=visible,\n tab_icons_html=(),\n tab_container_ids=(),\n )\n )\n return GuiTabGroupHandle(\n _tab_group_id=tab_group_id,\n _labels=[],\n _icons_html=[],\n _tabs=[],\n _gui_api=self,\n _parent_container_id=self._get_container_id(),\n _order=order,\n )\n\n def add_markdown(\n self,\n content: str,\n image_root: Path | None = None,\n order: float | None = None,\n visible: bool = True,\n ) -> GuiMarkdownHandle:\n \"\"\"Add markdown to the GUI.\n\n Args:\n content: Markdown content to display.\n image_root: Optional root directory to resolve relative image paths.\n order: Optional ordering, smallest values will be displayed first.\n visible: Whether the component is visible.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n handle = GuiMarkdownHandle(\n _gui_api=self,\n _id=_make_unique_id(),\n _visible=visible,\n _parent_container_id=self._get_container_id(),\n _order=_apply_default_order(order),\n _image_root=image_root,\n _content=None,\n )\n self._websock_interface.queue_message(\n _messages.GuiAddMarkdownMessage(\n order=handle._order,\n id=handle._id,\n markdown=\"\",\n container_id=handle._parent_container_id,\n visible=visible,\n )\n )\n\n # Logic for processing markdown, handling images, etc is all in the\n # `.content` setter, which should send a GuiUpdateMessage.\n handle.content = content\n return handle\n\n def add_plotly(\n self,\n figure: go.Figure,\n aspect: float = 1.0,\n order: float | None = None,\n visible: bool = True,\n ) -> GuiPlotlyHandle:\n \"\"\"Add a Plotly figure to the GUI. Requires the `plotly` package to be\n installed.\n\n Args:\n figure: Plotly figure to display.\n aspect: Aspect ratio of the plot in the control panel (width/height).\n order: Optional ordering, smallest values will be displayed first.\n visible: Whether the component is visible.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n handle = GuiPlotlyHandle(\n _gui_api=self,\n _id=_make_unique_id(),\n _visible=visible,\n _parent_container_id=self._get_container_id(),\n _order=_apply_default_order(order),\n _figure=None,\n _aspect=None,\n )\n\n # If plotly.min.js hasn't been sent to the client yet, the client won't be able\n # to render the plot. Send this large file now! (~3MB)\n if not self._setup_plotly_js:\n # Check if plotly is installed.\n try:\n import plotly\n except ImportError:\n raise ImportError(\n \"You must have the `plotly` package installed to use the Plotly GUI element.\"\n )\n\n # Check that plotly.min.js exists.\n plotly_path = (\n Path(plotly.__file__).parent / \"package_data\" / \"plotly.min.js\"\n )\n assert (\n plotly_path.exists()\n ), f\"Could not find plotly.min.js at {plotly_path}.\"\n\n # Send it over!\n plotly_js = plotly_path.read_text(encoding=\"utf-8\")\n self._websock_interface.queue_message(\n _messages.RunJavascriptMessage(source=plotly_js)\n )\n\n # Update the flag so we don't send it again.\n self._setup_plotly_js = True\n\n # After plotly.min.js has been sent, we can send the plotly figure.\n # Empty string for `plotly_json_str` is a signal to the client to render nothing.\n self._websock_interface.queue_message(\n _messages.GuiAddPlotlyMessage(\n order=handle._order,\n id=handle._id,\n plotly_json_str=\"\",\n aspect=1.0,\n container_id=handle._parent_container_id,\n visible=visible,\n )\n )\n\n # Set the plotly handle properties.\n handle.figure = figure\n handle.aspect = aspect\n\n return handle\n\n def add_button(\n self,\n label: str,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n color: Color | None = None,\n icon: IconName | None = None,\n order: float | None = None,\n ) -> GuiButtonHandle:\n \"\"\"Add a button to the GUI. The value of this input is set to `True` every time\n it is clicked; to detect clicks, we can manually set it back to `False`.\n\n Args:\n label: Label to display on the button.\n visible: Whether the button is visible.\n disabled: Whether the button is disabled.\n hint: Optional hint to display on hover.\n color: Optional color to use for the button.\n icon: Optional icon to display on the button.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n\n # Re-wrap the GUI handle with a button interface.\n id = _make_unique_id()\n order = _apply_default_order(order)\n return GuiButtonHandle(\n self._create_gui_input(\n value=False,\n message=_messages.GuiAddButtonMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=False,\n color=color,\n icon_html=None if icon is None else svg_from_icon(icon),\n disabled=disabled,\n visible=visible,\n ),\n is_button=True,\n )._impl\n )\n\n def add_upload_button(\n self,\n label: str,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n color: Color | None = None,\n icon: IconName | None = None,\n mime_type: str = \"*/*\",\n order: float | None = None,\n ) -> GuiUploadButtonHandle:\n \"\"\"Add a button to the GUI. The value of this input is set to `True` every time\n it is clicked; to detect clicks, we can manually set it back to `False`.\n\n Args:\n label: Label to display on the button.\n visible: Whether the button is visible.\n disabled: Whether the button is disabled.\n hint: Optional hint to display on hover.\n color: Optional color to use for the button.\n icon: Optional icon to display on the button.\n mime_type: Optional MIME type to filter the files that can be uploaded.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n\n # Re-wrap the GUI handle with a button interface.\n id = _make_unique_id()\n order = _apply_default_order(order)\n return GuiUploadButtonHandle(\n self._create_gui_input(\n value=UploadedFile(\"\", b\"\"),\n message=_messages.GuiAddUploadButtonMessage(\n value=None,\n disabled=disabled,\n visible=visible,\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n color=color,\n mime_type=mime_type,\n icon_html=None if icon is None else svg_from_icon(icon),\n ),\n is_button=True,\n )._impl\n )\n\n # The TLiteralString overload tells pyright to resolve the value type to a Literal\n # whenever possible.\n #\n # TString is helpful when the input types are generic (could be str, could be\n # Literal).\n @overload\n def add_button_group(\n self,\n label: str,\n options: Sequence[TLiteralString],\n visible: bool = True,\n disabled: bool = False,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiButtonGroupHandle[TLiteralString]: ...\n\n @overload\n def add_button_group(\n self,\n label: str,\n options: Sequence[TString],\n visible: bool = True,\n disabled: bool = False,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiButtonGroupHandle[TString]: ...\n\n def add_button_group(\n self,\n label: str,\n options: Sequence[TLiteralString] | Sequence[TString],\n visible: bool = True,\n disabled: bool = False,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiButtonGroupHandle[Any]: # Return types are specified in overloads.\n \"\"\"Add a button group to the GUI.\n\n Args:\n label: Label to display on the button group.\n options: Sequence of options to display as buttons.\n visible: Whether the button group is visible.\n disabled: Whether the button group is disabled.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = options[0]\n id = _make_unique_id()\n order = _apply_default_order(order)\n return GuiButtonGroupHandle(\n self._create_gui_input(\n value,\n message=_messages.GuiAddButtonGroupMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n options=tuple(options),\n disabled=disabled,\n visible=visible,\n ),\n )._impl,\n )\n\n def add_checkbox(\n self,\n label: str,\n initial_value: bool,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[bool]:\n \"\"\"Add a checkbox to the GUI.\n\n Args:\n label: Label to display on the checkbox.\n initial_value: Initial value of the checkbox.\n disabled: Whether the checkbox is disabled.\n visible: Whether the checkbox is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = initial_value\n assert isinstance(value, bool)\n id = _make_unique_id()\n order = _apply_default_order(order)\n return self._create_gui_input(\n value,\n message=_messages.GuiAddCheckboxMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n disabled=disabled,\n visible=visible,\n ),\n )\n\n def add_text(\n self,\n label: str,\n initial_value: str,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[str]:\n \"\"\"Add a text input to the GUI.\n\n Args:\n label: Label to display on the text input.\n initial_value: Initial value of the text input.\n disabled: Whether the text input is disabled.\n visible: Whether the text input is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = initial_value\n assert isinstance(value, str)\n id = _make_unique_id()\n order = _apply_default_order(order)\n return self._create_gui_input(\n value,\n message=_messages.GuiAddTextMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n disabled=disabled,\n visible=visible,\n ),\n )\n\n def add_number(\n self,\n label: str,\n initial_value: IntOrFloat,\n min: IntOrFloat | None = None,\n max: IntOrFloat | None = None,\n step: IntOrFloat | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[IntOrFloat]:\n \"\"\"Add a number input to the GUI, with user-specifiable bound and precision parameters.\n\n Args:\n label: Label to display on the number input.\n initial_value: Initial value of the number input.\n min: Optional minimum value of the number input.\n max: Optional maximum value of the number input.\n step: Optional step size of the number input. Computed automatically if not\n specified.\n disabled: Whether the number input is disabled.\n visible: Whether the number input is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = initial_value\n\n assert isinstance(value, (int, float))\n\n if step is None:\n # It's ok that `step` is always a float, even if the value is an integer,\n # because things all become `number` types after serialization.\n step = float( # type: ignore\n onp.min(\n [\n _compute_step(value),\n _compute_step(min),\n _compute_step(max),\n ]\n )\n )\n\n assert step is not None\n\n id = _make_unique_id()\n order = _apply_default_order(order)\n return self._create_gui_input(\n value,\n message=_messages.GuiAddNumberMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n min=min,\n max=max,\n precision=_compute_precision_digits(step),\n step=step,\n disabled=disabled,\n visible=visible,\n ),\n is_button=False,\n )\n\n def add_vector2(\n self,\n label: str,\n initial_value: tuple[float, float] | onp.ndarray,\n min: tuple[float, float] | onp.ndarray | None = None,\n max: tuple[float, float] | onp.ndarray | None = None,\n step: float | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[tuple[float, float]]:\n \"\"\"Add a length-2 vector input to the GUI.\n\n Args:\n label: Label to display on the vector input.\n initial_value: Initial value of the vector input.\n min: Optional minimum value of the vector input.\n max: Optional maximum value of the vector input.\n step: Optional step size of the vector input. Computed automatically if not\n disabled: Whether the vector input is disabled.\n visible: Whether the vector input is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = initial_value\n value = cast_vector(value, 2)\n min = cast_vector(min, 2) if min is not None else None\n max = cast_vector(max, 2) if max is not None else None\n id = _make_unique_id()\n order = _apply_default_order(order)\n\n if step is None:\n possible_steps: list[float] = []\n possible_steps.extend([_compute_step(x) for x in value])\n if min is not None:\n possible_steps.extend([_compute_step(x) for x in min])\n if max is not None:\n possible_steps.extend([_compute_step(x) for x in max])\n step = float(onp.min(possible_steps))\n\n return self._create_gui_input(\n value,\n message=_messages.GuiAddVector2Message(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n min=min,\n max=max,\n step=step,\n precision=_compute_precision_digits(step),\n disabled=disabled,\n visible=visible,\n ),\n )\n\n def add_vector3(\n self,\n label: str,\n initial_value: tuple[float, float, float] | onp.ndarray,\n min: tuple[float, float, float] | onp.ndarray | None = None,\n max: tuple[float, float, float] | onp.ndarray | None = None,\n step: float | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[tuple[float, float, float]]:\n \"\"\"Add a length-3 vector input to the GUI.\n\n Args:\n label: Label to display on the vector input.\n initial_value: Initial value of the vector input.\n min: Optional minimum value of the vector input.\n max: Optional maximum value of the vector input.\n step: Optional step size of the vector input. Computed automatically if not\n disabled: Whether the vector input is disabled.\n visible: Whether the vector input is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = initial_value\n value = cast_vector(value, 3)\n min = cast_vector(min, 3) if min is not None else None\n max = cast_vector(max, 3) if max is not None else None\n id = _make_unique_id()\n order = _apply_default_order(order)\n\n if step is None:\n possible_steps: list[float] = []\n possible_steps.extend([_compute_step(x) for x in value])\n if min is not None:\n possible_steps.extend([_compute_step(x) for x in min])\n if max is not None:\n possible_steps.extend([_compute_step(x) for x in max])\n step = float(onp.min(possible_steps))\n\n return self._create_gui_input(\n value,\n message=_messages.GuiAddVector3Message(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n min=min,\n max=max,\n step=step,\n precision=_compute_precision_digits(step),\n disabled=disabled,\n visible=visible,\n ),\n )\n\n # See add_dropdown for notes on overloads.\n @overload\n def add_dropdown(\n self,\n label: str,\n options: Sequence[TLiteralString],\n initial_value: TLiteralString | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiDropdownHandle[TLiteralString]: ...\n\n @overload\n def add_dropdown(\n self,\n label: str,\n options: Sequence[TString],\n initial_value: TString | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiDropdownHandle[TString]: ...\n\n def add_dropdown(\n self,\n label: str,\n options: Sequence[TLiteralString] | Sequence[TString],\n initial_value: TLiteralString | TString | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiDropdownHandle[Any]: # Output type is specified in overloads.\n \"\"\"Add a dropdown to the GUI.\n\n Args:\n label: Label to display on the dropdown.\n options: Sequence of options to display in the dropdown.\n initial_value: Initial value of the dropdown.\n disabled: Whether the dropdown is disabled.\n visible: Whether the dropdown is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = initial_value\n if value is None:\n value = options[0]\n id = _make_unique_id()\n order = _apply_default_order(order)\n return GuiDropdownHandle(\n self._create_gui_input(\n value,\n message=_messages.GuiAddDropdownMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n options=tuple(options),\n disabled=disabled,\n visible=visible,\n ),\n )._impl,\n _impl_options=tuple(options),\n )\n\n def add_progress_bar(\n self,\n value: float,\n visible: bool = True,\n animated: bool = False,\n color: Color | None = None,\n order: float | None = None,\n ) -> GuiProgressBarHandle:\n \"\"\"Add a progress bar to the GUI.\n\n Args:\n value: Value of the progress bar. (0 - 100)\n visible: Whether the progress bar is visible.\n animated: Whether the progress bar is in a loading state (animated, striped).\n color: The color of the progress bar.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n assert value >= 0 and value <= 100\n handle = GuiProgressBarHandle(\n _gui_api=self,\n _id=_make_unique_id(),\n _visible=visible,\n _animated=animated,\n _parent_container_id=self._get_container_id(),\n _order=_apply_default_order(order),\n _value=value,\n )\n self._websock_interface.queue_message(\n _messages.GuiAddProgressBarMessage(\n order=handle._order,\n id=handle._id,\n value=value,\n animated=animated,\n color=color,\n container_id=handle._parent_container_id,\n visible=visible,\n )\n )\n return handle\n\n def add_slider(\n self,\n label: str,\n min: IntOrFloat,\n max: IntOrFloat,\n step: IntOrFloat,\n initial_value: IntOrFloat,\n marks: tuple[IntOrFloat | tuple[IntOrFloat, str], ...] | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[IntOrFloat]:\n \"\"\"Add a slider to the GUI. Types of the min, max, step, and initial value should match.\n\n Args:\n label: Label to display on the slider.\n min: Minimum value of the slider.\n max: Maximum value of the slider.\n step: Step size of the slider.\n initial_value: Initial value of the slider.\n marks: tuple of marks to display below the slider. Each mark should\n either be a numerical or a (number, label) tuple, where the\n label is provided as a string.\n disabled: Whether the slider is disabled.\n visible: Whether the slider is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value: IntOrFloat = initial_value\n assert max >= min\n step = builtins.min(step, max - min)\n assert max >= value >= min\n\n # GUI callbacks cast incoming values to match the type of the initial value. If\n # the min, max, or step is a float, we should cast to a float.\n #\n # This should also match what the IntOrFloat TypeVar resolves to.\n if type(value) is int and (\n type(min) is float or type(max) is float or type(step) is float\n ):\n value = float(value) # type: ignore\n\n # TODO: as of 6/5/2023, this assert will break something in nerfstudio. (at\n # least LERF)\n #\n # assert type(min) == type(max) == type(step) == type(value)\n\n id = _make_unique_id()\n order = _apply_default_order(order)\n return self._create_gui_input(\n value,\n message=_messages.GuiAddSliderMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n min=min,\n max=max,\n step=step,\n value=value,\n precision=_compute_precision_digits(step),\n visible=visible,\n disabled=disabled,\n marks=tuple(\n {\"value\": float(x[0]), \"label\": x[1]}\n if isinstance(x, tuple)\n else {\"value\": float(x)}\n for x in marks\n )\n if marks is not None\n else None,\n ),\n is_button=False,\n )\n\n def add_multi_slider(\n self,\n label: str,\n min: IntOrFloat,\n max: IntOrFloat,\n step: IntOrFloat,\n initial_value: tuple[IntOrFloat, ...],\n min_range: IntOrFloat | None = None,\n fixed_endpoints: bool = False,\n marks: tuple[IntOrFloat | tuple[IntOrFloat, str], ...] | None = None,\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[tuple[IntOrFloat, ...]]:\n \"\"\"Add a multi slider to the GUI. Types of the min, max, step, and initial value should match.\n\n Args:\n label: Label to display on the slider.\n min: Minimum value of the slider.\n max: Maximum value of the slider.\n step: Step size of the slider.\n initial_value: Initial values of the slider.\n min_range: Optional minimum difference between two values of the slider.\n fixed_endpoints: Whether the endpoints of the slider are fixed.\n marks: tuple of marks to display below the slider. Each mark should\n either be a numerical or a (number, label) tuple, where the\n label is provided as a string.\n disabled: Whether the slider is disabled.\n visible: Whether the slider is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n assert max >= min\n step = builtins.min(step, max - min)\n assert all(max >= x >= min for x in initial_value)\n\n # GUI callbacks cast incoming values to match the type of the initial value. If\n # any of the arguments are floats, we should always use a float value.\n #\n # This should also match what the IntOrFloat TypeVar resolves to.\n if (\n type(min) is float\n or type(max) is float\n or type(step) is float\n or type(min_range) is float\n ):\n initial_value = tuple(float(x) for x in initial_value) # type: ignore\n\n id = _make_unique_id()\n order = _apply_default_order(order)\n return self._create_gui_input(\n value=initial_value,\n message=_messages.GuiAddMultiSliderMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n min=min,\n min_range=min_range,\n max=max,\n step=step,\n value=initial_value,\n visible=visible,\n disabled=disabled,\n fixed_endpoints=fixed_endpoints,\n precision=_compute_precision_digits(step),\n marks=tuple(\n {\"value\": float(x[0]), \"label\": x[1]}\n if isinstance(x, tuple)\n else {\"value\": float(x)}\n for x in marks\n )\n if marks is not None\n else None,\n ),\n is_button=False,\n )\n\n def add_rgb(\n self,\n label: str,\n initial_value: tuple[int, int, int],\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[tuple[int, int, int]]:\n \"\"\"Add an RGB picker to the GUI.\n\n Args:\n label: Label to display on the RGB picker.\n initial_value: Initial value of the RGB picker.\n disabled: Whether the RGB picker is disabled.\n visible: Whether the RGB picker is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n\n value = initial_value\n id = _make_unique_id()\n order = _apply_default_order(order)\n return self._create_gui_input(\n value,\n message=_messages.GuiAddRgbMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n disabled=disabled,\n visible=visible,\n ),\n )\n\n def add_rgba(\n self,\n label: str,\n initial_value: tuple[int, int, int, int],\n disabled: bool = False,\n visible: bool = True,\n hint: str | None = None,\n order: float | None = None,\n ) -> GuiInputHandle[tuple[int, int, int, int]]:\n \"\"\"Add an RGBA picker to the GUI.\n\n Args:\n label: Label to display on the RGBA picker.\n initial_value: Initial value of the RGBA picker.\n disabled: Whether the RGBA picker is disabled.\n visible: Whether the RGBA picker is visible.\n hint: Optional hint to display on hover.\n order: Optional ordering, smallest values will be displayed first.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n value = initial_value\n id = _make_unique_id()\n order = _apply_default_order(order)\n return self._create_gui_input(\n value,\n message=_messages.GuiAddRgbaMessage(\n order=order,\n id=id,\n label=label,\n container_id=self._get_container_id(),\n hint=hint,\n value=value,\n disabled=disabled,\n visible=visible,\n ),\n )\n\n def _create_gui_input(\n self,\n value: T,\n message: _messages._GuiAddInputBase,\n is_button: bool = False,\n ) -> GuiInputHandle[T]:\n \"\"\"Private helper for adding a simple GUI element.\"\"\"\n\n # Send add GUI input message.\n self._websock_interface.queue_message(message)\n\n # Construct handle.\n handle_state = _GuiHandleState(\n label=message.label,\n message_type=type(message),\n gui_api=self,\n value=value,\n update_timestamp=time.time(),\n parent_container_id=self._get_container_id(),\n update_cb=[],\n is_button=is_button,\n sync_cb=None,\n disabled=message.disabled,\n visible=message.visible,\n id=message.id,\n order=message.order,\n hint=message.hint,\n )\n\n # For broadcasted GUI handles, we should synchronize all clients.\n # This will be a no-op for client handles.\n if not is_button:\n\n def sync_other_clients(\n client_id: ClientId, updates: dict[str, Any]\n ) -> None:\n message = _messages.GuiUpdateMessage(handle_state.id, updates)\n message.excluded_self_client = client_id\n self._websock_interface.queue_message(message)\n\n handle_state.sync_cb = sync_other_clients\n\n handle = GuiInputHandle(handle_state)\n\n return handle\n" }, { "path": "viser/src/viser/_gui_handles.py", "content": "from __future__ import annotations\n\nimport base64\nimport dataclasses\nimport re\nimport time\nimport uuid\nimport warnings\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Any, Callable, Generic, Iterable, TypeVar\n\nimport imageio.v3 as iio\nimport numpy as onp\nfrom typing_extensions import Protocol\n\nfrom ._icons import svg_from_icon\nfrom ._icons_enum import IconName\nfrom ._messages import GuiCloseModalMessage, GuiRemoveMessage, GuiUpdateMessage, Message\nfrom ._scene_api import _encode_image_binary\nfrom .infra import ClientId\n\nif TYPE_CHECKING:\n import plotly.graph_objects as go\n\n from ._gui_api import GuiApi\n from ._viser import ClientHandle\n\n\nT = TypeVar(\"T\")\nTGuiHandle = TypeVar(\"TGuiHandle\", bound=\"_GuiInputHandle\")\n\n\ndef _make_unique_id() -> str:\n \"\"\"Return a unique ID for referencing GUI elements.\"\"\"\n return str(uuid.uuid4())\n\n\nclass GuiContainerProtocol(Protocol):\n _children: dict[str, SupportsRemoveProtocol] = dataclasses.field(\n default_factory=dict\n )\n\n\nclass SupportsRemoveProtocol(Protocol):\n def remove(self) -> None: ...\n\n\n@dataclasses.dataclass\nclass _GuiHandleState(Generic[T]):\n \"\"\"Internal API for GUI elements.\"\"\"\n\n label: str\n gui_api: GuiApi\n value: T\n update_timestamp: float\n\n parent_container_id: str\n \"\"\"Container that this GUI input was placed into.\"\"\"\n\n update_cb: list[Callable[[GuiEvent], None]]\n \"\"\"Registered functions to call when this input is updated.\"\"\"\n\n is_button: bool\n \"\"\"Indicates a button element, which requires special handling.\"\"\"\n\n sync_cb: Callable[[ClientId, dict[str, Any]], None] | None\n \"\"\"Callback for synchronizing inputs across clients.\"\"\"\n\n disabled: bool\n visible: bool\n\n order: float\n id: str\n hint: str | None\n\n message_type: type[Message]\n\n\n@dataclasses.dataclass\nclass _GuiInputHandle(Generic[T]):\n # Let's shove private implementation details in here...\n _impl: _GuiHandleState[T]\n\n # Should we use @property for get_value / set_value, set_hidden, etc?\n #\n # Benefits:\n # @property is syntactically very nice.\n # `gui.value = ...` is really tempting!\n # Feels a bit more magical.\n #\n # Downsides:\n # Consistency: not everything that can be written can be read, and not everything\n # that can be read can be written. `get_`/`set_` makes this really clear.\n # Clarity: some things that we read (like client mappings) are copied before\n # they're returned. An attribute access obfuscates the overhead here.\n # Flexibility: getter/setter types should match. https://github.com/python/mypy/issues/3004\n # Feels a bit more magical.\n #\n # Is this worth the tradeoff?\n\n @property\n def order(self) -> float:\n \"\"\"Read-only order value, which dictates the position of the GUI element.\"\"\"\n return self._impl.order\n\n @property\n def value(self) -> T:\n \"\"\"Value of the GUI input. Synchronized automatically when assigned.\"\"\"\n return self._impl.value\n\n @value.setter\n def value(self, value: T | onp.ndarray) -> None:\n if isinstance(value, onp.ndarray):\n assert len(value.shape) <= 1, f\"{value.shape} should be at most 1D!\"\n value = tuple(map(float, value)) # type: ignore\n\n # Send to client, except for buttons.\n if not self._impl.is_button:\n self._impl.gui_api._websock_interface.queue_message(\n GuiUpdateMessage(self._impl.id, {\"value\": value})\n )\n\n # Set internal state. We automatically convert numpy arrays to the expected\n # internal type. (eg 1D arrays to tuples)\n self._impl.value = type(self._impl.value)(value) # type: ignore\n self._impl.update_timestamp = time.time()\n\n # Call update callbacks.\n for cb in self._impl.update_cb:\n # Pushing callbacks into separate threads helps prevent deadlocks when we\n # have a lock in a callback. TODO: revisit other callbacks.\n self._impl.gui_api._thread_executor.submit(\n lambda: cb(\n GuiEvent(\n client_id=None,\n client=None,\n target=self,\n )\n )\n )\n\n @property\n def update_timestamp(self) -> float:\n \"\"\"Read-only timestamp when this input was last updated.\"\"\"\n return self._impl.update_timestamp\n\n @property\n def disabled(self) -> bool:\n \"\"\"Allow/disallow user interaction with the input. Synchronized automatically\n when assigned.\"\"\"\n return self._impl.disabled\n\n @disabled.setter\n def disabled(self, disabled: bool) -> None:\n if disabled == self.disabled:\n return\n\n self._impl.gui_api._websock_interface.queue_message(\n GuiUpdateMessage(self._impl.id, {\"disabled\": disabled})\n )\n self._impl.disabled = disabled\n\n @property\n def visible(self) -> bool:\n \"\"\"Temporarily show or hide this GUI element from the visualizer. Synchronized\n automatically when assigned.\"\"\"\n return self._impl.visible\n\n @visible.setter\n def visible(self, visible: bool) -> None:\n if visible == self.visible:\n return\n\n self._impl.gui_api._websock_interface.queue_message(\n GuiUpdateMessage(self._impl.id, {\"visible\": visible})\n )\n self._impl.visible = visible\n\n def __post_init__(self) -> None:\n \"\"\"We need to register ourself after construction for callbacks to work.\"\"\"\n gui_api = self._impl.gui_api\n\n # TODO: the current way we track GUI handles and children is very manual +\n # error-prone. We should revist this design.\n gui_api._gui_input_handle_from_id[self._impl.id] = self\n parent = gui_api._container_handle_from_id[self._impl.parent_container_id]\n parent._children[self._impl.id] = self\n\n def remove(self) -> None:\n \"\"\"Permanently remove this GUI element from the visualizer.\"\"\"\n gui_api = self._impl.gui_api\n gui_api._websock_interface.queue_message(GuiRemoveMessage(self._impl.id))\n gui_api._gui_input_handle_from_id.pop(self._impl.id)\n parent = gui_api._container_handle_from_id[self._impl.parent_container_id]\n parent._children.pop(self._impl.id)\n\n\nStringType = TypeVar(\"StringType\", bound=str)\n\n\n# GuiInputHandle[T] is used for all inputs except for buttons.\n#\n# We inherit from _GuiInputHandle to special-case buttons because the usage semantics\n# are slightly different: we have `on_click()` instead of `on_update()`.\n@dataclasses.dataclass\nclass GuiInputHandle(_GuiInputHandle[T], Generic[T]):\n \"\"\"A handle is created for each GUI element that is added in `viser`.\n Handles can be used to read and write state.\n\n When a GUI element is added via :attr:`ViserServer.gui`, state is\n synchronized between all connected clients. When a GUI element is added via\n :attr:`ClientHandle.gui`, state is local to a specific client.\n \"\"\"\n\n def on_update(\n self: TGuiHandle, func: Callable[[GuiEvent[TGuiHandle]], None]\n ) -> Callable[[GuiEvent[TGuiHandle]], None]:\n \"\"\"Attach a function to call when a GUI input is updated. Happens in a thread.\"\"\"\n self._impl.update_cb.append(func)\n return func\n\n\n@dataclasses.dataclass(frozen=True)\nclass GuiEvent(Generic[TGuiHandle]):\n \"\"\"Information associated with a GUI event, such as an update or click.\n\n Passed as input to callback functions.\"\"\"\n\n client: ClientHandle | None\n \"\"\"Client that triggered this event.\"\"\"\n client_id: int | None\n \"\"\"ID of client that triggered this event.\"\"\"\n target: TGuiHandle\n \"\"\"GUI element that was affected.\"\"\"\n\n\n@dataclasses.dataclass\nclass GuiButtonHandle(_GuiInputHandle[bool]):\n \"\"\"Handle for a button input in our visualizer.\n\n Lets us detect clicks.\"\"\"\n\n def on_click(\n self: TGuiHandle, func: Callable[[GuiEvent[TGuiHandle]], None]\n ) -> Callable[[GuiEvent[TGuiHandle]], None]:\n \"\"\"Attach a function to call when a button is pressed. Happens in a thread.\"\"\"\n self._impl.update_cb.append(func)\n return func\n\n\n@dataclasses.dataclass\nclass UploadedFile:\n \"\"\"Result of a file upload.\"\"\"\n\n name: str\n \"\"\"Name of the file.\"\"\"\n content: bytes\n \"\"\"Contents of the file.\"\"\"\n\n\n@dataclasses.dataclass\nclass GuiUploadButtonHandle(_GuiInputHandle[UploadedFile]):\n \"\"\"Handle for an upload file button in our visualizer.\n\n The `.value` attribute will be updated with the contents of uploaded files.\n \"\"\"\n\n def on_upload(\n self: TGuiHandle, func: Callable[[GuiEvent[TGuiHandle]], None]\n ) -> Callable[[GuiEvent[TGuiHandle]], None]:\n \"\"\"Attach a function to call when a button is pressed. Happens in a thread.\"\"\"\n self._impl.update_cb.append(func)\n return func\n\n\n@dataclasses.dataclass\nclass GuiButtonGroupHandle(_GuiInputHandle[StringType], Generic[StringType]):\n \"\"\"Handle for a button group input in our visualizer.\n\n Lets us detect clicks.\"\"\"\n\n def on_click(\n self: TGuiHandle, func: Callable[[GuiEvent[TGuiHandle]], None]\n ) -> Callable[[GuiEvent[TGuiHandle]], None]:\n \"\"\"Attach a function to call when a button is pressed. Happens in a thread.\"\"\"\n self._impl.update_cb.append(func)\n return func\n\n @property\n def disabled(self) -> bool:\n \"\"\"Button groups cannot be disabled.\"\"\"\n return False\n\n @disabled.setter\n def disabled(self, disabled: bool) -> None:\n \"\"\"Button groups cannot be disabled.\"\"\"\n assert not disabled, \"Button groups cannot be disabled.\"\n\n\n@dataclasses.dataclass\nclass GuiDropdownHandle(GuiInputHandle[StringType], Generic[StringType]):\n \"\"\"Handle for a dropdown-style GUI input in our visualizer.\n\n Lets us get values, set values, and detect updates.\"\"\"\n\n _impl_options: tuple[StringType, ...]\n\n @property\n def options(self) -> tuple[StringType, ...]:\n \"\"\"Options for our dropdown. Synchronized automatically when assigned.\n\n For projects that care about typing: the static type of `options` should be\n consistent with the `StringType` associated with a handle. Literal types will be\n inferred where possible when handles are instantiated; for the most flexibility,\n we can declare handles as `GuiDropdownHandle[str]`.\n \"\"\"\n return self._impl_options\n\n @options.setter\n def options(self, options: Iterable[StringType]) -> None:\n self._impl_options = tuple(options)\n\n need_to_overwrite_value = self.value not in self._impl_options\n if need_to_overwrite_value:\n self._impl.gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._impl.id,\n {\"options\": self._impl_options, \"value\": self._impl_options[0]},\n )\n )\n self._impl.value = self._impl_options[0]\n else:\n self._impl.gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._impl.id,\n {\"options\": self._impl_options},\n )\n )\n\n\n@dataclasses.dataclass(frozen=True)\nclass GuiTabGroupHandle:\n \"\"\"Handle for a tab group. Call :meth:`add_tab()` to add a tab.\"\"\"\n\n _tab_group_id: str\n _labels: list[str]\n _icons_html: list[str | None]\n _tabs: list[GuiTabHandle]\n _gui_api: GuiApi\n _parent_container_id: str\n _order: float\n\n @property\n def order(self) -> float:\n \"\"\"Read-only order value, which dictates the position of the GUI element.\"\"\"\n return self._order\n\n def add_tab(self, label: str, icon: IconName | None = None) -> GuiTabHandle:\n \"\"\"Add a tab. Returns a handle we can use to add GUI elements to it.\"\"\"\n\n id = _make_unique_id()\n\n # We may want to make this thread-safe in the future.\n out = GuiTabHandle(_parent=self, _id=id)\n\n self._labels.append(label)\n self._icons_html.append(None if icon is None else svg_from_icon(icon))\n self._tabs.append(out)\n\n self._sync_with_client()\n return out\n\n def __post_init__(self) -> None:\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children[self._tab_group_id] = self\n\n def remove(self) -> None:\n \"\"\"Remove this tab group and all contained GUI elements.\"\"\"\n for tab in tuple(self._tabs):\n tab.remove()\n gui_api = self._gui_api\n gui_api._websock_interface.queue_message(GuiRemoveMessage(self._tab_group_id))\n parent = gui_api._container_handle_from_id[self._parent_container_id]\n parent._children.pop(self._tab_group_id)\n\n def _sync_with_client(self) -> None:\n \"\"\"Send messages for syncing tab state with the client.\"\"\"\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._tab_group_id,\n {\n \"tab_labels\": tuple(self._labels),\n \"tab_icons_html\": tuple(self._icons_html),\n \"tab_container_ids\": tuple(tab._id for tab in self._tabs),\n },\n )\n )\n\n\n@dataclasses.dataclass\nclass GuiFolderHandle:\n \"\"\"Use as a context to place GUI elements into a folder.\"\"\"\n\n _gui_api: GuiApi\n _id: str # Used as container ID for children.\n _order: float\n _parent_container_id: str # Container ID of parent.\n _container_id_restore: str | None = None\n _children: dict[str, SupportsRemoveProtocol] = dataclasses.field(\n default_factory=dict\n )\n\n @property\n def order(self) -> float:\n \"\"\"Read-only order value, which dictates the position of the GUI element.\"\"\"\n return self._order\n\n def __enter__(self) -> GuiFolderHandle:\n self._container_id_restore = self._gui_api._get_container_id()\n self._gui_api._set_container_id(self._id)\n return self\n\n def __exit__(self, *args) -> None:\n del args\n assert self._container_id_restore is not None\n self._gui_api._set_container_id(self._container_id_restore)\n self._container_id_restore = None\n\n def __post_init__(self) -> None:\n self._gui_api._container_handle_from_id[self._id] = self\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children[self._id] = self\n\n def remove(self) -> None:\n \"\"\"Permanently remove this folder and all contained GUI elements from the\n visualizer.\"\"\"\n self._gui_api._websock_interface.queue_message(GuiRemoveMessage(self._id))\n for child in tuple(self._children.values()):\n child.remove()\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children.pop(self._id)\n self._gui_api._container_handle_from_id.pop(self._id)\n\n\n@dataclasses.dataclass\nclass GuiModalHandle:\n \"\"\"Use as a context to place GUI elements into a modal.\"\"\"\n\n _gui_api: GuiApi\n _id: str # Used as container ID of children.\n _container_id_restore: str | None = None\n _children: dict[str, SupportsRemoveProtocol] = dataclasses.field(\n default_factory=dict\n )\n\n def __enter__(self) -> GuiModalHandle:\n self._container_id_restore = self._gui_api._get_container_id()\n self._gui_api._set_container_id(self._id)\n return self\n\n def __exit__(self, *args) -> None:\n del args\n assert self._container_id_restore is not None\n self._gui_api._set_container_id(self._container_id_restore)\n self._container_id_restore = None\n\n def __post_init__(self) -> None:\n self._gui_api._container_handle_from_id[self._id] = self\n\n def close(self) -> None:\n \"\"\"Close this modal and permananently remove all contained GUI elements.\"\"\"\n self._gui_api._websock_interface.queue_message(\n GuiCloseModalMessage(self._id),\n )\n for child in tuple(self._children.values()):\n child.remove()\n self._gui_api._container_handle_from_id.pop(self._id)\n\n\n@dataclasses.dataclass\nclass GuiTabHandle:\n \"\"\"Use as a context to place GUI elements into a tab.\"\"\"\n\n _parent: GuiTabGroupHandle\n _id: str # Used as container ID of children.\n _container_id_restore: str | None = None\n _children: dict[str, SupportsRemoveProtocol] = dataclasses.field(\n default_factory=dict\n )\n\n def __enter__(self) -> GuiTabHandle:\n self._container_id_restore = self._parent._gui_api._get_container_id()\n self._parent._gui_api._set_container_id(self._id)\n return self\n\n def __exit__(self, *args) -> None:\n del args\n assert self._container_id_restore is not None\n self._parent._gui_api._set_container_id(self._container_id_restore)\n self._container_id_restore = None\n\n def __post_init__(self) -> None:\n self._parent._gui_api._container_handle_from_id[self._id] = self\n\n def remove(self) -> None:\n \"\"\"Permanently remove this tab and all contained GUI elements from the\n visualizer.\"\"\"\n # We may want to make this thread-safe in the future.\n container_index = -1\n for i, tab in enumerate(self._parent._tabs):\n if tab is self:\n container_index = i\n break\n assert container_index != -1, \"Tab already removed!\"\n\n self._parent._labels.pop(container_index)\n self._parent._icons_html.pop(container_index)\n self._parent._tabs.pop(container_index)\n self._parent._sync_with_client()\n for child in tuple(self._children.values()):\n child.remove()\n self._parent._gui_api._container_handle_from_id.pop(self._id)\n\n\ndef _get_data_url(url: str, image_root: Path | None) -> str:\n if not url.startswith(\"http\") and not image_root:\n warnings.warn(\n (\n \"No `image_root` provided. All relative paths will be scoped to viser's\"\n \" installation path.\"\n ),\n stacklevel=2,\n )\n if url.startswith(\"http\") or url.startswith(\"data:\"):\n return url\n if image_root is None:\n image_root = Path(__file__).parent\n try:\n image = iio.imread(image_root / url)\n media_type, binary = _encode_image_binary(image, \"png\")\n url = base64.b64encode(binary).decode(\"utf-8\")\n return f\"data:{media_type};base64,{url}\"\n except (IOError, FileNotFoundError):\n warnings.warn(\n f\"Failed to read image {url}, with image_root set to {image_root}.\",\n stacklevel=2,\n )\n return url\n\n\ndef _parse_markdown(markdown: str, image_root: Path | None) -> str:\n markdown = re.sub(\n r\"\\!\\[([^]]*)\\]\\(([^]]*)\\)\",\n lambda match: (\n f\"![{match.group(1)}]({_get_data_url(match.group(2), image_root)})\"\n ),\n markdown,\n )\n return markdown\n\n\n@dataclasses.dataclass\nclass GuiProgressBarHandle:\n \"\"\"Use to remove markdown.\"\"\"\n\n _gui_api: GuiApi\n _id: str\n _visible: bool\n _animated: bool\n _parent_container_id: str\n _order: float\n _value: float\n\n @property\n def value(self) -> float:\n \"\"\"Current content of this progress bar element, 0 - 100. Synchronized\n automatically when assigned.\"\"\"\n return self._value\n\n @value.setter\n def value(self, value: float) -> None:\n assert value >= 0 and value <= 100\n self._value = value\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._id,\n {\"value\": value},\n )\n )\n\n @property\n def animated(self) -> bool:\n \"\"\"Show this progress bar as loading (animated, striped).\"\"\"\n return self._animated\n\n @animated.setter\n def animated(self, animated: bool) -> None:\n self._animated = animated\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._id,\n {\"animated\": animated},\n )\n )\n\n @property\n def order(self) -> float:\n \"\"\"Read-only order value, which dictates the position of the GUI element.\"\"\"\n return self._order\n\n @property\n def visible(self) -> bool:\n \"\"\"Temporarily show or hide this GUI element from the visualizer. Synchronized\n automatically when assigned.\"\"\"\n return self._visible\n\n @visible.setter\n def visible(self, visible: bool) -> None:\n if visible == self.visible:\n return\n\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(self._id, {\"visible\": visible})\n )\n self._visible = visible\n\n def __post_init__(self) -> None:\n \"\"\"We need to register ourself after construction for callbacks to work.\"\"\"\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children[self._id] = self\n\n def remove(self) -> None:\n \"\"\"Permanently remove this progress bar from the visualizer.\"\"\"\n self._gui_api._websock_interface.queue_message(GuiRemoveMessage(self._id))\n\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children.pop(self._id)\n\n\n@dataclasses.dataclass\nclass GuiMarkdownHandle:\n \"\"\"Use to remove markdown.\"\"\"\n\n _gui_api: GuiApi\n _id: str\n _visible: bool\n _parent_container_id: str\n _order: float\n _image_root: Path | None\n _content: str | None\n\n @property\n def content(self) -> str:\n \"\"\"Current content of this markdown element. Synchronized automatically when assigned.\"\"\"\n assert self._content is not None\n return self._content\n\n @content.setter\n def content(self, content: str) -> None:\n self._content = content\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._id,\n {\"markdown\": _parse_markdown(content, self._image_root)},\n )\n )\n\n @property\n def order(self) -> float:\n \"\"\"Read-only order value, which dictates the position of the GUI element.\"\"\"\n return self._order\n\n @property\n def visible(self) -> bool:\n \"\"\"Temporarily show or hide this GUI element from the visualizer. Synchronized\n automatically when assigned.\"\"\"\n return self._visible\n\n @visible.setter\n def visible(self, visible: bool) -> None:\n if visible == self.visible:\n return\n\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(self._id, {\"visible\": visible})\n )\n self._visible = visible\n\n def __post_init__(self) -> None:\n \"\"\"We need to register ourself after construction for callbacks to work.\"\"\"\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children[self._id] = self\n\n def remove(self) -> None:\n \"\"\"Permanently remove this markdown from the visualizer.\"\"\"\n self._gui_api._websock_interface.queue_message(GuiRemoveMessage(self._id))\n\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children.pop(self._id)\n\n\n@dataclasses.dataclass\nclass GuiPlotlyHandle:\n \"\"\"Use to update or remove markdown elements.\"\"\"\n\n _gui_api: GuiApi\n _id: str\n _visible: bool\n _parent_container_id: str\n _order: float\n _figure: go.Figure | None\n _aspect: float | None\n\n @property\n def figure(self) -> go.Figure:\n \"\"\"Current content of this markdown element. Synchronized automatically when assigned.\"\"\"\n assert self._figure is not None\n return self._figure\n\n @figure.setter\n def figure(self, figure: go.Figure) -> None:\n self._figure = figure\n\n json_str = figure.to_json()\n assert isinstance(json_str, str)\n\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._id,\n {\"plotly_json_str\": json_str},\n )\n )\n\n @property\n def aspect(self) -> float:\n \"\"\"Aspect ratio of the plotly figure, in the control panel.\"\"\"\n assert self._aspect is not None\n return self._aspect\n\n @aspect.setter\n def aspect(self, aspect: float) -> None:\n self._aspect = aspect\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(\n self._id,\n {\"aspect\": aspect},\n )\n )\n\n @property\n def order(self) -> float:\n \"\"\"Read-only order value, which dictates the position of the GUI element.\"\"\"\n return self._order\n\n @property\n def visible(self) -> bool:\n \"\"\"Temporarily show or hide this GUI element from the visualizer. Synchronized\n automatically when assigned.\"\"\"\n return self._visible\n\n @visible.setter\n def visible(self, visible: bool) -> None:\n if visible == self.visible:\n return\n\n self._gui_api._websock_interface.queue_message(\n GuiUpdateMessage(self._id, {\"visible\": visible})\n )\n self._visible = visible\n\n def __post_init__(self) -> None:\n \"\"\"We need to register ourself after construction for callbacks to work.\"\"\"\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children[self._id] = self\n\n def remove(self) -> None:\n \"\"\"Permanently remove this figure from the visualizer.\"\"\"\n self._gui_api._websock_interface.queue_message(GuiRemoveMessage(self._id))\n parent = self._gui_api._container_handle_from_id[self._parent_container_id]\n parent._children.pop(self._id)\n" }, { "path": "viser/src/viser/_icons.py", "content": "import tarfile\nfrom pathlib import Path\n\nfrom ._icons_enum import IconName\n\nICONS_DIR = Path(__file__).absolute().parent / \"_icons\"\n\n\ndef svg_from_icon(icon_name: IconName) -> str:\n \"\"\"Read an icon and return it as a UTF string; we expect this to be an\n tag.\"\"\"\n assert isinstance(icon_name, str)\n icons_tarball = ICONS_DIR / \"tabler-icons.tar\"\n\n with tarfile.open(icons_tarball) as tar:\n icon_file = tar.extractfile(f\"{icon_name}.svg\")\n assert icon_file is not None\n out = icon_file.read()\n\n return out.decode(\"utf-8\")\n" }, { "path": "viser/src/viser/_icons_enum.py", "content": "# Automatically generated by `_icons_generate_enum.py`\n# See https://tabler-icons.io/\nfrom typing import NewType\n\nIconName = NewType(\"IconName\", str)\n\"\"\"Name of an icon. Should be generated via `viser.Icon.*`.\"\"\"\n\n\nclass _IconStringConverter(type):\n def __getattr__(self, __name: str) -> IconName:\n if not __name.startswith(\"_\"):\n return IconName(__name.lower().replace(\"_\", \"-\"))\n else:\n raise AttributeError()\n\n\nclass Icon(metaclass=_IconStringConverter):\n \"\"\"'Enum' class for referencing Tabler icons.\n\n We don't subclass enum.Enum for performance reasons -- importing an enum with\n thousands of names can result in import times in the hundreds of milliseconds.\n\n Attributes:\n ICON_123 (IconName): The :code:`123` icon.\n ICON_24_HOURS (IconName): The :code:`24-hours` icon.\n ICON_2FA (IconName): The :code:`2fa` icon.\n ICON_360 (IconName): The :code:`360` icon.\n ICON_360_VIEW (IconName): The :code:`360-view` icon.\n ICON_3D_CUBE_SPHERE (IconName): The :code:`3d-cube-sphere` icon.\n ICON_3D_CUBE_SPHERE_OFF (IconName): The :code:`3d-cube-sphere-off` icon.\n ICON_3D_ROTATE (IconName): The :code:`3d-rotate` icon.\n A_B (IconName): The :code:`a-b` icon.\n A_B_2 (IconName): The :code:`a-b-2` icon.\n A_B_OFF (IconName): The :code:`a-b-off` icon.\n ABACUS (IconName): The :code:`abacus` icon.\n ABACUS_OFF (IconName): The :code:`abacus-off` icon.\n ABC (IconName): The :code:`abc` icon.\n ACCESS_POINT (IconName): The :code:`access-point` icon.\n ACCESS_POINT_OFF (IconName): The :code:`access-point-off` icon.\n ACCESSIBLE (IconName): The :code:`accessible` icon.\n ACCESSIBLE_OFF (IconName): The :code:`accessible-off` icon.\n ACCESSIBLE_OFF_FILLED (IconName): The :code:`accessible-off-filled` icon.\n ACTIVITY (IconName): The :code:`activity` icon.\n ACTIVITY_HEARTBEAT (IconName): The :code:`activity-heartbeat` icon.\n AD (IconName): The :code:`ad` icon.\n AD_2 (IconName): The :code:`ad-2` icon.\n AD_CIRCLE (IconName): The :code:`ad-circle` icon.\n AD_CIRCLE_FILLED (IconName): The :code:`ad-circle-filled` icon.\n AD_CIRCLE_OFF (IconName): The :code:`ad-circle-off` icon.\n AD_FILLED (IconName): The :code:`ad-filled` icon.\n AD_OFF (IconName): The :code:`ad-off` icon.\n ADDRESS_BOOK (IconName): The :code:`address-book` icon.\n ADDRESS_BOOK_OFF (IconName): The :code:`address-book-off` icon.\n ADJUSTMENTS (IconName): The :code:`adjustments` icon.\n ADJUSTMENTS_ALT (IconName): The :code:`adjustments-alt` icon.\n ADJUSTMENTS_BOLT (IconName): The :code:`adjustments-bolt` icon.\n ADJUSTMENTS_CANCEL (IconName): The :code:`adjustments-cancel` icon.\n ADJUSTMENTS_CHECK (IconName): The :code:`adjustments-check` icon.\n ADJUSTMENTS_CODE (IconName): The :code:`adjustments-code` icon.\n ADJUSTMENTS_COG (IconName): The :code:`adjustments-cog` icon.\n ADJUSTMENTS_DOLLAR (IconName): The :code:`adjustments-dollar` icon.\n ADJUSTMENTS_DOWN (IconName): The :code:`adjustments-down` icon.\n ADJUSTMENTS_EXCLAMATION (IconName): The :code:`adjustments-exclamation` icon.\n ADJUSTMENTS_FILLED (IconName): The :code:`adjustments-filled` icon.\n ADJUSTMENTS_HEART (IconName): The :code:`adjustments-heart` icon.\n ADJUSTMENTS_HORIZONTAL (IconName): The :code:`adjustments-horizontal` icon.\n ADJUSTMENTS_MINUS (IconName): The :code:`adjustments-minus` icon.\n ADJUSTMENTS_OFF (IconName): The :code:`adjustments-off` icon.\n ADJUSTMENTS_PAUSE (IconName): The :code:`adjustments-pause` icon.\n ADJUSTMENTS_PIN (IconName): The :code:`adjustments-pin` icon.\n ADJUSTMENTS_PLUS (IconName): The :code:`adjustments-plus` icon.\n ADJUSTMENTS_QUESTION (IconName): The :code:`adjustments-question` icon.\n ADJUSTMENTS_SEARCH (IconName): The :code:`adjustments-search` icon.\n ADJUSTMENTS_SHARE (IconName): The :code:`adjustments-share` icon.\n ADJUSTMENTS_STAR (IconName): The :code:`adjustments-star` icon.\n ADJUSTMENTS_UP (IconName): The :code:`adjustments-up` icon.\n ADJUSTMENTS_X (IconName): The :code:`adjustments-x` icon.\n AERIAL_LIFT (IconName): The :code:`aerial-lift` icon.\n AFFILIATE (IconName): The :code:`affiliate` icon.\n AFFILIATE_FILLED (IconName): The :code:`affiliate-filled` icon.\n AIR_BALLOON (IconName): The :code:`air-balloon` icon.\n AIR_CONDITIONING (IconName): The :code:`air-conditioning` icon.\n AIR_CONDITIONING_DISABLED (IconName): The :code:`air-conditioning-disabled` icon.\n ALARM (IconName): The :code:`alarm` icon.\n ALARM_FILLED (IconName): The :code:`alarm-filled` icon.\n ALARM_MINUS (IconName): The :code:`alarm-minus` icon.\n ALARM_MINUS_FILLED (IconName): The :code:`alarm-minus-filled` icon.\n ALARM_OFF (IconName): The :code:`alarm-off` icon.\n ALARM_PLUS (IconName): The :code:`alarm-plus` icon.\n ALARM_PLUS_FILLED (IconName): The :code:`alarm-plus-filled` icon.\n ALARM_SNOOZE (IconName): The :code:`alarm-snooze` icon.\n ALARM_SNOOZE_FILLED (IconName): The :code:`alarm-snooze-filled` icon.\n ALBUM (IconName): The :code:`album` icon.\n ALBUM_OFF (IconName): The :code:`album-off` icon.\n ALERT_CIRCLE (IconName): The :code:`alert-circle` icon.\n ALERT_CIRCLE_FILLED (IconName): The :code:`alert-circle-filled` icon.\n ALERT_HEXAGON (IconName): The :code:`alert-hexagon` icon.\n ALERT_HEXAGON_FILLED (IconName): The :code:`alert-hexagon-filled` icon.\n ALERT_OCTAGON (IconName): The :code:`alert-octagon` icon.\n ALERT_OCTAGON_FILLED (IconName): The :code:`alert-octagon-filled` icon.\n ALERT_SMALL (IconName): The :code:`alert-small` icon.\n ALERT_SQUARE (IconName): The :code:`alert-square` icon.\n ALERT_SQUARE_FILLED (IconName): The :code:`alert-square-filled` icon.\n ALERT_SQUARE_ROUNDED (IconName): The :code:`alert-square-rounded` icon.\n ALERT_SQUARE_ROUNDED_FILLED (IconName): The :code:`alert-square-rounded-filled` icon.\n ALERT_TRIANGLE (IconName): The :code:`alert-triangle` icon.\n ALERT_TRIANGLE_FILLED (IconName): The :code:`alert-triangle-filled` icon.\n ALIEN (IconName): The :code:`alien` icon.\n ALIEN_FILLED (IconName): The :code:`alien-filled` icon.\n ALIGN_BOX_BOTTOM_CENTER (IconName): The :code:`align-box-bottom-center` icon.\n ALIGN_BOX_BOTTOM_CENTER_FILLED (IconName): The :code:`align-box-bottom-center-filled` icon.\n ALIGN_BOX_BOTTOM_LEFT (IconName): The :code:`align-box-bottom-left` icon.\n ALIGN_BOX_BOTTOM_LEFT_FILLED (IconName): The :code:`align-box-bottom-left-filled` icon.\n ALIGN_BOX_BOTTOM_RIGHT (IconName): The :code:`align-box-bottom-right` icon.\n ALIGN_BOX_BOTTOM_RIGHT_FILLED (IconName): The :code:`align-box-bottom-right-filled` icon.\n ALIGN_BOX_CENTER_BOTTOM (IconName): The :code:`align-box-center-bottom` icon.\n ALIGN_BOX_CENTER_MIDDLE (IconName): The :code:`align-box-center-middle` icon.\n ALIGN_BOX_CENTER_MIDDLE_FILLED (IconName): The :code:`align-box-center-middle-filled` icon.\n ALIGN_BOX_CENTER_STRETCH (IconName): The :code:`align-box-center-stretch` icon.\n ALIGN_BOX_CENTER_TOP (IconName): The :code:`align-box-center-top` icon.\n ALIGN_BOX_LEFT_BOTTOM (IconName): The :code:`align-box-left-bottom` icon.\n ALIGN_BOX_LEFT_BOTTOM_FILLED (IconName): The :code:`align-box-left-bottom-filled` icon.\n ALIGN_BOX_LEFT_MIDDLE (IconName): The :code:`align-box-left-middle` icon.\n ALIGN_BOX_LEFT_MIDDLE_FILLED (IconName): The :code:`align-box-left-middle-filled` icon.\n ALIGN_BOX_LEFT_STRETCH (IconName): The :code:`align-box-left-stretch` icon.\n ALIGN_BOX_LEFT_TOP (IconName): The :code:`align-box-left-top` icon.\n ALIGN_BOX_LEFT_TOP_FILLED (IconName): The :code:`align-box-left-top-filled` icon.\n ALIGN_BOX_RIGHT_BOTTOM (IconName): The :code:`align-box-right-bottom` icon.\n ALIGN_BOX_RIGHT_BOTTOM_FILLED (IconName): The :code:`align-box-right-bottom-filled` icon.\n ALIGN_BOX_RIGHT_MIDDLE (IconName): The :code:`align-box-right-middle` icon.\n ALIGN_BOX_RIGHT_MIDDLE_FILLED (IconName): The :code:`align-box-right-middle-filled` icon.\n ALIGN_BOX_RIGHT_STRETCH (IconName): The :code:`align-box-right-stretch` icon.\n ALIGN_BOX_RIGHT_TOP (IconName): The :code:`align-box-right-top` icon.\n ALIGN_BOX_RIGHT_TOP_FILLED (IconName): The :code:`align-box-right-top-filled` icon.\n ALIGN_BOX_TOP_CENTER (IconName): The :code:`align-box-top-center` icon.\n ALIGN_BOX_TOP_CENTER_FILLED (IconName): The :code:`align-box-top-center-filled` icon.\n ALIGN_BOX_TOP_LEFT (IconName): The :code:`align-box-top-left` icon.\n ALIGN_BOX_TOP_LEFT_FILLED (IconName): The :code:`align-box-top-left-filled` icon.\n ALIGN_BOX_TOP_RIGHT (IconName): The :code:`align-box-top-right` icon.\n ALIGN_BOX_TOP_RIGHT_FILLED (IconName): The :code:`align-box-top-right-filled` icon.\n ALIGN_CENTER (IconName): The :code:`align-center` icon.\n ALIGN_JUSTIFIED (IconName): The :code:`align-justified` icon.\n ALIGN_LEFT (IconName): The :code:`align-left` icon.\n ALIGN_RIGHT (IconName): The :code:`align-right` icon.\n ALPHA (IconName): The :code:`alpha` icon.\n ALPHABET_CYRILLIC (IconName): The :code:`alphabet-cyrillic` icon.\n ALPHABET_GREEK (IconName): The :code:`alphabet-greek` icon.\n ALPHABET_LATIN (IconName): The :code:`alphabet-latin` icon.\n AMBULANCE (IconName): The :code:`ambulance` icon.\n AMPERSAND (IconName): The :code:`ampersand` icon.\n ANALYZE (IconName): The :code:`analyze` icon.\n ANALYZE_FILLED (IconName): The :code:`analyze-filled` icon.\n ANALYZE_OFF (IconName): The :code:`analyze-off` icon.\n ANCHOR (IconName): The :code:`anchor` icon.\n ANCHOR_OFF (IconName): The :code:`anchor-off` icon.\n ANGLE (IconName): The :code:`angle` icon.\n ANKH (IconName): The :code:`ankh` icon.\n ANTENNA (IconName): The :code:`antenna` icon.\n ANTENNA_BARS_1 (IconName): The :code:`antenna-bars-1` icon.\n ANTENNA_BARS_2 (IconName): The :code:`antenna-bars-2` icon.\n ANTENNA_BARS_3 (IconName): The :code:`antenna-bars-3` icon.\n ANTENNA_BARS_4 (IconName): The :code:`antenna-bars-4` icon.\n ANTENNA_BARS_5 (IconName): The :code:`antenna-bars-5` icon.\n ANTENNA_BARS_OFF (IconName): The :code:`antenna-bars-off` icon.\n ANTENNA_OFF (IconName): The :code:`antenna-off` icon.\n APERTURE (IconName): The :code:`aperture` icon.\n APERTURE_OFF (IconName): The :code:`aperture-off` icon.\n API (IconName): The :code:`api` icon.\n API_APP (IconName): The :code:`api-app` icon.\n API_APP_OFF (IconName): The :code:`api-app-off` icon.\n API_OFF (IconName): The :code:`api-off` icon.\n APP_WINDOW (IconName): The :code:`app-window` icon.\n APP_WINDOW_FILLED (IconName): The :code:`app-window-filled` icon.\n APPLE (IconName): The :code:`apple` icon.\n APPS (IconName): The :code:`apps` icon.\n APPS_FILLED (IconName): The :code:`apps-filled` icon.\n APPS_OFF (IconName): The :code:`apps-off` icon.\n ARCHIVE (IconName): The :code:`archive` icon.\n ARCHIVE_FILLED (IconName): The :code:`archive-filled` icon.\n ARCHIVE_OFF (IconName): The :code:`archive-off` icon.\n ARMCHAIR (IconName): The :code:`armchair` icon.\n ARMCHAIR_2 (IconName): The :code:`armchair-2` icon.\n ARMCHAIR_2_OFF (IconName): The :code:`armchair-2-off` icon.\n ARMCHAIR_OFF (IconName): The :code:`armchair-off` icon.\n ARROW_AUTOFIT_CONTENT (IconName): The :code:`arrow-autofit-content` icon.\n ARROW_AUTOFIT_CONTENT_FILLED (IconName): The :code:`arrow-autofit-content-filled` icon.\n ARROW_AUTOFIT_DOWN (IconName): The :code:`arrow-autofit-down` icon.\n ARROW_AUTOFIT_HEIGHT (IconName): The :code:`arrow-autofit-height` icon.\n ARROW_AUTOFIT_LEFT (IconName): The :code:`arrow-autofit-left` icon.\n ARROW_AUTOFIT_RIGHT (IconName): The :code:`arrow-autofit-right` icon.\n ARROW_AUTOFIT_UP (IconName): The :code:`arrow-autofit-up` icon.\n ARROW_AUTOFIT_WIDTH (IconName): The :code:`arrow-autofit-width` icon.\n ARROW_BACK (IconName): The :code:`arrow-back` icon.\n ARROW_BACK_UP (IconName): The :code:`arrow-back-up` icon.\n ARROW_BACK_UP_DOUBLE (IconName): The :code:`arrow-back-up-double` icon.\n ARROW_BADGE_DOWN (IconName): The :code:`arrow-badge-down` icon.\n ARROW_BADGE_DOWN_FILLED (IconName): The :code:`arrow-badge-down-filled` icon.\n ARROW_BADGE_LEFT (IconName): The :code:`arrow-badge-left` icon.\n ARROW_BADGE_LEFT_FILLED (IconName): The :code:`arrow-badge-left-filled` icon.\n ARROW_BADGE_RIGHT (IconName): The :code:`arrow-badge-right` icon.\n ARROW_BADGE_RIGHT_FILLED (IconName): The :code:`arrow-badge-right-filled` icon.\n ARROW_BADGE_UP (IconName): The :code:`arrow-badge-up` icon.\n ARROW_BADGE_UP_FILLED (IconName): The :code:`arrow-badge-up-filled` icon.\n ARROW_BAR_BOTH (IconName): The :code:`arrow-bar-both` icon.\n ARROW_BAR_DOWN (IconName): The :code:`arrow-bar-down` icon.\n ARROW_BAR_LEFT (IconName): The :code:`arrow-bar-left` icon.\n ARROW_BAR_RIGHT (IconName): The :code:`arrow-bar-right` icon.\n ARROW_BAR_TO_DOWN (IconName): The :code:`arrow-bar-to-down` icon.\n ARROW_BAR_TO_LEFT (IconName): The :code:`arrow-bar-to-left` icon.\n ARROW_BAR_TO_RIGHT (IconName): The :code:`arrow-bar-to-right` icon.\n ARROW_BAR_TO_UP (IconName): The :code:`arrow-bar-to-up` icon.\n ARROW_BAR_UP (IconName): The :code:`arrow-bar-up` icon.\n ARROW_BEAR_LEFT (IconName): The :code:`arrow-bear-left` icon.\n ARROW_BEAR_LEFT_2 (IconName): The :code:`arrow-bear-left-2` icon.\n ARROW_BEAR_RIGHT (IconName): The :code:`arrow-bear-right` icon.\n ARROW_BEAR_RIGHT_2 (IconName): The :code:`arrow-bear-right-2` icon.\n ARROW_BIG_DOWN (IconName): The :code:`arrow-big-down` icon.\n ARROW_BIG_DOWN_FILLED (IconName): The :code:`arrow-big-down-filled` icon.\n ARROW_BIG_DOWN_LINE (IconName): The :code:`arrow-big-down-line` icon.\n ARROW_BIG_DOWN_LINE_FILLED (IconName): The :code:`arrow-big-down-line-filled` icon.\n ARROW_BIG_DOWN_LINES (IconName): The :code:`arrow-big-down-lines` icon.\n ARROW_BIG_DOWN_LINES_FILLED (IconName): The :code:`arrow-big-down-lines-filled` icon.\n ARROW_BIG_LEFT (IconName): The :code:`arrow-big-left` icon.\n ARROW_BIG_LEFT_FILLED (IconName): The :code:`arrow-big-left-filled` icon.\n ARROW_BIG_LEFT_LINE (IconName): The :code:`arrow-big-left-line` icon.\n ARROW_BIG_LEFT_LINE_FILLED (IconName): The :code:`arrow-big-left-line-filled` icon.\n ARROW_BIG_LEFT_LINES (IconName): The :code:`arrow-big-left-lines` icon.\n ARROW_BIG_LEFT_LINES_FILLED (IconName): The :code:`arrow-big-left-lines-filled` icon.\n ARROW_BIG_RIGHT (IconName): The :code:`arrow-big-right` icon.\n ARROW_BIG_RIGHT_FILLED (IconName): The :code:`arrow-big-right-filled` icon.\n ARROW_BIG_RIGHT_LINE (IconName): The :code:`arrow-big-right-line` icon.\n ARROW_BIG_RIGHT_LINE_FILLED (IconName): The :code:`arrow-big-right-line-filled` icon.\n ARROW_BIG_RIGHT_LINES (IconName): The :code:`arrow-big-right-lines` icon.\n ARROW_BIG_RIGHT_LINES_FILLED (IconName): The :code:`arrow-big-right-lines-filled` icon.\n ARROW_BIG_UP (IconName): The :code:`arrow-big-up` icon.\n ARROW_BIG_UP_FILLED (IconName): The :code:`arrow-big-up-filled` icon.\n ARROW_BIG_UP_LINE (IconName): The :code:`arrow-big-up-line` icon.\n ARROW_BIG_UP_LINE_FILLED (IconName): The :code:`arrow-big-up-line-filled` icon.\n ARROW_BIG_UP_LINES (IconName): The :code:`arrow-big-up-lines` icon.\n ARROW_BIG_UP_LINES_FILLED (IconName): The :code:`arrow-big-up-lines-filled` icon.\n ARROW_BOUNCE (IconName): The :code:`arrow-bounce` icon.\n ARROW_CAPSULE (IconName): The :code:`arrow-capsule` icon.\n ARROW_CURVE_LEFT (IconName): The :code:`arrow-curve-left` icon.\n ARROW_CURVE_RIGHT (IconName): The :code:`arrow-curve-right` icon.\n ARROW_DOWN (IconName): The :code:`arrow-down` icon.\n ARROW_DOWN_BAR (IconName): The :code:`arrow-down-bar` icon.\n ARROW_DOWN_CIRCLE (IconName): The :code:`arrow-down-circle` icon.\n ARROW_DOWN_LEFT (IconName): The :code:`arrow-down-left` icon.\n ARROW_DOWN_LEFT_CIRCLE (IconName): The :code:`arrow-down-left-circle` icon.\n ARROW_DOWN_RHOMBUS (IconName): The :code:`arrow-down-rhombus` icon.\n ARROW_DOWN_RIGHT (IconName): The :code:`arrow-down-right` icon.\n ARROW_DOWN_RIGHT_CIRCLE (IconName): The :code:`arrow-down-right-circle` icon.\n ARROW_DOWN_SQUARE (IconName): The :code:`arrow-down-square` icon.\n ARROW_DOWN_TAIL (IconName): The :code:`arrow-down-tail` icon.\n ARROW_ELBOW_LEFT (IconName): The :code:`arrow-elbow-left` icon.\n ARROW_ELBOW_RIGHT (IconName): The :code:`arrow-elbow-right` icon.\n ARROW_FORK (IconName): The :code:`arrow-fork` icon.\n ARROW_FORWARD (IconName): The :code:`arrow-forward` icon.\n ARROW_FORWARD_UP (IconName): The :code:`arrow-forward-up` icon.\n ARROW_FORWARD_UP_DOUBLE (IconName): The :code:`arrow-forward-up-double` icon.\n ARROW_GUIDE (IconName): The :code:`arrow-guide` icon.\n ARROW_ITERATION (IconName): The :code:`arrow-iteration` icon.\n ARROW_LEFT (IconName): The :code:`arrow-left` icon.\n ARROW_LEFT_BAR (IconName): The :code:`arrow-left-bar` icon.\n ARROW_LEFT_CIRCLE (IconName): The :code:`arrow-left-circle` icon.\n ARROW_LEFT_RHOMBUS (IconName): The :code:`arrow-left-rhombus` icon.\n ARROW_LEFT_RIGHT (IconName): The :code:`arrow-left-right` icon.\n ARROW_LEFT_SQUARE (IconName): The :code:`arrow-left-square` icon.\n ARROW_LEFT_TAIL (IconName): The :code:`arrow-left-tail` icon.\n ARROW_LOOP_LEFT (IconName): The :code:`arrow-loop-left` icon.\n ARROW_LOOP_LEFT_2 (IconName): The :code:`arrow-loop-left-2` icon.\n ARROW_LOOP_RIGHT (IconName): The :code:`arrow-loop-right` icon.\n ARROW_LOOP_RIGHT_2 (IconName): The :code:`arrow-loop-right-2` icon.\n ARROW_MERGE (IconName): The :code:`arrow-merge` icon.\n ARROW_MERGE_BOTH (IconName): The :code:`arrow-merge-both` icon.\n ARROW_MERGE_LEFT (IconName): The :code:`arrow-merge-left` icon.\n ARROW_MERGE_RIGHT (IconName): The :code:`arrow-merge-right` icon.\n ARROW_MOVE_DOWN (IconName): The :code:`arrow-move-down` icon.\n ARROW_MOVE_LEFT (IconName): The :code:`arrow-move-left` icon.\n ARROW_MOVE_RIGHT (IconName): The :code:`arrow-move-right` icon.\n ARROW_MOVE_UP (IconName): The :code:`arrow-move-up` icon.\n ARROW_NARROW_DOWN (IconName): The :code:`arrow-narrow-down` icon.\n ARROW_NARROW_LEFT (IconName): The :code:`arrow-narrow-left` icon.\n ARROW_NARROW_RIGHT (IconName): The :code:`arrow-narrow-right` icon.\n ARROW_NARROW_UP (IconName): The :code:`arrow-narrow-up` icon.\n ARROW_RAMP_LEFT (IconName): The :code:`arrow-ramp-left` icon.\n ARROW_RAMP_LEFT_2 (IconName): The :code:`arrow-ramp-left-2` icon.\n ARROW_RAMP_LEFT_3 (IconName): The :code:`arrow-ramp-left-3` icon.\n ARROW_RAMP_RIGHT (IconName): The :code:`arrow-ramp-right` icon.\n ARROW_RAMP_RIGHT_2 (IconName): The :code:`arrow-ramp-right-2` icon.\n ARROW_RAMP_RIGHT_3 (IconName): The :code:`arrow-ramp-right-3` icon.\n ARROW_RIGHT (IconName): The :code:`arrow-right` icon.\n ARROW_RIGHT_BAR (IconName): The :code:`arrow-right-bar` icon.\n ARROW_RIGHT_CIRCLE (IconName): The :code:`arrow-right-circle` icon.\n ARROW_RIGHT_RHOMBUS (IconName): The :code:`arrow-right-rhombus` icon.\n ARROW_RIGHT_SQUARE (IconName): The :code:`arrow-right-square` icon.\n ARROW_RIGHT_TAIL (IconName): The :code:`arrow-right-tail` icon.\n ARROW_ROTARY_FIRST_LEFT (IconName): The :code:`arrow-rotary-first-left` icon.\n ARROW_ROTARY_FIRST_RIGHT (IconName): The :code:`arrow-rotary-first-right` icon.\n ARROW_ROTARY_LAST_LEFT (IconName): The :code:`arrow-rotary-last-left` icon.\n ARROW_ROTARY_LAST_RIGHT (IconName): The :code:`arrow-rotary-last-right` icon.\n ARROW_ROTARY_LEFT (IconName): The :code:`arrow-rotary-left` icon.\n ARROW_ROTARY_RIGHT (IconName): The :code:`arrow-rotary-right` icon.\n ARROW_ROTARY_STRAIGHT (IconName): The :code:`arrow-rotary-straight` icon.\n ARROW_ROUNDABOUT_LEFT (IconName): The :code:`arrow-roundabout-left` icon.\n ARROW_ROUNDABOUT_RIGHT (IconName): The :code:`arrow-roundabout-right` icon.\n ARROW_SHARP_TURN_LEFT (IconName): The :code:`arrow-sharp-turn-left` icon.\n ARROW_SHARP_TURN_RIGHT (IconName): The :code:`arrow-sharp-turn-right` icon.\n ARROW_UP (IconName): The :code:`arrow-up` icon.\n ARROW_UP_BAR (IconName): The :code:`arrow-up-bar` icon.\n ARROW_UP_CIRCLE (IconName): The :code:`arrow-up-circle` icon.\n ARROW_UP_LEFT (IconName): The :code:`arrow-up-left` icon.\n ARROW_UP_LEFT_CIRCLE (IconName): The :code:`arrow-up-left-circle` icon.\n ARROW_UP_RHOMBUS (IconName): The :code:`arrow-up-rhombus` icon.\n ARROW_UP_RIGHT (IconName): The :code:`arrow-up-right` icon.\n ARROW_UP_RIGHT_CIRCLE (IconName): The :code:`arrow-up-right-circle` icon.\n ARROW_UP_SQUARE (IconName): The :code:`arrow-up-square` icon.\n ARROW_UP_TAIL (IconName): The :code:`arrow-up-tail` icon.\n ARROW_WAVE_LEFT_DOWN (IconName): The :code:`arrow-wave-left-down` icon.\n ARROW_WAVE_LEFT_UP (IconName): The :code:`arrow-wave-left-up` icon.\n ARROW_WAVE_RIGHT_DOWN (IconName): The :code:`arrow-wave-right-down` icon.\n ARROW_WAVE_RIGHT_UP (IconName): The :code:`arrow-wave-right-up` icon.\n ARROW_ZIG_ZAG (IconName): The :code:`arrow-zig-zag` icon.\n ARROWS_CROSS (IconName): The :code:`arrows-cross` icon.\n ARROWS_DIAGONAL (IconName): The :code:`arrows-diagonal` icon.\n ARROWS_DIAGONAL_2 (IconName): The :code:`arrows-diagonal-2` icon.\n ARROWS_DIAGONAL_MINIMIZE (IconName): The :code:`arrows-diagonal-minimize` icon.\n ARROWS_DIAGONAL_MINIMIZE_2 (IconName): The :code:`arrows-diagonal-minimize-2` icon.\n ARROWS_DIFF (IconName): The :code:`arrows-diff` icon.\n ARROWS_DOUBLE_NE_SW (IconName): The :code:`arrows-double-ne-sw` icon.\n ARROWS_DOUBLE_NW_SE (IconName): The :code:`arrows-double-nw-se` icon.\n ARROWS_DOUBLE_SE_NW (IconName): The :code:`arrows-double-se-nw` icon.\n ARROWS_DOUBLE_SW_NE (IconName): The :code:`arrows-double-sw-ne` icon.\n ARROWS_DOWN (IconName): The :code:`arrows-down` icon.\n ARROWS_DOWN_UP (IconName): The :code:`arrows-down-up` icon.\n ARROWS_EXCHANGE (IconName): The :code:`arrows-exchange` icon.\n ARROWS_EXCHANGE_2 (IconName): The :code:`arrows-exchange-2` icon.\n ARROWS_HORIZONTAL (IconName): The :code:`arrows-horizontal` icon.\n ARROWS_JOIN (IconName): The :code:`arrows-join` icon.\n ARROWS_JOIN_2 (IconName): The :code:`arrows-join-2` icon.\n ARROWS_LEFT (IconName): The :code:`arrows-left` icon.\n ARROWS_LEFT_DOWN (IconName): The :code:`arrows-left-down` icon.\n ARROWS_LEFT_RIGHT (IconName): The :code:`arrows-left-right` icon.\n ARROWS_MAXIMIZE (IconName): The :code:`arrows-maximize` icon.\n ARROWS_MINIMIZE (IconName): The :code:`arrows-minimize` icon.\n ARROWS_MOVE (IconName): The :code:`arrows-move` icon.\n ARROWS_MOVE_HORIZONTAL (IconName): The :code:`arrows-move-horizontal` icon.\n ARROWS_MOVE_VERTICAL (IconName): The :code:`arrows-move-vertical` icon.\n ARROWS_RANDOM (IconName): The :code:`arrows-random` icon.\n ARROWS_RIGHT (IconName): The :code:`arrows-right` icon.\n ARROWS_RIGHT_DOWN (IconName): The :code:`arrows-right-down` icon.\n ARROWS_RIGHT_LEFT (IconName): The :code:`arrows-right-left` icon.\n ARROWS_SHUFFLE (IconName): The :code:`arrows-shuffle` icon.\n ARROWS_SHUFFLE_2 (IconName): The :code:`arrows-shuffle-2` icon.\n ARROWS_SORT (IconName): The :code:`arrows-sort` icon.\n ARROWS_SPLIT (IconName): The :code:`arrows-split` icon.\n ARROWS_SPLIT_2 (IconName): The :code:`arrows-split-2` icon.\n ARROWS_TRANSFER_DOWN (IconName): The :code:`arrows-transfer-down` icon.\n ARROWS_TRANSFER_UP (IconName): The :code:`arrows-transfer-up` icon.\n ARROWS_UP (IconName): The :code:`arrows-up` icon.\n ARROWS_UP_DOWN (IconName): The :code:`arrows-up-down` icon.\n ARROWS_UP_LEFT (IconName): The :code:`arrows-up-left` icon.\n ARROWS_UP_RIGHT (IconName): The :code:`arrows-up-right` icon.\n ARROWS_VERTICAL (IconName): The :code:`arrows-vertical` icon.\n ARTBOARD (IconName): The :code:`artboard` icon.\n ARTBOARD_FILLED (IconName): The :code:`artboard-filled` icon.\n ARTBOARD_OFF (IconName): The :code:`artboard-off` icon.\n ARTICLE (IconName): The :code:`article` icon.\n ARTICLE_FILLED_FILLED (IconName): The :code:`article-filled-filled` icon.\n ARTICLE_OFF (IconName): The :code:`article-off` icon.\n ASPECT_RATIO (IconName): The :code:`aspect-ratio` icon.\n ASPECT_RATIO_FILLED (IconName): The :code:`aspect-ratio-filled` icon.\n ASPECT_RATIO_OFF (IconName): The :code:`aspect-ratio-off` icon.\n ASSEMBLY (IconName): The :code:`assembly` icon.\n ASSEMBLY_OFF (IconName): The :code:`assembly-off` icon.\n ASSET (IconName): The :code:`asset` icon.\n ASTERISK (IconName): The :code:`asterisk` icon.\n ASTERISK_SIMPLE (IconName): The :code:`asterisk-simple` icon.\n AT (IconName): The :code:`at` icon.\n AT_OFF (IconName): The :code:`at-off` icon.\n ATOM (IconName): The :code:`atom` icon.\n ATOM_2 (IconName): The :code:`atom-2` icon.\n ATOM_2_FILLED (IconName): The :code:`atom-2-filled` icon.\n ATOM_OFF (IconName): The :code:`atom-off` icon.\n AUGMENTED_REALITY (IconName): The :code:`augmented-reality` icon.\n AUGMENTED_REALITY_2 (IconName): The :code:`augmented-reality-2` icon.\n AUGMENTED_REALITY_OFF (IconName): The :code:`augmented-reality-off` icon.\n AWARD (IconName): The :code:`award` icon.\n AWARD_FILLED (IconName): The :code:`award-filled` icon.\n AWARD_OFF (IconName): The :code:`award-off` icon.\n AXE (IconName): The :code:`axe` icon.\n AXIS_X (IconName): The :code:`axis-x` icon.\n AXIS_Y (IconName): The :code:`axis-y` icon.\n BABY_BOTTLE (IconName): The :code:`baby-bottle` icon.\n BABY_CARRIAGE (IconName): The :code:`baby-carriage` icon.\n BACKHOE (IconName): The :code:`backhoe` icon.\n BACKPACK (IconName): The :code:`backpack` icon.\n BACKPACK_OFF (IconName): The :code:`backpack-off` icon.\n BACKSLASH (IconName): The :code:`backslash` icon.\n BACKSPACE (IconName): The :code:`backspace` icon.\n BACKSPACE_FILLED (IconName): The :code:`backspace-filled` icon.\n BADGE (IconName): The :code:`badge` icon.\n BADGE_3D (IconName): The :code:`badge-3d` icon.\n BADGE_4K (IconName): The :code:`badge-4k` icon.\n BADGE_8K (IconName): The :code:`badge-8k` icon.\n BADGE_AD (IconName): The :code:`badge-ad` icon.\n BADGE_AR (IconName): The :code:`badge-ar` icon.\n BADGE_CC (IconName): The :code:`badge-cc` icon.\n BADGE_FILLED (IconName): The :code:`badge-filled` icon.\n BADGE_HD (IconName): The :code:`badge-hd` icon.\n BADGE_OFF (IconName): The :code:`badge-off` icon.\n BADGE_SD (IconName): The :code:`badge-sd` icon.\n BADGE_TM (IconName): The :code:`badge-tm` icon.\n BADGE_VO (IconName): The :code:`badge-vo` icon.\n BADGE_VR (IconName): The :code:`badge-vr` icon.\n BADGE_WC (IconName): The :code:`badge-wc` icon.\n BADGES (IconName): The :code:`badges` icon.\n BADGES_FILLED (IconName): The :code:`badges-filled` icon.\n BADGES_OFF (IconName): The :code:`badges-off` icon.\n BAGUETTE (IconName): The :code:`baguette` icon.\n BALL_AMERICAN_FOOTBALL (IconName): The :code:`ball-american-football` icon.\n BALL_AMERICAN_FOOTBALL_OFF (IconName): The :code:`ball-american-football-off` icon.\n BALL_BASEBALL (IconName): The :code:`ball-baseball` icon.\n BALL_BASKETBALL (IconName): The :code:`ball-basketball` icon.\n BALL_BOWLING (IconName): The :code:`ball-bowling` icon.\n BALL_FOOTBALL (IconName): The :code:`ball-football` icon.\n BALL_FOOTBALL_OFF (IconName): The :code:`ball-football-off` icon.\n BALL_TENNIS (IconName): The :code:`ball-tennis` icon.\n BALL_VOLLEYBALL (IconName): The :code:`ball-volleyball` icon.\n BALLOON (IconName): The :code:`balloon` icon.\n BALLOON_FILLED (IconName): The :code:`balloon-filled` icon.\n BALLOON_OFF (IconName): The :code:`balloon-off` icon.\n BALLPEN (IconName): The :code:`ballpen` icon.\n BALLPEN_FILLED (IconName): The :code:`ballpen-filled` icon.\n BALLPEN_OFF (IconName): The :code:`ballpen-off` icon.\n BAN (IconName): The :code:`ban` icon.\n BANDAGE (IconName): The :code:`bandage` icon.\n BANDAGE_FILLED (IconName): The :code:`bandage-filled` icon.\n BANDAGE_OFF (IconName): The :code:`bandage-off` icon.\n BARBELL (IconName): The :code:`barbell` icon.\n BARBELL_OFF (IconName): The :code:`barbell-off` icon.\n BARCODE (IconName): The :code:`barcode` icon.\n BARCODE_OFF (IconName): The :code:`barcode-off` icon.\n BARREL (IconName): The :code:`barrel` icon.\n BARREL_OFF (IconName): The :code:`barrel-off` icon.\n BARRIER_BLOCK (IconName): The :code:`barrier-block` icon.\n BARRIER_BLOCK_OFF (IconName): The :code:`barrier-block-off` icon.\n BASELINE (IconName): The :code:`baseline` icon.\n BASELINE_DENSITY_LARGE (IconName): The :code:`baseline-density-large` icon.\n BASELINE_DENSITY_MEDIUM (IconName): The :code:`baseline-density-medium` icon.\n BASELINE_DENSITY_SMALL (IconName): The :code:`baseline-density-small` icon.\n BASKET (IconName): The :code:`basket` icon.\n BASKET_FILLED (IconName): The :code:`basket-filled` icon.\n BASKET_OFF (IconName): The :code:`basket-off` icon.\n BAT (IconName): The :code:`bat` icon.\n BATH (IconName): The :code:`bath` icon.\n BATH_FILLED (IconName): The :code:`bath-filled` icon.\n BATH_OFF (IconName): The :code:`bath-off` icon.\n BATTERY (IconName): The :code:`battery` icon.\n BATTERY_1 (IconName): The :code:`battery-1` icon.\n BATTERY_1_FILLED (IconName): The :code:`battery-1-filled` icon.\n BATTERY_2 (IconName): The :code:`battery-2` icon.\n BATTERY_2_FILLED (IconName): The :code:`battery-2-filled` icon.\n BATTERY_3 (IconName): The :code:`battery-3` icon.\n BATTERY_3_FILLED (IconName): The :code:`battery-3-filled` icon.\n BATTERY_4 (IconName): The :code:`battery-4` icon.\n BATTERY_4_FILLED (IconName): The :code:`battery-4-filled` icon.\n BATTERY_AUTOMOTIVE (IconName): The :code:`battery-automotive` icon.\n BATTERY_CHARGING (IconName): The :code:`battery-charging` icon.\n BATTERY_CHARGING_2 (IconName): The :code:`battery-charging-2` icon.\n BATTERY_ECO (IconName): The :code:`battery-eco` icon.\n BATTERY_FILLED (IconName): The :code:`battery-filled` icon.\n BATTERY_OFF (IconName): The :code:`battery-off` icon.\n BEACH (IconName): The :code:`beach` icon.\n BEACH_OFF (IconName): The :code:`beach-off` icon.\n BED (IconName): The :code:`bed` icon.\n BED_FILLED (IconName): The :code:`bed-filled` icon.\n BED_OFF (IconName): The :code:`bed-off` icon.\n BEER (IconName): The :code:`beer` icon.\n BEER_FILLED (IconName): The :code:`beer-filled` icon.\n BEER_OFF (IconName): The :code:`beer-off` icon.\n BELL (IconName): The :code:`bell` icon.\n BELL_BOLT (IconName): The :code:`bell-bolt` icon.\n BELL_CANCEL (IconName): The :code:`bell-cancel` icon.\n BELL_CHECK (IconName): The :code:`bell-check` icon.\n BELL_CODE (IconName): The :code:`bell-code` icon.\n BELL_COG (IconName): The :code:`bell-cog` icon.\n BELL_DOLLAR (IconName): The :code:`bell-dollar` icon.\n BELL_DOWN (IconName): The :code:`bell-down` icon.\n BELL_EXCLAMATION (IconName): The :code:`bell-exclamation` icon.\n BELL_FILLED (IconName): The :code:`bell-filled` icon.\n BELL_HEART (IconName): The :code:`bell-heart` icon.\n BELL_MINUS (IconName): The :code:`bell-minus` icon.\n BELL_MINUS_FILLED (IconName): The :code:`bell-minus-filled` icon.\n BELL_OFF (IconName): The :code:`bell-off` icon.\n BELL_PAUSE (IconName): The :code:`bell-pause` icon.\n BELL_PIN (IconName): The :code:`bell-pin` icon.\n BELL_PLUS (IconName): The :code:`bell-plus` icon.\n BELL_PLUS_FILLED (IconName): The :code:`bell-plus-filled` icon.\n BELL_QUESTION (IconName): The :code:`bell-question` icon.\n BELL_RINGING (IconName): The :code:`bell-ringing` icon.\n BELL_RINGING_2 (IconName): The :code:`bell-ringing-2` icon.\n BELL_RINGING_2_FILLED (IconName): The :code:`bell-ringing-2-filled` icon.\n BELL_RINGING_FILLED (IconName): The :code:`bell-ringing-filled` icon.\n BELL_SCHOOL (IconName): The :code:`bell-school` icon.\n BELL_SEARCH (IconName): The :code:`bell-search` icon.\n BELL_SHARE (IconName): The :code:`bell-share` icon.\n BELL_STAR (IconName): The :code:`bell-star` icon.\n BELL_UP (IconName): The :code:`bell-up` icon.\n BELL_X (IconName): The :code:`bell-x` icon.\n BELL_X_FILLED (IconName): The :code:`bell-x-filled` icon.\n BELL_Z (IconName): The :code:`bell-z` icon.\n BELL_Z_FILLED (IconName): The :code:`bell-z-filled` icon.\n BETA (IconName): The :code:`beta` icon.\n BIBLE (IconName): The :code:`bible` icon.\n BIKE (IconName): The :code:`bike` icon.\n BIKE_OFF (IconName): The :code:`bike-off` icon.\n BINARY (IconName): The :code:`binary` icon.\n BINARY_OFF (IconName): The :code:`binary-off` icon.\n BINARY_TREE (IconName): The :code:`binary-tree` icon.\n BINARY_TREE_2 (IconName): The :code:`binary-tree-2` icon.\n BIOHAZARD (IconName): The :code:`biohazard` icon.\n BIOHAZARD_OFF (IconName): The :code:`biohazard-off` icon.\n BLADE (IconName): The :code:`blade` icon.\n BLADE_FILLED (IconName): The :code:`blade-filled` icon.\n BLEACH (IconName): The :code:`bleach` icon.\n BLEACH_CHLORINE (IconName): The :code:`bleach-chlorine` icon.\n BLEACH_NO_CHLORINE (IconName): The :code:`bleach-no-chlorine` icon.\n BLEACH_OFF (IconName): The :code:`bleach-off` icon.\n BLOCKQUOTE (IconName): The :code:`blockquote` icon.\n BLUETOOTH (IconName): The :code:`bluetooth` icon.\n BLUETOOTH_CONNECTED (IconName): The :code:`bluetooth-connected` icon.\n BLUETOOTH_OFF (IconName): The :code:`bluetooth-off` icon.\n BLUETOOTH_X (IconName): The :code:`bluetooth-x` icon.\n BLUR (IconName): The :code:`blur` icon.\n BLUR_OFF (IconName): The :code:`blur-off` icon.\n BMP (IconName): The :code:`bmp` icon.\n BOLD (IconName): The :code:`bold` icon.\n BOLD_OFF (IconName): The :code:`bold-off` icon.\n BOLT (IconName): The :code:`bolt` icon.\n BOLT_OFF (IconName): The :code:`bolt-off` icon.\n BOMB (IconName): The :code:`bomb` icon.\n BOMB_FILLED (IconName): The :code:`bomb-filled` icon.\n BONE (IconName): The :code:`bone` icon.\n BONE_OFF (IconName): The :code:`bone-off` icon.\n BONG (IconName): The :code:`bong` icon.\n BONG_OFF (IconName): The :code:`bong-off` icon.\n BOOK (IconName): The :code:`book` icon.\n BOOK_2 (IconName): The :code:`book-2` icon.\n BOOK_DOWNLOAD (IconName): The :code:`book-download` icon.\n BOOK_FILLED (IconName): The :code:`book-filled` icon.\n BOOK_OFF (IconName): The :code:`book-off` icon.\n BOOK_UPLOAD (IconName): The :code:`book-upload` icon.\n BOOKMARK (IconName): The :code:`bookmark` icon.\n BOOKMARK_EDIT (IconName): The :code:`bookmark-edit` icon.\n BOOKMARK_FILLED (IconName): The :code:`bookmark-filled` icon.\n BOOKMARK_MINUS (IconName): The :code:`bookmark-minus` icon.\n BOOKMARK_OFF (IconName): The :code:`bookmark-off` icon.\n BOOKMARK_PLUS (IconName): The :code:`bookmark-plus` icon.\n BOOKMARK_QUESTION (IconName): The :code:`bookmark-question` icon.\n BOOKMARKS (IconName): The :code:`bookmarks` icon.\n BOOKMARKS_OFF (IconName): The :code:`bookmarks-off` icon.\n BOOKS (IconName): The :code:`books` icon.\n BOOKS_OFF (IconName): The :code:`books-off` icon.\n BORDER_ALL (IconName): The :code:`border-all` icon.\n BORDER_BOTTOM (IconName): The :code:`border-bottom` icon.\n BORDER_CORNERS (IconName): The :code:`border-corners` icon.\n BORDER_HORIZONTAL (IconName): The :code:`border-horizontal` icon.\n BORDER_INNER (IconName): The :code:`border-inner` icon.\n BORDER_LEFT (IconName): The :code:`border-left` icon.\n BORDER_NONE (IconName): The :code:`border-none` icon.\n BORDER_OUTER (IconName): The :code:`border-outer` icon.\n BORDER_RADIUS (IconName): The :code:`border-radius` icon.\n BORDER_RIGHT (IconName): The :code:`border-right` icon.\n BORDER_SIDES (IconName): The :code:`border-sides` icon.\n BORDER_STYLE (IconName): The :code:`border-style` icon.\n BORDER_STYLE_2 (IconName): The :code:`border-style-2` icon.\n BORDER_TOP (IconName): The :code:`border-top` icon.\n BORDER_VERTICAL (IconName): The :code:`border-vertical` icon.\n BOTTLE (IconName): The :code:`bottle` icon.\n BOTTLE_FILLED (IconName): The :code:`bottle-filled` icon.\n BOTTLE_OFF (IconName): The :code:`bottle-off` icon.\n BOUNCE_LEFT (IconName): The :code:`bounce-left` icon.\n BOUNCE_RIGHT (IconName): The :code:`bounce-right` icon.\n BOW (IconName): The :code:`bow` icon.\n BOWL (IconName): The :code:`bowl` icon.\n BOX (IconName): The :code:`box` icon.\n BOX_ALIGN_BOTTOM (IconName): The :code:`box-align-bottom` icon.\n BOX_ALIGN_BOTTOM_FILLED (IconName): The :code:`box-align-bottom-filled` icon.\n BOX_ALIGN_BOTTOM_LEFT (IconName): The :code:`box-align-bottom-left` icon.\n BOX_ALIGN_BOTTOM_LEFT_FILLED (IconName): The :code:`box-align-bottom-left-filled` icon.\n BOX_ALIGN_BOTTOM_RIGHT (IconName): The :code:`box-align-bottom-right` icon.\n BOX_ALIGN_BOTTOM_RIGHT_FILLED (IconName): The :code:`box-align-bottom-right-filled` icon.\n BOX_ALIGN_LEFT (IconName): The :code:`box-align-left` icon.\n BOX_ALIGN_LEFT_FILLED (IconName): The :code:`box-align-left-filled` icon.\n BOX_ALIGN_RIGHT (IconName): The :code:`box-align-right` icon.\n BOX_ALIGN_RIGHT_FILLED (IconName): The :code:`box-align-right-filled` icon.\n BOX_ALIGN_TOP (IconName): The :code:`box-align-top` icon.\n BOX_ALIGN_TOP_FILLED (IconName): The :code:`box-align-top-filled` icon.\n BOX_ALIGN_TOP_LEFT (IconName): The :code:`box-align-top-left` icon.\n BOX_ALIGN_TOP_LEFT_FILLED (IconName): The :code:`box-align-top-left-filled` icon.\n BOX_ALIGN_TOP_RIGHT (IconName): The :code:`box-align-top-right` icon.\n BOX_ALIGN_TOP_RIGHT_FILLED (IconName): The :code:`box-align-top-right-filled` icon.\n BOX_MARGIN (IconName): The :code:`box-margin` icon.\n BOX_MODEL (IconName): The :code:`box-model` icon.\n BOX_MODEL_2 (IconName): The :code:`box-model-2` icon.\n BOX_MODEL_2_OFF (IconName): The :code:`box-model-2-off` icon.\n BOX_MODEL_OFF (IconName): The :code:`box-model-off` icon.\n BOX_MULTIPLE (IconName): The :code:`box-multiple` icon.\n BOX_MULTIPLE_0 (IconName): The :code:`box-multiple-0` icon.\n BOX_MULTIPLE_1 (IconName): The :code:`box-multiple-1` icon.\n BOX_MULTIPLE_2 (IconName): The :code:`box-multiple-2` icon.\n BOX_MULTIPLE_3 (IconName): The :code:`box-multiple-3` icon.\n BOX_MULTIPLE_4 (IconName): The :code:`box-multiple-4` icon.\n BOX_MULTIPLE_5 (IconName): The :code:`box-multiple-5` icon.\n BOX_MULTIPLE_6 (IconName): The :code:`box-multiple-6` icon.\n BOX_MULTIPLE_7 (IconName): The :code:`box-multiple-7` icon.\n BOX_MULTIPLE_8 (IconName): The :code:`box-multiple-8` icon.\n BOX_MULTIPLE_9 (IconName): The :code:`box-multiple-9` icon.\n BOX_OFF (IconName): The :code:`box-off` icon.\n BOX_PADDING (IconName): The :code:`box-padding` icon.\n BOX_SEAM (IconName): The :code:`box-seam` icon.\n BRACES (IconName): The :code:`braces` icon.\n BRACES_OFF (IconName): The :code:`braces-off` icon.\n BRACKETS (IconName): The :code:`brackets` icon.\n BRACKETS_CONTAIN (IconName): The :code:`brackets-contain` icon.\n BRACKETS_CONTAIN_END (IconName): The :code:`brackets-contain-end` icon.\n BRACKETS_CONTAIN_START (IconName): The :code:`brackets-contain-start` icon.\n BRACKETS_OFF (IconName): The :code:`brackets-off` icon.\n BRAILLE (IconName): The :code:`braille` icon.\n BRAIN (IconName): The :code:`brain` icon.\n BRAND_4CHAN (IconName): The :code:`brand-4chan` icon.\n BRAND_ABSTRACT (IconName): The :code:`brand-abstract` icon.\n BRAND_ADOBE (IconName): The :code:`brand-adobe` icon.\n BRAND_ADONIS_JS (IconName): The :code:`brand-adonis-js` icon.\n BRAND_AIRBNB (IconName): The :code:`brand-airbnb` icon.\n BRAND_AIRTABLE (IconName): The :code:`brand-airtable` icon.\n BRAND_ALGOLIA (IconName): The :code:`brand-algolia` icon.\n BRAND_ALIPAY (IconName): The :code:`brand-alipay` icon.\n BRAND_ALPINE_JS (IconName): The :code:`brand-alpine-js` icon.\n BRAND_AMAZON (IconName): The :code:`brand-amazon` icon.\n BRAND_AMD (IconName): The :code:`brand-amd` icon.\n BRAND_AMIGO (IconName): The :code:`brand-amigo` icon.\n BRAND_AMONG_US (IconName): The :code:`brand-among-us` icon.\n BRAND_ANDROID (IconName): The :code:`brand-android` icon.\n BRAND_ANGULAR (IconName): The :code:`brand-angular` icon.\n BRAND_ANSIBLE (IconName): The :code:`brand-ansible` icon.\n BRAND_AO3 (IconName): The :code:`brand-ao3` icon.\n BRAND_APPGALLERY (IconName): The :code:`brand-appgallery` icon.\n BRAND_APPLE (IconName): The :code:`brand-apple` icon.\n BRAND_APPLE_ARCADE (IconName): The :code:`brand-apple-arcade` icon.\n BRAND_APPLE_PODCAST (IconName): The :code:`brand-apple-podcast` icon.\n BRAND_APPSTORE (IconName): The :code:`brand-appstore` icon.\n BRAND_ASANA (IconName): The :code:`brand-asana` icon.\n BRAND_AWS (IconName): The :code:`brand-aws` icon.\n BRAND_AZURE (IconName): The :code:`brand-azure` icon.\n BRAND_BACKBONE (IconName): The :code:`brand-backbone` icon.\n BRAND_BADOO (IconName): The :code:`brand-badoo` icon.\n BRAND_BAIDU (IconName): The :code:`brand-baidu` icon.\n BRAND_BANDCAMP (IconName): The :code:`brand-bandcamp` icon.\n BRAND_BANDLAB (IconName): The :code:`brand-bandlab` icon.\n BRAND_BEATS (IconName): The :code:`brand-beats` icon.\n BRAND_BEHANCE (IconName): The :code:`brand-behance` icon.\n BRAND_BILIBILI (IconName): The :code:`brand-bilibili` icon.\n BRAND_BINANCE (IconName): The :code:`brand-binance` icon.\n BRAND_BING (IconName): The :code:`brand-bing` icon.\n BRAND_BITBUCKET (IconName): The :code:`brand-bitbucket` icon.\n BRAND_BLACKBERRY (IconName): The :code:`brand-blackberry` icon.\n BRAND_BLENDER (IconName): The :code:`brand-blender` icon.\n BRAND_BLOGGER (IconName): The :code:`brand-blogger` icon.\n BRAND_BOOKING (IconName): The :code:`brand-booking` icon.\n BRAND_BOOTSTRAP (IconName): The :code:`brand-bootstrap` icon.\n BRAND_BULMA (IconName): The :code:`brand-bulma` icon.\n BRAND_BUMBLE (IconName): The :code:`brand-bumble` icon.\n BRAND_BUNPO (IconName): The :code:`brand-bunpo` icon.\n BRAND_C_SHARP (IconName): The :code:`brand-c-sharp` icon.\n BRAND_CAKE (IconName): The :code:`brand-cake` icon.\n BRAND_CAKEPHP (IconName): The :code:`brand-cakephp` icon.\n BRAND_CAMPAIGNMONITOR (IconName): The :code:`brand-campaignmonitor` icon.\n BRAND_CARBON (IconName): The :code:`brand-carbon` icon.\n BRAND_CASHAPP (IconName): The :code:`brand-cashapp` icon.\n BRAND_CHROME (IconName): The :code:`brand-chrome` icon.\n BRAND_CINEMA_4D (IconName): The :code:`brand-cinema-4d` icon.\n BRAND_CITYMAPPER (IconName): The :code:`brand-citymapper` icon.\n BRAND_CLOUDFLARE (IconName): The :code:`brand-cloudflare` icon.\n BRAND_CODECOV (IconName): The :code:`brand-codecov` icon.\n BRAND_CODEPEN (IconName): The :code:`brand-codepen` icon.\n BRAND_CODESANDBOX (IconName): The :code:`brand-codesandbox` icon.\n BRAND_COHOST (IconName): The :code:`brand-cohost` icon.\n BRAND_COINBASE (IconName): The :code:`brand-coinbase` icon.\n BRAND_COMEDY_CENTRAL (IconName): The :code:`brand-comedy-central` icon.\n BRAND_COREOS (IconName): The :code:`brand-coreos` icon.\n BRAND_COUCHDB (IconName): The :code:`brand-couchdb` icon.\n BRAND_COUCHSURFING (IconName): The :code:`brand-couchsurfing` icon.\n BRAND_CPP (IconName): The :code:`brand-cpp` icon.\n BRAND_CRAFT (IconName): The :code:`brand-craft` icon.\n BRAND_CRUNCHBASE (IconName): The :code:`brand-crunchbase` icon.\n BRAND_CSS3 (IconName): The :code:`brand-css3` icon.\n BRAND_CTEMPLAR (IconName): The :code:`brand-ctemplar` icon.\n BRAND_CUCUMBER (IconName): The :code:`brand-cucumber` icon.\n BRAND_CUPRA (IconName): The :code:`brand-cupra` icon.\n BRAND_CYPRESS (IconName): The :code:`brand-cypress` icon.\n BRAND_D3 (IconName): The :code:`brand-d3` icon.\n BRAND_DAYS_COUNTER (IconName): The :code:`brand-days-counter` icon.\n BRAND_DCOS (IconName): The :code:`brand-dcos` icon.\n BRAND_DEBIAN (IconName): The :code:`brand-debian` icon.\n BRAND_DEEZER (IconName): The :code:`brand-deezer` icon.\n BRAND_DELIVEROO (IconName): The :code:`brand-deliveroo` icon.\n BRAND_DENO (IconName): The :code:`brand-deno` icon.\n BRAND_DENODO (IconName): The :code:`brand-denodo` icon.\n BRAND_DEVIANTART (IconName): The :code:`brand-deviantart` icon.\n BRAND_DIGG (IconName): The :code:`brand-digg` icon.\n BRAND_DINGTALK (IconName): The :code:`brand-dingtalk` icon.\n BRAND_DISCORD (IconName): The :code:`brand-discord` icon.\n BRAND_DISCORD_FILLED (IconName): The :code:`brand-discord-filled` icon.\n BRAND_DISNEY (IconName): The :code:`brand-disney` icon.\n BRAND_DISQUS (IconName): The :code:`brand-disqus` icon.\n BRAND_DJANGO (IconName): The :code:`brand-django` icon.\n BRAND_DOCKER (IconName): The :code:`brand-docker` icon.\n BRAND_DOCTRINE (IconName): The :code:`brand-doctrine` icon.\n BRAND_DOLBY_DIGITAL (IconName): The :code:`brand-dolby-digital` icon.\n BRAND_DOUBAN (IconName): The :code:`brand-douban` icon.\n BRAND_DRIBBBLE (IconName): The :code:`brand-dribbble` icon.\n BRAND_DRIBBBLE_FILLED (IconName): The :code:`brand-dribbble-filled` icon.\n BRAND_DROPS (IconName): The :code:`brand-drops` icon.\n BRAND_DRUPAL (IconName): The :code:`brand-drupal` icon.\n BRAND_EDGE (IconName): The :code:`brand-edge` icon.\n BRAND_ELASTIC (IconName): The :code:`brand-elastic` icon.\n BRAND_ELECTRONIC_ARTS (IconName): The :code:`brand-electronic-arts` icon.\n BRAND_EMBER (IconName): The :code:`brand-ember` icon.\n BRAND_ENVATO (IconName): The :code:`brand-envato` icon.\n BRAND_ETSY (IconName): The :code:`brand-etsy` icon.\n BRAND_EVERNOTE (IconName): The :code:`brand-evernote` icon.\n BRAND_FACEBOOK (IconName): The :code:`brand-facebook` icon.\n BRAND_FACEBOOK_FILLED (IconName): The :code:`brand-facebook-filled` icon.\n BRAND_FEEDLY (IconName): The :code:`brand-feedly` icon.\n BRAND_FIGMA (IconName): The :code:`brand-figma` icon.\n BRAND_FILEZILLA (IconName): The :code:`brand-filezilla` icon.\n BRAND_FINDER (IconName): The :code:`brand-finder` icon.\n BRAND_FIREBASE (IconName): The :code:`brand-firebase` icon.\n BRAND_FIREFOX (IconName): The :code:`brand-firefox` icon.\n BRAND_FIVERR (IconName): The :code:`brand-fiverr` icon.\n BRAND_FLICKR (IconName): The :code:`brand-flickr` icon.\n BRAND_FLIGHTRADAR24 (IconName): The :code:`brand-flightradar24` icon.\n BRAND_FLIPBOARD (IconName): The :code:`brand-flipboard` icon.\n BRAND_FLUTTER (IconName): The :code:`brand-flutter` icon.\n BRAND_FORTNITE (IconName): The :code:`brand-fortnite` icon.\n BRAND_FOURSQUARE (IconName): The :code:`brand-foursquare` icon.\n BRAND_FRAMER (IconName): The :code:`brand-framer` icon.\n BRAND_FRAMER_MOTION (IconName): The :code:`brand-framer-motion` icon.\n BRAND_FUNIMATION (IconName): The :code:`brand-funimation` icon.\n BRAND_GATSBY (IconName): The :code:`brand-gatsby` icon.\n BRAND_GIT (IconName): The :code:`brand-git` icon.\n BRAND_GITHUB (IconName): The :code:`brand-github` icon.\n BRAND_GITHUB_COPILOT (IconName): The :code:`brand-github-copilot` icon.\n BRAND_GITHUB_FILLED (IconName): The :code:`brand-github-filled` icon.\n BRAND_GITLAB (IconName): The :code:`brand-gitlab` icon.\n BRAND_GMAIL (IconName): The :code:`brand-gmail` icon.\n BRAND_GOLANG (IconName): The :code:`brand-golang` icon.\n BRAND_GOOGLE (IconName): The :code:`brand-google` icon.\n BRAND_GOOGLE_ANALYTICS (IconName): The :code:`brand-google-analytics` icon.\n BRAND_GOOGLE_BIG_QUERY (IconName): The :code:`brand-google-big-query` icon.\n BRAND_GOOGLE_DRIVE (IconName): The :code:`brand-google-drive` icon.\n BRAND_GOOGLE_FIT (IconName): The :code:`brand-google-fit` icon.\n BRAND_GOOGLE_HOME (IconName): The :code:`brand-google-home` icon.\n BRAND_GOOGLE_MAPS (IconName): The :code:`brand-google-maps` icon.\n BRAND_GOOGLE_ONE (IconName): The :code:`brand-google-one` icon.\n BRAND_GOOGLE_PHOTOS (IconName): The :code:`brand-google-photos` icon.\n BRAND_GOOGLE_PLAY (IconName): The :code:`brand-google-play` icon.\n BRAND_GOOGLE_PODCASTS (IconName): The :code:`brand-google-podcasts` icon.\n BRAND_GRAMMARLY (IconName): The :code:`brand-grammarly` icon.\n BRAND_GRAPHQL (IconName): The :code:`brand-graphql` icon.\n BRAND_GRAVATAR (IconName): The :code:`brand-gravatar` icon.\n BRAND_GRINDR (IconName): The :code:`brand-grindr` icon.\n BRAND_GUARDIAN (IconName): The :code:`brand-guardian` icon.\n BRAND_GUMROAD (IconName): The :code:`brand-gumroad` icon.\n BRAND_HBO (IconName): The :code:`brand-hbo` icon.\n BRAND_HEADLESSUI (IconName): The :code:`brand-headlessui` icon.\n BRAND_HEXO (IconName): The :code:`brand-hexo` icon.\n BRAND_HIPCHAT (IconName): The :code:`brand-hipchat` icon.\n BRAND_HTML5 (IconName): The :code:`brand-html5` icon.\n BRAND_INERTIA (IconName): The :code:`brand-inertia` icon.\n BRAND_INSTAGRAM (IconName): The :code:`brand-instagram` icon.\n BRAND_INTERCOM (IconName): The :code:`brand-intercom` icon.\n BRAND_ITCH (IconName): The :code:`brand-itch` icon.\n BRAND_JAVASCRIPT (IconName): The :code:`brand-javascript` icon.\n BRAND_JUEJIN (IconName): The :code:`brand-juejin` icon.\n BRAND_KBIN (IconName): The :code:`brand-kbin` icon.\n BRAND_KICK (IconName): The :code:`brand-kick` icon.\n BRAND_KICKSTARTER (IconName): The :code:`brand-kickstarter` icon.\n BRAND_KOTLIN (IconName): The :code:`brand-kotlin` icon.\n BRAND_LARAVEL (IconName): The :code:`brand-laravel` icon.\n BRAND_LASTFM (IconName): The :code:`brand-lastfm` icon.\n BRAND_LEETCODE (IconName): The :code:`brand-leetcode` icon.\n BRAND_LETTERBOXD (IconName): The :code:`brand-letterboxd` icon.\n BRAND_LINE (IconName): The :code:`brand-line` icon.\n BRAND_LINKEDIN (IconName): The :code:`brand-linkedin` icon.\n BRAND_LINKTREE (IconName): The :code:`brand-linktree` icon.\n BRAND_LINQPAD (IconName): The :code:`brand-linqpad` icon.\n BRAND_LOOM (IconName): The :code:`brand-loom` icon.\n BRAND_MAILGUN (IconName): The :code:`brand-mailgun` icon.\n BRAND_MANTINE (IconName): The :code:`brand-mantine` icon.\n BRAND_MASTERCARD (IconName): The :code:`brand-mastercard` icon.\n BRAND_MASTODON (IconName): The :code:`brand-mastodon` icon.\n BRAND_MATRIX (IconName): The :code:`brand-matrix` icon.\n BRAND_MCDONALDS (IconName): The :code:`brand-mcdonalds` icon.\n BRAND_MEDIUM (IconName): The :code:`brand-medium` icon.\n BRAND_MERCEDES (IconName): The :code:`brand-mercedes` icon.\n BRAND_MESSENGER (IconName): The :code:`brand-messenger` icon.\n BRAND_META (IconName): The :code:`brand-meta` icon.\n BRAND_MICROSOFT_TEAMS (IconName): The :code:`brand-microsoft-teams` icon.\n BRAND_MINECRAFT (IconName): The :code:`brand-minecraft` icon.\n BRAND_MINIPROGRAM (IconName): The :code:`brand-miniprogram` icon.\n BRAND_MIXPANEL (IconName): The :code:`brand-mixpanel` icon.\n BRAND_MONDAY (IconName): The :code:`brand-monday` icon.\n BRAND_MONGODB (IconName): The :code:`brand-mongodb` icon.\n BRAND_MY_OPPO (IconName): The :code:`brand-my-oppo` icon.\n BRAND_MYSQL (IconName): The :code:`brand-mysql` icon.\n BRAND_NATIONAL_GEOGRAPHIC (IconName): The :code:`brand-national-geographic` icon.\n BRAND_NEM (IconName): The :code:`brand-nem` icon.\n BRAND_NETBEANS (IconName): The :code:`brand-netbeans` icon.\n BRAND_NETEASE_MUSIC (IconName): The :code:`brand-netease-music` icon.\n BRAND_NETFLIX (IconName): The :code:`brand-netflix` icon.\n BRAND_NEXO (IconName): The :code:`brand-nexo` icon.\n BRAND_NEXTCLOUD (IconName): The :code:`brand-nextcloud` icon.\n BRAND_NEXTJS (IconName): The :code:`brand-nextjs` icon.\n BRAND_NODEJS (IconName): The :code:`brand-nodejs` icon.\n BRAND_NORD_VPN (IconName): The :code:`brand-nord-vpn` icon.\n BRAND_NOTION (IconName): The :code:`brand-notion` icon.\n BRAND_NPM (IconName): The :code:`brand-npm` icon.\n BRAND_NUXT (IconName): The :code:`brand-nuxt` icon.\n BRAND_NYTIMES (IconName): The :code:`brand-nytimes` icon.\n BRAND_OAUTH (IconName): The :code:`brand-oauth` icon.\n BRAND_OFFICE (IconName): The :code:`brand-office` icon.\n BRAND_OK_RU (IconName): The :code:`brand-ok-ru` icon.\n BRAND_ONEDRIVE (IconName): The :code:`brand-onedrive` icon.\n BRAND_ONLYFANS (IconName): The :code:`brand-onlyfans` icon.\n BRAND_OPEN_SOURCE (IconName): The :code:`brand-open-source` icon.\n BRAND_OPENAI (IconName): The :code:`brand-openai` icon.\n BRAND_OPENVPN (IconName): The :code:`brand-openvpn` icon.\n BRAND_OPERA (IconName): The :code:`brand-opera` icon.\n BRAND_PAGEKIT (IconName): The :code:`brand-pagekit` icon.\n BRAND_PATREON (IconName): The :code:`brand-patreon` icon.\n BRAND_PAYPAL (IconName): The :code:`brand-paypal` icon.\n BRAND_PAYPAL_FILLED (IconName): The :code:`brand-paypal-filled` icon.\n BRAND_PAYPAY (IconName): The :code:`brand-paypay` icon.\n BRAND_PEANUT (IconName): The :code:`brand-peanut` icon.\n BRAND_PEPSI (IconName): The :code:`brand-pepsi` icon.\n BRAND_PHP (IconName): The :code:`brand-php` icon.\n BRAND_PICSART (IconName): The :code:`brand-picsart` icon.\n BRAND_PINTEREST (IconName): The :code:`brand-pinterest` icon.\n BRAND_PLANETSCALE (IconName): The :code:`brand-planetscale` icon.\n BRAND_POCKET (IconName): The :code:`brand-pocket` icon.\n BRAND_POLYMER (IconName): The :code:`brand-polymer` icon.\n BRAND_POWERSHELL (IconName): The :code:`brand-powershell` icon.\n BRAND_PRISMA (IconName): The :code:`brand-prisma` icon.\n BRAND_PRODUCTHUNT (IconName): The :code:`brand-producthunt` icon.\n BRAND_PUSHBULLET (IconName): The :code:`brand-pushbullet` icon.\n BRAND_PUSHOVER (IconName): The :code:`brand-pushover` icon.\n BRAND_PYTHON (IconName): The :code:`brand-python` icon.\n BRAND_QQ (IconName): The :code:`brand-qq` icon.\n BRAND_RADIX_UI (IconName): The :code:`brand-radix-ui` icon.\n BRAND_REACT (IconName): The :code:`brand-react` icon.\n BRAND_REACT_NATIVE (IconName): The :code:`brand-react-native` icon.\n BRAND_REASON (IconName): The :code:`brand-reason` icon.\n BRAND_REDDIT (IconName): The :code:`brand-reddit` icon.\n BRAND_REDHAT (IconName): The :code:`brand-redhat` icon.\n BRAND_REDUX (IconName): The :code:`brand-redux` icon.\n BRAND_REVOLUT (IconName): The :code:`brand-revolut` icon.\n BRAND_RUMBLE (IconName): The :code:`brand-rumble` icon.\n BRAND_RUST (IconName): The :code:`brand-rust` icon.\n BRAND_SAFARI (IconName): The :code:`brand-safari` icon.\n BRAND_SAMSUNGPASS (IconName): The :code:`brand-samsungpass` icon.\n BRAND_SASS (IconName): The :code:`brand-sass` icon.\n BRAND_SENTRY (IconName): The :code:`brand-sentry` icon.\n BRAND_SHARIK (IconName): The :code:`brand-sharik` icon.\n BRAND_SHAZAM (IconName): The :code:`brand-shazam` icon.\n BRAND_SHOPEE (IconName): The :code:`brand-shopee` icon.\n BRAND_SKETCH (IconName): The :code:`brand-sketch` icon.\n BRAND_SKYPE (IconName): The :code:`brand-skype` icon.\n BRAND_SLACK (IconName): The :code:`brand-slack` icon.\n BRAND_SNAPCHAT (IconName): The :code:`brand-snapchat` icon.\n BRAND_SNAPSEED (IconName): The :code:`brand-snapseed` icon.\n BRAND_SNOWFLAKE (IconName): The :code:`brand-snowflake` icon.\n BRAND_SOCKET_IO (IconName): The :code:`brand-socket-io` icon.\n BRAND_SOLIDJS (IconName): The :code:`brand-solidjs` icon.\n BRAND_SOUNDCLOUD (IconName): The :code:`brand-soundcloud` icon.\n BRAND_SPACEHEY (IconName): The :code:`brand-spacehey` icon.\n BRAND_SPEEDTEST (IconName): The :code:`brand-speedtest` icon.\n BRAND_SPOTIFY (IconName): The :code:`brand-spotify` icon.\n BRAND_STACKOVERFLOW (IconName): The :code:`brand-stackoverflow` icon.\n BRAND_STACKSHARE (IconName): The :code:`brand-stackshare` icon.\n BRAND_STEAM (IconName): The :code:`brand-steam` icon.\n BRAND_STORJ (IconName): The :code:`brand-storj` icon.\n BRAND_STORYBOOK (IconName): The :code:`brand-storybook` icon.\n BRAND_STORYTEL (IconName): The :code:`brand-storytel` icon.\n BRAND_STRAVA (IconName): The :code:`brand-strava` icon.\n BRAND_STRIPE (IconName): The :code:`brand-stripe` icon.\n BRAND_SUBLIME_TEXT (IconName): The :code:`brand-sublime-text` icon.\n BRAND_SUGARIZER (IconName): The :code:`brand-sugarizer` icon.\n BRAND_SUPABASE (IconName): The :code:`brand-supabase` icon.\n BRAND_SUPERHUMAN (IconName): The :code:`brand-superhuman` icon.\n BRAND_SUPERNOVA (IconName): The :code:`brand-supernova` icon.\n BRAND_SURFSHARK (IconName): The :code:`brand-surfshark` icon.\n BRAND_SVELTE (IconName): The :code:`brand-svelte` icon.\n BRAND_SWIFT (IconName): The :code:`brand-swift` icon.\n BRAND_SYMFONY (IconName): The :code:`brand-symfony` icon.\n BRAND_TABLER (IconName): The :code:`brand-tabler` icon.\n BRAND_TAILWIND (IconName): The :code:`brand-tailwind` icon.\n BRAND_TAOBAO (IconName): The :code:`brand-taobao` icon.\n BRAND_TED (IconName): The :code:`brand-ted` icon.\n BRAND_TELEGRAM (IconName): The :code:`brand-telegram` icon.\n BRAND_TERRAFORM (IconName): The :code:`brand-terraform` icon.\n BRAND_TETHER (IconName): The :code:`brand-tether` icon.\n BRAND_THREEJS (IconName): The :code:`brand-threejs` icon.\n BRAND_TIDAL (IconName): The :code:`brand-tidal` icon.\n BRAND_TIKTO_FILLED (IconName): The :code:`brand-tikto-filled` icon.\n BRAND_TIKTOK (IconName): The :code:`brand-tiktok` icon.\n BRAND_TINDER (IconName): The :code:`brand-tinder` icon.\n BRAND_TOPBUZZ (IconName): The :code:`brand-topbuzz` icon.\n BRAND_TORCHAIN (IconName): The :code:`brand-torchain` icon.\n BRAND_TOYOTA (IconName): The :code:`brand-toyota` icon.\n BRAND_TRELLO (IconName): The :code:`brand-trello` icon.\n BRAND_TRIPADVISOR (IconName): The :code:`brand-tripadvisor` icon.\n BRAND_TUMBLR (IconName): The :code:`brand-tumblr` icon.\n BRAND_TWILIO (IconName): The :code:`brand-twilio` icon.\n BRAND_TWITCH (IconName): The :code:`brand-twitch` icon.\n BRAND_TWITTER (IconName): The :code:`brand-twitter` icon.\n BRAND_TWITTER_FILLED (IconName): The :code:`brand-twitter-filled` icon.\n BRAND_TYPESCRIPT (IconName): The :code:`brand-typescript` icon.\n BRAND_UBER (IconName): The :code:`brand-uber` icon.\n BRAND_UBUNTU (IconName): The :code:`brand-ubuntu` icon.\n BRAND_UNITY (IconName): The :code:`brand-unity` icon.\n BRAND_UNSPLASH (IconName): The :code:`brand-unsplash` icon.\n BRAND_UPWORK (IconName): The :code:`brand-upwork` icon.\n BRAND_VALORANT (IconName): The :code:`brand-valorant` icon.\n BRAND_VERCEL (IconName): The :code:`brand-vercel` icon.\n BRAND_VIMEO (IconName): The :code:`brand-vimeo` icon.\n BRAND_VINTED (IconName): The :code:`brand-vinted` icon.\n BRAND_VISA (IconName): The :code:`brand-visa` icon.\n BRAND_VISUAL_STUDIO (IconName): The :code:`brand-visual-studio` icon.\n BRAND_VITE (IconName): The :code:`brand-vite` icon.\n BRAND_VIVALDI (IconName): The :code:`brand-vivaldi` icon.\n BRAND_VK (IconName): The :code:`brand-vk` icon.\n BRAND_VLC (IconName): The :code:`brand-vlc` icon.\n BRAND_VOLKSWAGEN (IconName): The :code:`brand-volkswagen` icon.\n BRAND_VSCO (IconName): The :code:`brand-vsco` icon.\n BRAND_VSCODE (IconName): The :code:`brand-vscode` icon.\n BRAND_VUE (IconName): The :code:`brand-vue` icon.\n BRAND_WALMART (IconName): The :code:`brand-walmart` icon.\n BRAND_WAZE (IconName): The :code:`brand-waze` icon.\n BRAND_WEBFLOW (IconName): The :code:`brand-webflow` icon.\n BRAND_WECHAT (IconName): The :code:`brand-wechat` icon.\n BRAND_WEIBO (IconName): The :code:`brand-weibo` icon.\n BRAND_WHATSAPP (IconName): The :code:`brand-whatsapp` icon.\n BRAND_WIKIPEDIA (IconName): The :code:`brand-wikipedia` icon.\n BRAND_WINDOWS (IconName): The :code:`brand-windows` icon.\n BRAND_WINDY (IconName): The :code:`brand-windy` icon.\n BRAND_WISH (IconName): The :code:`brand-wish` icon.\n BRAND_WIX (IconName): The :code:`brand-wix` icon.\n BRAND_WORDPRESS (IconName): The :code:`brand-wordpress` icon.\n BRAND_XAMARIN (IconName): The :code:`brand-xamarin` icon.\n BRAND_XBOX (IconName): The :code:`brand-xbox` icon.\n BRAND_XING (IconName): The :code:`brand-xing` icon.\n BRAND_YAHOO (IconName): The :code:`brand-yahoo` icon.\n BRAND_YANDEX (IconName): The :code:`brand-yandex` icon.\n BRAND_YATSE (IconName): The :code:`brand-yatse` icon.\n BRAND_YCOMBINATOR (IconName): The :code:`brand-ycombinator` icon.\n BRAND_YOUTUBE (IconName): The :code:`brand-youtube` icon.\n BRAND_YOUTUBE_KIDS (IconName): The :code:`brand-youtube-kids` icon.\n BRAND_ZALANDO (IconName): The :code:`brand-zalando` icon.\n BRAND_ZAPIER (IconName): The :code:`brand-zapier` icon.\n BRAND_ZEIT (IconName): The :code:`brand-zeit` icon.\n BRAND_ZHIHU (IconName): The :code:`brand-zhihu` icon.\n BRAND_ZOOM (IconName): The :code:`brand-zoom` icon.\n BRAND_ZULIP (IconName): The :code:`brand-zulip` icon.\n BRAND_ZWIFT (IconName): The :code:`brand-zwift` icon.\n BREAD (IconName): The :code:`bread` icon.\n BREAD_OFF (IconName): The :code:`bread-off` icon.\n BRIEFCASE (IconName): The :code:`briefcase` icon.\n BRIEFCASE_OFF (IconName): The :code:`briefcase-off` icon.\n BRIGHTNESS (IconName): The :code:`brightness` icon.\n BRIGHTNESS_2 (IconName): The :code:`brightness-2` icon.\n BRIGHTNESS_DOWN (IconName): The :code:`brightness-down` icon.\n BRIGHTNESS_HALF (IconName): The :code:`brightness-half` icon.\n BRIGHTNESS_OFF (IconName): The :code:`brightness-off` icon.\n BRIGHTNESS_UP (IconName): The :code:`brightness-up` icon.\n BROADCAST (IconName): The :code:`broadcast` icon.\n BROADCAST_OFF (IconName): The :code:`broadcast-off` icon.\n BROWSER (IconName): The :code:`browser` icon.\n BROWSER_CHECK (IconName): The :code:`browser-check` icon.\n BROWSER_OFF (IconName): The :code:`browser-off` icon.\n BROWSER_PLUS (IconName): The :code:`browser-plus` icon.\n BROWSER_X (IconName): The :code:`browser-x` icon.\n BRUSH (IconName): The :code:`brush` icon.\n BRUSH_OFF (IconName): The :code:`brush-off` icon.\n BUCKET (IconName): The :code:`bucket` icon.\n BUCKET_DROPLET (IconName): The :code:`bucket-droplet` icon.\n BUCKET_OFF (IconName): The :code:`bucket-off` icon.\n BUG (IconName): The :code:`bug` icon.\n BUG_OFF (IconName): The :code:`bug-off` icon.\n BUILDING (IconName): The :code:`building` icon.\n BUILDING_ARCH (IconName): The :code:`building-arch` icon.\n BUILDING_BANK (IconName): The :code:`building-bank` icon.\n BUILDING_BRIDGE (IconName): The :code:`building-bridge` icon.\n BUILDING_BRIDGE_2 (IconName): The :code:`building-bridge-2` icon.\n BUILDING_BROADCAST_TOWER (IconName): The :code:`building-broadcast-tower` icon.\n BUILDING_CAROUSEL (IconName): The :code:`building-carousel` icon.\n BUILDING_CASTLE (IconName): The :code:`building-castle` icon.\n BUILDING_CHURCH (IconName): The :code:`building-church` icon.\n BUILDING_CIRCUS (IconName): The :code:`building-circus` icon.\n BUILDING_COMMUNITY (IconName): The :code:`building-community` icon.\n BUILDING_COTTAGE (IconName): The :code:`building-cottage` icon.\n BUILDING_ESTATE (IconName): The :code:`building-estate` icon.\n BUILDING_FACTORY (IconName): The :code:`building-factory` icon.\n BUILDING_FACTORY_2 (IconName): The :code:`building-factory-2` icon.\n BUILDING_FORTRESS (IconName): The :code:`building-fortress` icon.\n BUILDING_HOSPITAL (IconName): The :code:`building-hospital` icon.\n BUILDING_LIGHTHOUSE (IconName): The :code:`building-lighthouse` icon.\n BUILDING_MONUMENT (IconName): The :code:`building-monument` icon.\n BUILDING_MOSQUE (IconName): The :code:`building-mosque` icon.\n BUILDING_PAVILION (IconName): The :code:`building-pavilion` icon.\n BUILDING_SKYSCRAPER (IconName): The :code:`building-skyscraper` icon.\n BUILDING_STADIUM (IconName): The :code:`building-stadium` icon.\n BUILDING_STORE (IconName): The :code:`building-store` icon.\n BUILDING_TUNNEL (IconName): The :code:`building-tunnel` icon.\n BUILDING_WAREHOUSE (IconName): The :code:`building-warehouse` icon.\n BUILDING_WIND_TURBINE (IconName): The :code:`building-wind-turbine` icon.\n BULB (IconName): The :code:`bulb` icon.\n BULB_FILLED (IconName): The :code:`bulb-filled` icon.\n BULB_OFF (IconName): The :code:`bulb-off` icon.\n BULLDOZER (IconName): The :code:`bulldozer` icon.\n BUS (IconName): The :code:`bus` icon.\n BUS_OFF (IconName): The :code:`bus-off` icon.\n BUS_STOP (IconName): The :code:`bus-stop` icon.\n BUSINESSPLAN (IconName): The :code:`businessplan` icon.\n BUTTERFLY (IconName): The :code:`butterfly` icon.\n CACTUS (IconName): The :code:`cactus` icon.\n CACTUS_OFF (IconName): The :code:`cactus-off` icon.\n CAKE (IconName): The :code:`cake` icon.\n CAKE_OFF (IconName): The :code:`cake-off` icon.\n CALCULATOR (IconName): The :code:`calculator` icon.\n CALCULATOR_OFF (IconName): The :code:`calculator-off` icon.\n CALENDAR (IconName): The :code:`calendar` icon.\n CALENDAR_BOLT (IconName): The :code:`calendar-bolt` icon.\n CALENDAR_CANCEL (IconName): The :code:`calendar-cancel` icon.\n CALENDAR_CHECK (IconName): The :code:`calendar-check` icon.\n CALENDAR_CODE (IconName): The :code:`calendar-code` icon.\n CALENDAR_COG (IconName): The :code:`calendar-cog` icon.\n CALENDAR_DOLLAR (IconName): The :code:`calendar-dollar` icon.\n CALENDAR_DOWN (IconName): The :code:`calendar-down` icon.\n CALENDAR_DUE (IconName): The :code:`calendar-due` icon.\n CALENDAR_EVENT (IconName): The :code:`calendar-event` icon.\n CALENDAR_EXCLAMATION (IconName): The :code:`calendar-exclamation` icon.\n CALENDAR_HEART (IconName): The :code:`calendar-heart` icon.\n CALENDAR_MINUS (IconName): The :code:`calendar-minus` icon.\n CALENDAR_OFF (IconName): The :code:`calendar-off` icon.\n CALENDAR_PAUSE (IconName): The :code:`calendar-pause` icon.\n CALENDAR_PIN (IconName): The :code:`calendar-pin` icon.\n CALENDAR_PLUS (IconName): The :code:`calendar-plus` icon.\n CALENDAR_QUESTION (IconName): The :code:`calendar-question` icon.\n CALENDAR_REPEAT (IconName): The :code:`calendar-repeat` icon.\n CALENDAR_SEARCH (IconName): The :code:`calendar-search` icon.\n CALENDAR_SHARE (IconName): The :code:`calendar-share` icon.\n CALENDAR_STAR (IconName): The :code:`calendar-star` icon.\n CALENDAR_STATS (IconName): The :code:`calendar-stats` icon.\n CALENDAR_TIME (IconName): The :code:`calendar-time` icon.\n CALENDAR_UP (IconName): The :code:`calendar-up` icon.\n CALENDAR_X (IconName): The :code:`calendar-x` icon.\n CAMERA (IconName): The :code:`camera` icon.\n CAMERA_BOLT (IconName): The :code:`camera-bolt` icon.\n CAMERA_CANCEL (IconName): The :code:`camera-cancel` icon.\n CAMERA_CHECK (IconName): The :code:`camera-check` icon.\n CAMERA_CODE (IconName): The :code:`camera-code` icon.\n CAMERA_COG (IconName): The :code:`camera-cog` icon.\n CAMERA_DOLLAR (IconName): The :code:`camera-dollar` icon.\n CAMERA_DOWN (IconName): The :code:`camera-down` icon.\n CAMERA_EXCLAMATION (IconName): The :code:`camera-exclamation` icon.\n CAMERA_FILLED (IconName): The :code:`camera-filled` icon.\n CAMERA_HEART (IconName): The :code:`camera-heart` icon.\n CAMERA_MINUS (IconName): The :code:`camera-minus` icon.\n CAMERA_OFF (IconName): The :code:`camera-off` icon.\n CAMERA_PAUSE (IconName): The :code:`camera-pause` icon.\n CAMERA_PIN (IconName): The :code:`camera-pin` icon.\n CAMERA_PLUS (IconName): The :code:`camera-plus` icon.\n CAMERA_QUESTION (IconName): The :code:`camera-question` icon.\n CAMERA_ROTATE (IconName): The :code:`camera-rotate` icon.\n CAMERA_SEARCH (IconName): The :code:`camera-search` icon.\n CAMERA_SELFIE (IconName): The :code:`camera-selfie` icon.\n CAMERA_SHARE (IconName): The :code:`camera-share` icon.\n CAMERA_STAR (IconName): The :code:`camera-star` icon.\n CAMERA_UP (IconName): The :code:`camera-up` icon.\n CAMERA_X (IconName): The :code:`camera-x` icon.\n CAMPER (IconName): The :code:`camper` icon.\n CAMPFIRE (IconName): The :code:`campfire` icon.\n CANDLE (IconName): The :code:`candle` icon.\n CANDY (IconName): The :code:`candy` icon.\n CANDY_OFF (IconName): The :code:`candy-off` icon.\n CANE (IconName): The :code:`cane` icon.\n CANNABIS (IconName): The :code:`cannabis` icon.\n CAPSULE (IconName): The :code:`capsule` icon.\n CAPSULE_HORIZONTAL (IconName): The :code:`capsule-horizontal` icon.\n CAPTURE (IconName): The :code:`capture` icon.\n CAPTURE_OFF (IconName): The :code:`capture-off` icon.\n CAR (IconName): The :code:`car` icon.\n CAR_CRANE (IconName): The :code:`car-crane` icon.\n CAR_CRASH (IconName): The :code:`car-crash` icon.\n CAR_OFF (IconName): The :code:`car-off` icon.\n CAR_TURBINE (IconName): The :code:`car-turbine` icon.\n CARAVAN (IconName): The :code:`caravan` icon.\n CARDBOARDS (IconName): The :code:`cardboards` icon.\n CARDBOARDS_OFF (IconName): The :code:`cardboards-off` icon.\n CARDS (IconName): The :code:`cards` icon.\n CARET_DOWN (IconName): The :code:`caret-down` icon.\n CARET_LEFT (IconName): The :code:`caret-left` icon.\n CARET_RIGHT (IconName): The :code:`caret-right` icon.\n CARET_UP (IconName): The :code:`caret-up` icon.\n CAROUSEL_HORIZONTAL (IconName): The :code:`carousel-horizontal` icon.\n CAROUSEL_HORIZONTAL_FILLED (IconName): The :code:`carousel-horizontal-filled` icon.\n CAROUSEL_VERTICAL (IconName): The :code:`carousel-vertical` icon.\n CAROUSEL_VERTICAL_FILLED (IconName): The :code:`carousel-vertical-filled` icon.\n CARROT (IconName): The :code:`carrot` icon.\n CARROT_OFF (IconName): The :code:`carrot-off` icon.\n CASH (IconName): The :code:`cash` icon.\n CASH_BANKNOTE (IconName): The :code:`cash-banknote` icon.\n CASH_BANKNOTE_OFF (IconName): The :code:`cash-banknote-off` icon.\n CASH_OFF (IconName): The :code:`cash-off` icon.\n CAST (IconName): The :code:`cast` icon.\n CAST_OFF (IconName): The :code:`cast-off` icon.\n CAT (IconName): The :code:`cat` icon.\n CATEGORY (IconName): The :code:`category` icon.\n CATEGORY_2 (IconName): The :code:`category-2` icon.\n CE (IconName): The :code:`ce` icon.\n CE_OFF (IconName): The :code:`ce-off` icon.\n CELL (IconName): The :code:`cell` icon.\n CELL_SIGNAL_1 (IconName): The :code:`cell-signal-1` icon.\n CELL_SIGNAL_2 (IconName): The :code:`cell-signal-2` icon.\n CELL_SIGNAL_3 (IconName): The :code:`cell-signal-3` icon.\n CELL_SIGNAL_4 (IconName): The :code:`cell-signal-4` icon.\n CELL_SIGNAL_5 (IconName): The :code:`cell-signal-5` icon.\n CELL_SIGNAL_OFF (IconName): The :code:`cell-signal-off` icon.\n CERTIFICATE (IconName): The :code:`certificate` icon.\n CERTIFICATE_2 (IconName): The :code:`certificate-2` icon.\n CERTIFICATE_2_OFF (IconName): The :code:`certificate-2-off` icon.\n CERTIFICATE_OFF (IconName): The :code:`certificate-off` icon.\n CHAIR_DIRECTOR (IconName): The :code:`chair-director` icon.\n CHALKBOARD (IconName): The :code:`chalkboard` icon.\n CHALKBOARD_OFF (IconName): The :code:`chalkboard-off` icon.\n CHARGING_PILE (IconName): The :code:`charging-pile` icon.\n CHART_ARCS (IconName): The :code:`chart-arcs` icon.\n CHART_ARCS_3 (IconName): The :code:`chart-arcs-3` icon.\n CHART_AREA (IconName): The :code:`chart-area` icon.\n CHART_AREA_FILLED (IconName): The :code:`chart-area-filled` icon.\n CHART_AREA_LINE (IconName): The :code:`chart-area-line` icon.\n CHART_AREA_LINE_FILLED (IconName): The :code:`chart-area-line-filled` icon.\n CHART_ARROWS (IconName): The :code:`chart-arrows` icon.\n CHART_ARROWS_VERTICAL (IconName): The :code:`chart-arrows-vertical` icon.\n CHART_BAR (IconName): The :code:`chart-bar` icon.\n CHART_BAR_OFF (IconName): The :code:`chart-bar-off` icon.\n CHART_BUBBLE (IconName): The :code:`chart-bubble` icon.\n CHART_BUBBLE_FILLED (IconName): The :code:`chart-bubble-filled` icon.\n CHART_CANDLE (IconName): The :code:`chart-candle` icon.\n CHART_CANDLE_FILLED (IconName): The :code:`chart-candle-filled` icon.\n CHART_CIRCLES (IconName): The :code:`chart-circles` icon.\n CHART_DONUT (IconName): The :code:`chart-donut` icon.\n CHART_DONUT_2 (IconName): The :code:`chart-donut-2` icon.\n CHART_DONUT_3 (IconName): The :code:`chart-donut-3` icon.\n CHART_DONUT_4 (IconName): The :code:`chart-donut-4` icon.\n CHART_DONUT_FILLED (IconName): The :code:`chart-donut-filled` icon.\n CHART_DOTS (IconName): The :code:`chart-dots` icon.\n CHART_DOTS_2 (IconName): The :code:`chart-dots-2` icon.\n CHART_DOTS_3 (IconName): The :code:`chart-dots-3` icon.\n CHART_GRID_DOTS (IconName): The :code:`chart-grid-dots` icon.\n CHART_HISTOGRAM (IconName): The :code:`chart-histogram` icon.\n CHART_INFOGRAPHIC (IconName): The :code:`chart-infographic` icon.\n CHART_LINE (IconName): The :code:`chart-line` icon.\n CHART_PIE (IconName): The :code:`chart-pie` icon.\n CHART_PIE_2 (IconName): The :code:`chart-pie-2` icon.\n CHART_PIE_3 (IconName): The :code:`chart-pie-3` icon.\n CHART_PIE_4 (IconName): The :code:`chart-pie-4` icon.\n CHART_PIE_FILLED (IconName): The :code:`chart-pie-filled` icon.\n CHART_PIE_OFF (IconName): The :code:`chart-pie-off` icon.\n CHART_PPF (IconName): The :code:`chart-ppf` icon.\n CHART_RADAR (IconName): The :code:`chart-radar` icon.\n CHART_SANKEY (IconName): The :code:`chart-sankey` icon.\n CHART_TREEMAP (IconName): The :code:`chart-treemap` icon.\n CHECK (IconName): The :code:`check` icon.\n CHECKBOX (IconName): The :code:`checkbox` icon.\n CHECKLIST (IconName): The :code:`checklist` icon.\n CHECKS (IconName): The :code:`checks` icon.\n CHECKUP_LIST (IconName): The :code:`checkup-list` icon.\n CHEESE (IconName): The :code:`cheese` icon.\n CHEF_HAT (IconName): The :code:`chef-hat` icon.\n CHEF_HAT_OFF (IconName): The :code:`chef-hat-off` icon.\n CHERRY (IconName): The :code:`cherry` icon.\n CHERRY_FILLED (IconName): The :code:`cherry-filled` icon.\n CHESS (IconName): The :code:`chess` icon.\n CHESS_BISHOP (IconName): The :code:`chess-bishop` icon.\n CHESS_BISHOP_FILLED (IconName): The :code:`chess-bishop-filled` icon.\n CHESS_FILLED (IconName): The :code:`chess-filled` icon.\n CHESS_KING (IconName): The :code:`chess-king` icon.\n CHESS_KING_FILLED (IconName): The :code:`chess-king-filled` icon.\n CHESS_KNIGHT (IconName): The :code:`chess-knight` icon.\n CHESS_KNIGHT_FILLED (IconName): The :code:`chess-knight-filled` icon.\n CHESS_QUEEN (IconName): The :code:`chess-queen` icon.\n CHESS_QUEEN_FILLED (IconName): The :code:`chess-queen-filled` icon.\n CHESS_ROOK (IconName): The :code:`chess-rook` icon.\n CHESS_ROOK_FILLED (IconName): The :code:`chess-rook-filled` icon.\n CHEVRON_COMPACT_DOWN (IconName): The :code:`chevron-compact-down` icon.\n CHEVRON_COMPACT_LEFT (IconName): The :code:`chevron-compact-left` icon.\n CHEVRON_COMPACT_RIGHT (IconName): The :code:`chevron-compact-right` icon.\n CHEVRON_COMPACT_UP (IconName): The :code:`chevron-compact-up` icon.\n CHEVRON_DOWN (IconName): The :code:`chevron-down` icon.\n CHEVRON_DOWN_LEFT (IconName): The :code:`chevron-down-left` icon.\n CHEVRON_DOWN_RIGHT (IconName): The :code:`chevron-down-right` icon.\n CHEVRON_LEFT (IconName): The :code:`chevron-left` icon.\n CHEVRON_LEFT_PIPE (IconName): The :code:`chevron-left-pipe` icon.\n CHEVRON_RIGHT (IconName): The :code:`chevron-right` icon.\n CHEVRON_RIGHT_PIPE (IconName): The :code:`chevron-right-pipe` icon.\n CHEVRON_UP (IconName): The :code:`chevron-up` icon.\n CHEVRON_UP_LEFT (IconName): The :code:`chevron-up-left` icon.\n CHEVRON_UP_RIGHT (IconName): The :code:`chevron-up-right` icon.\n CHEVRONS_DOWN (IconName): The :code:`chevrons-down` icon.\n CHEVRONS_DOWN_LEFT (IconName): The :code:`chevrons-down-left` icon.\n CHEVRONS_DOWN_RIGHT (IconName): The :code:`chevrons-down-right` icon.\n CHEVRONS_LEFT (IconName): The :code:`chevrons-left` icon.\n CHEVRONS_RIGHT (IconName): The :code:`chevrons-right` icon.\n CHEVRONS_UP (IconName): The :code:`chevrons-up` icon.\n CHEVRONS_UP_LEFT (IconName): The :code:`chevrons-up-left` icon.\n CHEVRONS_UP_RIGHT (IconName): The :code:`chevrons-up-right` icon.\n CHISEL (IconName): The :code:`chisel` icon.\n CHRISTMAS_TREE (IconName): The :code:`christmas-tree` icon.\n CHRISTMAS_TREE_OFF (IconName): The :code:`christmas-tree-off` icon.\n CIRCLE (IconName): The :code:`circle` icon.\n CIRCLE_0_FILLED (IconName): The :code:`circle-0-filled` icon.\n CIRCLE_1_FILLED (IconName): The :code:`circle-1-filled` icon.\n CIRCLE_2_FILLED (IconName): The :code:`circle-2-filled` icon.\n CIRCLE_3_FILLED (IconName): The :code:`circle-3-filled` icon.\n CIRCLE_4_FILLED (IconName): The :code:`circle-4-filled` icon.\n CIRCLE_5_FILLED (IconName): The :code:`circle-5-filled` icon.\n CIRCLE_6_FILLED (IconName): The :code:`circle-6-filled` icon.\n CIRCLE_7_FILLED (IconName): The :code:`circle-7-filled` icon.\n CIRCLE_8_FILLED (IconName): The :code:`circle-8-filled` icon.\n CIRCLE_9_FILLED (IconName): The :code:`circle-9-filled` icon.\n CIRCLE_ARROW_DOWN (IconName): The :code:`circle-arrow-down` icon.\n CIRCLE_ARROW_DOWN_FILLED (IconName): The :code:`circle-arrow-down-filled` icon.\n CIRCLE_ARROW_DOWN_LEFT (IconName): The :code:`circle-arrow-down-left` icon.\n CIRCLE_ARROW_DOWN_LEFT_FILLED (IconName): The :code:`circle-arrow-down-left-filled` icon.\n CIRCLE_ARROW_DOWN_RIGHT (IconName): The :code:`circle-arrow-down-right` icon.\n CIRCLE_ARROW_DOWN_RIGHT_FILLED (IconName): The :code:`circle-arrow-down-right-filled` icon.\n CIRCLE_ARROW_LEFT (IconName): The :code:`circle-arrow-left` icon.\n CIRCLE_ARROW_LEFT_FILLED (IconName): The :code:`circle-arrow-left-filled` icon.\n CIRCLE_ARROW_RIGHT (IconName): The :code:`circle-arrow-right` icon.\n CIRCLE_ARROW_RIGHT_FILLED (IconName): The :code:`circle-arrow-right-filled` icon.\n CIRCLE_ARROW_UP (IconName): The :code:`circle-arrow-up` icon.\n CIRCLE_ARROW_UP_FILLED (IconName): The :code:`circle-arrow-up-filled` icon.\n CIRCLE_ARROW_UP_LEFT (IconName): The :code:`circle-arrow-up-left` icon.\n CIRCLE_ARROW_UP_LEFT_FILLED (IconName): The :code:`circle-arrow-up-left-filled` icon.\n CIRCLE_ARROW_UP_RIGHT (IconName): The :code:`circle-arrow-up-right` icon.\n CIRCLE_ARROW_UP_RIGHT_FILLED (IconName): The :code:`circle-arrow-up-right-filled` icon.\n CIRCLE_CARET_DOWN (IconName): The :code:`circle-caret-down` icon.\n CIRCLE_CARET_LEFT (IconName): The :code:`circle-caret-left` icon.\n CIRCLE_CARET_RIGHT (IconName): The :code:`circle-caret-right` icon.\n CIRCLE_CARET_UP (IconName): The :code:`circle-caret-up` icon.\n CIRCLE_CHECK (IconName): The :code:`circle-check` icon.\n CIRCLE_CHECK_FILLED (IconName): The :code:`circle-check-filled` icon.\n CIRCLE_CHEVRON_DOWN (IconName): The :code:`circle-chevron-down` icon.\n CIRCLE_CHEVRON_LEFT (IconName): The :code:`circle-chevron-left` icon.\n CIRCLE_CHEVRON_RIGHT (IconName): The :code:`circle-chevron-right` icon.\n CIRCLE_CHEVRON_UP (IconName): The :code:`circle-chevron-up` icon.\n CIRCLE_CHEVRONS_DOWN (IconName): The :code:`circle-chevrons-down` icon.\n CIRCLE_CHEVRONS_LEFT (IconName): The :code:`circle-chevrons-left` icon.\n CIRCLE_CHEVRONS_RIGHT (IconName): The :code:`circle-chevrons-right` icon.\n CIRCLE_CHEVRONS_UP (IconName): The :code:`circle-chevrons-up` icon.\n CIRCLE_DASHED (IconName): The :code:`circle-dashed` icon.\n CIRCLE_DOT (IconName): The :code:`circle-dot` icon.\n CIRCLE_DOT_FILLED (IconName): The :code:`circle-dot-filled` icon.\n CIRCLE_DOTTED (IconName): The :code:`circle-dotted` icon.\n CIRCLE_FILLED (IconName): The :code:`circle-filled` icon.\n CIRCLE_HALF (IconName): The :code:`circle-half` icon.\n CIRCLE_HALF_2 (IconName): The :code:`circle-half-2` icon.\n CIRCLE_HALF_VERTICAL (IconName): The :code:`circle-half-vertical` icon.\n CIRCLE_KEY (IconName): The :code:`circle-key` icon.\n CIRCLE_KEY_FILLED (IconName): The :code:`circle-key-filled` icon.\n CIRCLE_LETTER_A (IconName): The :code:`circle-letter-a` icon.\n CIRCLE_LETTER_B (IconName): The :code:`circle-letter-b` icon.\n CIRCLE_LETTER_C (IconName): The :code:`circle-letter-c` icon.\n CIRCLE_LETTER_D (IconName): The :code:`circle-letter-d` icon.\n CIRCLE_LETTER_E (IconName): The :code:`circle-letter-e` icon.\n CIRCLE_LETTER_F (IconName): The :code:`circle-letter-f` icon.\n CIRCLE_LETTER_G (IconName): The :code:`circle-letter-g` icon.\n CIRCLE_LETTER_H (IconName): The :code:`circle-letter-h` icon.\n CIRCLE_LETTER_I (IconName): The :code:`circle-letter-i` icon.\n CIRCLE_LETTER_J (IconName): The :code:`circle-letter-j` icon.\n CIRCLE_LETTER_K (IconName): The :code:`circle-letter-k` icon.\n CIRCLE_LETTER_L (IconName): The :code:`circle-letter-l` icon.\n CIRCLE_LETTER_M (IconName): The :code:`circle-letter-m` icon.\n CIRCLE_LETTER_N (IconName): The :code:`circle-letter-n` icon.\n CIRCLE_LETTER_O (IconName): The :code:`circle-letter-o` icon.\n CIRCLE_LETTER_P (IconName): The :code:`circle-letter-p` icon.\n CIRCLE_LETTER_Q (IconName): The :code:`circle-letter-q` icon.\n CIRCLE_LETTER_R (IconName): The :code:`circle-letter-r` icon.\n CIRCLE_LETTER_S (IconName): The :code:`circle-letter-s` icon.\n CIRCLE_LETTER_T (IconName): The :code:`circle-letter-t` icon.\n CIRCLE_LETTER_U (IconName): The :code:`circle-letter-u` icon.\n CIRCLE_LETTER_V (IconName): The :code:`circle-letter-v` icon.\n CIRCLE_LETTER_W (IconName): The :code:`circle-letter-w` icon.\n CIRCLE_LETTER_X (IconName): The :code:`circle-letter-x` icon.\n CIRCLE_LETTER_Y (IconName): The :code:`circle-letter-y` icon.\n CIRCLE_LETTER_Z (IconName): The :code:`circle-letter-z` icon.\n CIRCLE_MINUS (IconName): The :code:`circle-minus` icon.\n CIRCLE_NUMBER_0 (IconName): The :code:`circle-number-0` icon.\n CIRCLE_NUMBER_1 (IconName): The :code:`circle-number-1` icon.\n CIRCLE_NUMBER_2 (IconName): The :code:`circle-number-2` icon.\n CIRCLE_NUMBER_3 (IconName): The :code:`circle-number-3` icon.\n CIRCLE_NUMBER_4 (IconName): The :code:`circle-number-4` icon.\n CIRCLE_NUMBER_5 (IconName): The :code:`circle-number-5` icon.\n CIRCLE_NUMBER_6 (IconName): The :code:`circle-number-6` icon.\n CIRCLE_NUMBER_7 (IconName): The :code:`circle-number-7` icon.\n CIRCLE_NUMBER_8 (IconName): The :code:`circle-number-8` icon.\n CIRCLE_NUMBER_9 (IconName): The :code:`circle-number-9` icon.\n CIRCLE_OFF (IconName): The :code:`circle-off` icon.\n CIRCLE_PLUS (IconName): The :code:`circle-plus` icon.\n CIRCLE_RECTANGLE (IconName): The :code:`circle-rectangle` icon.\n CIRCLE_RECTANGLE_OFF (IconName): The :code:`circle-rectangle-off` icon.\n CIRCLE_SQUARE (IconName): The :code:`circle-square` icon.\n CIRCLE_TRIANGLE (IconName): The :code:`circle-triangle` icon.\n CIRCLE_X (IconName): The :code:`circle-x` icon.\n CIRCLE_X_FILLED (IconName): The :code:`circle-x-filled` icon.\n CIRCLES (IconName): The :code:`circles` icon.\n CIRCLES_FILLED (IconName): The :code:`circles-filled` icon.\n CIRCLES_RELATION (IconName): The :code:`circles-relation` icon.\n CIRCUIT_AMMETER (IconName): The :code:`circuit-ammeter` icon.\n CIRCUIT_BATTERY (IconName): The :code:`circuit-battery` icon.\n CIRCUIT_BULB (IconName): The :code:`circuit-bulb` icon.\n CIRCUIT_CAPACITOR (IconName): The :code:`circuit-capacitor` icon.\n CIRCUIT_CAPACITOR_POLARIZED (IconName): The :code:`circuit-capacitor-polarized` icon.\n CIRCUIT_CELL (IconName): The :code:`circuit-cell` icon.\n CIRCUIT_CELL_PLUS (IconName): The :code:`circuit-cell-plus` icon.\n CIRCUIT_CHANGEOVER (IconName): The :code:`circuit-changeover` icon.\n CIRCUIT_DIODE (IconName): The :code:`circuit-diode` icon.\n CIRCUIT_DIODE_ZENER (IconName): The :code:`circuit-diode-zener` icon.\n CIRCUIT_GROUND (IconName): The :code:`circuit-ground` icon.\n CIRCUIT_GROUND_DIGITAL (IconName): The :code:`circuit-ground-digital` icon.\n CIRCUIT_INDUCTOR (IconName): The :code:`circuit-inductor` icon.\n CIRCUIT_MOTOR (IconName): The :code:`circuit-motor` icon.\n CIRCUIT_PUSHBUTTON (IconName): The :code:`circuit-pushbutton` icon.\n CIRCUIT_RESISTOR (IconName): The :code:`circuit-resistor` icon.\n CIRCUIT_SWITCH_CLOSED (IconName): The :code:`circuit-switch-closed` icon.\n CIRCUIT_SWITCH_OPEN (IconName): The :code:`circuit-switch-open` icon.\n CIRCUIT_VOLTMETER (IconName): The :code:`circuit-voltmeter` icon.\n CLEAR_ALL (IconName): The :code:`clear-all` icon.\n CLEAR_FORMATTING (IconName): The :code:`clear-formatting` icon.\n CLICK (IconName): The :code:`click` icon.\n CLIPBOARD (IconName): The :code:`clipboard` icon.\n CLIPBOARD_CHECK (IconName): The :code:`clipboard-check` icon.\n CLIPBOARD_COPY (IconName): The :code:`clipboard-copy` icon.\n CLIPBOARD_DATA (IconName): The :code:`clipboard-data` icon.\n CLIPBOARD_HEART (IconName): The :code:`clipboard-heart` icon.\n CLIPBOARD_LIST (IconName): The :code:`clipboard-list` icon.\n CLIPBOARD_OFF (IconName): The :code:`clipboard-off` icon.\n CLIPBOARD_PLUS (IconName): The :code:`clipboard-plus` icon.\n CLIPBOARD_TEXT (IconName): The :code:`clipboard-text` icon.\n CLIPBOARD_TYPOGRAPHY (IconName): The :code:`clipboard-typography` icon.\n CLIPBOARD_X (IconName): The :code:`clipboard-x` icon.\n CLOCK (IconName): The :code:`clock` icon.\n CLOCK_2 (IconName): The :code:`clock-2` icon.\n CLOCK_BOLT (IconName): The :code:`clock-bolt` icon.\n CLOCK_CANCEL (IconName): The :code:`clock-cancel` icon.\n CLOCK_CHECK (IconName): The :code:`clock-check` icon.\n CLOCK_CODE (IconName): The :code:`clock-code` icon.\n CLOCK_COG (IconName): The :code:`clock-cog` icon.\n CLOCK_DOLLAR (IconName): The :code:`clock-dollar` icon.\n CLOCK_DOWN (IconName): The :code:`clock-down` icon.\n CLOCK_EDIT (IconName): The :code:`clock-edit` icon.\n CLOCK_EXCLAMATION (IconName): The :code:`clock-exclamation` icon.\n CLOCK_FILLED (IconName): The :code:`clock-filled` icon.\n CLOCK_HEART (IconName): The :code:`clock-heart` icon.\n CLOCK_HOUR_1 (IconName): The :code:`clock-hour-1` icon.\n CLOCK_HOUR_10 (IconName): The :code:`clock-hour-10` icon.\n CLOCK_HOUR_11 (IconName): The :code:`clock-hour-11` icon.\n CLOCK_HOUR_12 (IconName): The :code:`clock-hour-12` icon.\n CLOCK_HOUR_2 (IconName): The :code:`clock-hour-2` icon.\n CLOCK_HOUR_3 (IconName): The :code:`clock-hour-3` icon.\n CLOCK_HOUR_4 (IconName): The :code:`clock-hour-4` icon.\n CLOCK_HOUR_5 (IconName): The :code:`clock-hour-5` icon.\n CLOCK_HOUR_6 (IconName): The :code:`clock-hour-6` icon.\n CLOCK_HOUR_7 (IconName): The :code:`clock-hour-7` icon.\n CLOCK_HOUR_8 (IconName): The :code:`clock-hour-8` icon.\n CLOCK_HOUR_9 (IconName): The :code:`clock-hour-9` icon.\n CLOCK_MINUS (IconName): The :code:`clock-minus` icon.\n CLOCK_OFF (IconName): The :code:`clock-off` icon.\n CLOCK_PAUSE (IconName): The :code:`clock-pause` icon.\n CLOCK_PIN (IconName): The :code:`clock-pin` icon.\n CLOCK_PLAY (IconName): The :code:`clock-play` icon.\n CLOCK_PLUS (IconName): The :code:`clock-plus` icon.\n CLOCK_QUESTION (IconName): The :code:`clock-question` icon.\n CLOCK_RECORD (IconName): The :code:`clock-record` icon.\n CLOCK_SEARCH (IconName): The :code:`clock-search` icon.\n CLOCK_SHARE (IconName): The :code:`clock-share` icon.\n CLOCK_SHIELD (IconName): The :code:`clock-shield` icon.\n CLOCK_STAR (IconName): The :code:`clock-star` icon.\n CLOCK_STOP (IconName): The :code:`clock-stop` icon.\n CLOCK_UP (IconName): The :code:`clock-up` icon.\n CLOCK_X (IconName): The :code:`clock-x` icon.\n CLOTHES_RACK (IconName): The :code:`clothes-rack` icon.\n CLOTHES_RACK_OFF (IconName): The :code:`clothes-rack-off` icon.\n CLOUD (IconName): The :code:`cloud` icon.\n CLOUD_BOLT (IconName): The :code:`cloud-bolt` icon.\n CLOUD_CANCEL (IconName): The :code:`cloud-cancel` icon.\n CLOUD_CHECK (IconName): The :code:`cloud-check` icon.\n CLOUD_CODE (IconName): The :code:`cloud-code` icon.\n CLOUD_COG (IconName): The :code:`cloud-cog` icon.\n CLOUD_COMPUTING (IconName): The :code:`cloud-computing` icon.\n CLOUD_DATA_CONNECTION (IconName): The :code:`cloud-data-connection` icon.\n CLOUD_DOLLAR (IconName): The :code:`cloud-dollar` icon.\n CLOUD_DOWN (IconName): The :code:`cloud-down` icon.\n CLOUD_DOWNLOAD (IconName): The :code:`cloud-download` icon.\n CLOUD_EXCLAMATION (IconName): The :code:`cloud-exclamation` icon.\n CLOUD_FILLED (IconName): The :code:`cloud-filled` icon.\n CLOUD_FOG (IconName): The :code:`cloud-fog` icon.\n CLOUD_HEART (IconName): The :code:`cloud-heart` icon.\n CLOUD_LOCK (IconName): The :code:`cloud-lock` icon.\n CLOUD_LOCK_OPEN (IconName): The :code:`cloud-lock-open` icon.\n CLOUD_MINUS (IconName): The :code:`cloud-minus` icon.\n CLOUD_OFF (IconName): The :code:`cloud-off` icon.\n CLOUD_PAUSE (IconName): The :code:`cloud-pause` icon.\n CLOUD_PIN (IconName): The :code:`cloud-pin` icon.\n CLOUD_PLUS (IconName): The :code:`cloud-plus` icon.\n CLOUD_QUESTION (IconName): The :code:`cloud-question` icon.\n CLOUD_RAIN (IconName): The :code:`cloud-rain` icon.\n CLOUD_SEARCH (IconName): The :code:`cloud-search` icon.\n CLOUD_SHARE (IconName): The :code:`cloud-share` icon.\n CLOUD_SNOW (IconName): The :code:`cloud-snow` icon.\n CLOUD_STAR (IconName): The :code:`cloud-star` icon.\n CLOUD_STORM (IconName): The :code:`cloud-storm` icon.\n CLOUD_UP (IconName): The :code:`cloud-up` icon.\n CLOUD_UPLOAD (IconName): The :code:`cloud-upload` icon.\n CLOUD_X (IconName): The :code:`cloud-x` icon.\n CLOVER (IconName): The :code:`clover` icon.\n CLOVER_2 (IconName): The :code:`clover-2` icon.\n CLUBS (IconName): The :code:`clubs` icon.\n CLUBS_FILLED (IconName): The :code:`clubs-filled` icon.\n CODE (IconName): The :code:`code` icon.\n CODE_ASTERIX (IconName): The :code:`code-asterix` icon.\n CODE_CIRCLE (IconName): The :code:`code-circle` icon.\n CODE_CIRCLE_2 (IconName): The :code:`code-circle-2` icon.\n CODE_DOTS (IconName): The :code:`code-dots` icon.\n CODE_MINUS (IconName): The :code:`code-minus` icon.\n CODE_OFF (IconName): The :code:`code-off` icon.\n CODE_PLUS (IconName): The :code:`code-plus` icon.\n COFFEE (IconName): The :code:`coffee` icon.\n COFFEE_OFF (IconName): The :code:`coffee-off` icon.\n COFFIN (IconName): The :code:`coffin` icon.\n COIN (IconName): The :code:`coin` icon.\n COIN_BITCOIN (IconName): The :code:`coin-bitcoin` icon.\n COIN_EURO (IconName): The :code:`coin-euro` icon.\n COIN_MONERO (IconName): The :code:`coin-monero` icon.\n COIN_OFF (IconName): The :code:`coin-off` icon.\n COIN_POUND (IconName): The :code:`coin-pound` icon.\n COIN_RUPEE (IconName): The :code:`coin-rupee` icon.\n COIN_YEN (IconName): The :code:`coin-yen` icon.\n COIN_YUAN (IconName): The :code:`coin-yuan` icon.\n COINS (IconName): The :code:`coins` icon.\n COLOR_FILTER (IconName): The :code:`color-filter` icon.\n COLOR_PICKER (IconName): The :code:`color-picker` icon.\n COLOR_PICKER_OFF (IconName): The :code:`color-picker-off` icon.\n COLOR_SWATCH (IconName): The :code:`color-swatch` icon.\n COLOR_SWATCH_OFF (IconName): The :code:`color-swatch-off` icon.\n COLUMN_INSERT_LEFT (IconName): The :code:`column-insert-left` icon.\n COLUMN_INSERT_RIGHT (IconName): The :code:`column-insert-right` icon.\n COLUMN_REMOVE (IconName): The :code:`column-remove` icon.\n COLUMNS (IconName): The :code:`columns` icon.\n COLUMNS_1 (IconName): The :code:`columns-1` icon.\n COLUMNS_2 (IconName): The :code:`columns-2` icon.\n COLUMNS_3 (IconName): The :code:`columns-3` icon.\n COLUMNS_OFF (IconName): The :code:`columns-off` icon.\n COMET (IconName): The :code:`comet` icon.\n COMMAND (IconName): The :code:`command` icon.\n COMMAND_OFF (IconName): The :code:`command-off` icon.\n COMPASS (IconName): The :code:`compass` icon.\n COMPASS_OFF (IconName): The :code:`compass-off` icon.\n COMPONENTS (IconName): The :code:`components` icon.\n COMPONENTS_OFF (IconName): The :code:`components-off` icon.\n CONE (IconName): The :code:`cone` icon.\n CONE_2 (IconName): The :code:`cone-2` icon.\n CONE_OFF (IconName): The :code:`cone-off` icon.\n CONE_PLUS (IconName): The :code:`cone-plus` icon.\n CONFETTI (IconName): The :code:`confetti` icon.\n CONFETTI_OFF (IconName): The :code:`confetti-off` icon.\n CONFUCIUS (IconName): The :code:`confucius` icon.\n CONTAINER (IconName): The :code:`container` icon.\n CONTAINER_OFF (IconName): The :code:`container-off` icon.\n CONTRAST (IconName): The :code:`contrast` icon.\n CONTRAST_2 (IconName): The :code:`contrast-2` icon.\n CONTRAST_2_OFF (IconName): The :code:`contrast-2-off` icon.\n CONTRAST_OFF (IconName): The :code:`contrast-off` icon.\n COOKER (IconName): The :code:`cooker` icon.\n COOKIE (IconName): The :code:`cookie` icon.\n COOKIE_MAN (IconName): The :code:`cookie-man` icon.\n COOKIE_OFF (IconName): The :code:`cookie-off` icon.\n COPY (IconName): The :code:`copy` icon.\n COPY_OFF (IconName): The :code:`copy-off` icon.\n COPYLEFT (IconName): The :code:`copyleft` icon.\n COPYLEFT_FILLED (IconName): The :code:`copyleft-filled` icon.\n COPYLEFT_OFF (IconName): The :code:`copyleft-off` icon.\n COPYRIGHT (IconName): The :code:`copyright` icon.\n COPYRIGHT_FILLED (IconName): The :code:`copyright-filled` icon.\n COPYRIGHT_OFF (IconName): The :code:`copyright-off` icon.\n CORNER_DOWN_LEFT (IconName): The :code:`corner-down-left` icon.\n CORNER_DOWN_LEFT_DOUBLE (IconName): The :code:`corner-down-left-double` icon.\n CORNER_DOWN_RIGHT (IconName): The :code:`corner-down-right` icon.\n CORNER_DOWN_RIGHT_DOUBLE (IconName): The :code:`corner-down-right-double` icon.\n CORNER_LEFT_DOWN (IconName): The :code:`corner-left-down` icon.\n CORNER_LEFT_DOWN_DOUBLE (IconName): The :code:`corner-left-down-double` icon.\n CORNER_LEFT_UP (IconName): The :code:`corner-left-up` icon.\n CORNER_LEFT_UP_DOUBLE (IconName): The :code:`corner-left-up-double` icon.\n CORNER_RIGHT_DOWN (IconName): The :code:`corner-right-down` icon.\n CORNER_RIGHT_DOWN_DOUBLE (IconName): The :code:`corner-right-down-double` icon.\n CORNER_RIGHT_UP (IconName): The :code:`corner-right-up` icon.\n CORNER_RIGHT_UP_DOUBLE (IconName): The :code:`corner-right-up-double` icon.\n CORNER_UP_LEFT (IconName): The :code:`corner-up-left` icon.\n CORNER_UP_LEFT_DOUBLE (IconName): The :code:`corner-up-left-double` icon.\n CORNER_UP_RIGHT (IconName): The :code:`corner-up-right` icon.\n CORNER_UP_RIGHT_DOUBLE (IconName): The :code:`corner-up-right-double` icon.\n CPU (IconName): The :code:`cpu` icon.\n CPU_2 (IconName): The :code:`cpu-2` icon.\n CPU_OFF (IconName): The :code:`cpu-off` icon.\n CRANE (IconName): The :code:`crane` icon.\n CRANE_OFF (IconName): The :code:`crane-off` icon.\n CREATIVE_COMMONS (IconName): The :code:`creative-commons` icon.\n CREATIVE_COMMONS_BY (IconName): The :code:`creative-commons-by` icon.\n CREATIVE_COMMONS_NC (IconName): The :code:`creative-commons-nc` icon.\n CREATIVE_COMMONS_ND (IconName): The :code:`creative-commons-nd` icon.\n CREATIVE_COMMONS_OFF (IconName): The :code:`creative-commons-off` icon.\n CREATIVE_COMMONS_SA (IconName): The :code:`creative-commons-sa` icon.\n CREATIVE_COMMONS_ZERO (IconName): The :code:`creative-commons-zero` icon.\n CREDIT_CARD (IconName): The :code:`credit-card` icon.\n CREDIT_CARD_OFF (IconName): The :code:`credit-card-off` icon.\n CRICKET (IconName): The :code:`cricket` icon.\n CROP (IconName): The :code:`crop` icon.\n CROSS (IconName): The :code:`cross` icon.\n CROSS_FILLED (IconName): The :code:`cross-filled` icon.\n CROSS_OFF (IconName): The :code:`cross-off` icon.\n CROSSHAIR (IconName): The :code:`crosshair` icon.\n CROWN (IconName): The :code:`crown` icon.\n CROWN_OFF (IconName): The :code:`crown-off` icon.\n CRUTCHES (IconName): The :code:`crutches` icon.\n CRUTCHES_OFF (IconName): The :code:`crutches-off` icon.\n CRYSTAL_BALL (IconName): The :code:`crystal-ball` icon.\n CSV (IconName): The :code:`csv` icon.\n CUBE (IconName): The :code:`cube` icon.\n CUBE_OFF (IconName): The :code:`cube-off` icon.\n CUBE_PLUS (IconName): The :code:`cube-plus` icon.\n CUBE_SEND (IconName): The :code:`cube-send` icon.\n CUBE_UNFOLDED (IconName): The :code:`cube-unfolded` icon.\n CUP (IconName): The :code:`cup` icon.\n CUP_OFF (IconName): The :code:`cup-off` icon.\n CURLING (IconName): The :code:`curling` icon.\n CURLY_LOOP (IconName): The :code:`curly-loop` icon.\n CURRENCY (IconName): The :code:`currency` icon.\n CURRENCY_AFGHANI (IconName): The :code:`currency-afghani` icon.\n CURRENCY_BAHRAINI (IconName): The :code:`currency-bahraini` icon.\n CURRENCY_BAHT (IconName): The :code:`currency-baht` icon.\n CURRENCY_BITCOIN (IconName): The :code:`currency-bitcoin` icon.\n CURRENCY_CENT (IconName): The :code:`currency-cent` icon.\n CURRENCY_DINAR (IconName): The :code:`currency-dinar` icon.\n CURRENCY_DIRHAM (IconName): The :code:`currency-dirham` icon.\n CURRENCY_DOGECOIN (IconName): The :code:`currency-dogecoin` icon.\n CURRENCY_DOLLAR (IconName): The :code:`currency-dollar` icon.\n CURRENCY_DOLLAR_AUSTRALIAN (IconName): The :code:`currency-dollar-australian` icon.\n CURRENCY_DOLLAR_BRUNEI (IconName): The :code:`currency-dollar-brunei` icon.\n CURRENCY_DOLLAR_CANADIAN (IconName): The :code:`currency-dollar-canadian` icon.\n CURRENCY_DOLLAR_GUYANESE (IconName): The :code:`currency-dollar-guyanese` icon.\n CURRENCY_DOLLAR_OFF (IconName): The :code:`currency-dollar-off` icon.\n CURRENCY_DOLLAR_SINGAPORE (IconName): The :code:`currency-dollar-singapore` icon.\n CURRENCY_DOLLAR_ZIMBABWEAN (IconName): The :code:`currency-dollar-zimbabwean` icon.\n CURRENCY_DONG (IconName): The :code:`currency-dong` icon.\n CURRENCY_DRAM (IconName): The :code:`currency-dram` icon.\n CURRENCY_ETHEREUM (IconName): The :code:`currency-ethereum` icon.\n CURRENCY_EURO (IconName): The :code:`currency-euro` icon.\n CURRENCY_EURO_OFF (IconName): The :code:`currency-euro-off` icon.\n CURRENCY_FLORIN (IconName): The :code:`currency-florin` icon.\n CURRENCY_FORINT (IconName): The :code:`currency-forint` icon.\n CURRENCY_FRANK (IconName): The :code:`currency-frank` icon.\n CURRENCY_GUARANI (IconName): The :code:`currency-guarani` icon.\n CURRENCY_HRYVNIA (IconName): The :code:`currency-hryvnia` icon.\n CURRENCY_IRANIAN_RIAL (IconName): The :code:`currency-iranian-rial` icon.\n CURRENCY_KIP (IconName): The :code:`currency-kip` icon.\n CURRENCY_KRONE_CZECH (IconName): The :code:`currency-krone-czech` icon.\n CURRENCY_KRONE_DANISH (IconName): The :code:`currency-krone-danish` icon.\n CURRENCY_KRONE_SWEDISH (IconName): The :code:`currency-krone-swedish` icon.\n CURRENCY_LARI (IconName): The :code:`currency-lari` icon.\n CURRENCY_LEU (IconName): The :code:`currency-leu` icon.\n CURRENCY_LIRA (IconName): The :code:`currency-lira` icon.\n CURRENCY_LITECOIN (IconName): The :code:`currency-litecoin` icon.\n CURRENCY_LYD (IconName): The :code:`currency-lyd` icon.\n CURRENCY_MANAT (IconName): The :code:`currency-manat` icon.\n CURRENCY_MONERO (IconName): The :code:`currency-monero` icon.\n CURRENCY_NAIRA (IconName): The :code:`currency-naira` icon.\n CURRENCY_NANO (IconName): The :code:`currency-nano` icon.\n CURRENCY_OFF (IconName): The :code:`currency-off` icon.\n CURRENCY_PAANGA (IconName): The :code:`currency-paanga` icon.\n CURRENCY_PESO (IconName): The :code:`currency-peso` icon.\n CURRENCY_POUND (IconName): The :code:`currency-pound` icon.\n CURRENCY_POUND_OFF (IconName): The :code:`currency-pound-off` icon.\n CURRENCY_QUETZAL (IconName): The :code:`currency-quetzal` icon.\n CURRENCY_REAL (IconName): The :code:`currency-real` icon.\n CURRENCY_RENMINBI (IconName): The :code:`currency-renminbi` icon.\n CURRENCY_RIPPLE (IconName): The :code:`currency-ripple` icon.\n CURRENCY_RIYAL (IconName): The :code:`currency-riyal` icon.\n CURRENCY_RUBEL (IconName): The :code:`currency-rubel` icon.\n CURRENCY_RUFIYAA (IconName): The :code:`currency-rufiyaa` icon.\n CURRENCY_RUPEE (IconName): The :code:`currency-rupee` icon.\n CURRENCY_RUPEE_NEPALESE (IconName): The :code:`currency-rupee-nepalese` icon.\n CURRENCY_SHEKEL (IconName): The :code:`currency-shekel` icon.\n CURRENCY_SOLANA (IconName): The :code:`currency-solana` icon.\n CURRENCY_SOM (IconName): The :code:`currency-som` icon.\n CURRENCY_TAKA (IconName): The :code:`currency-taka` icon.\n CURRENCY_TENGE (IconName): The :code:`currency-tenge` icon.\n CURRENCY_TUGRIK (IconName): The :code:`currency-tugrik` icon.\n CURRENCY_WON (IconName): The :code:`currency-won` icon.\n CURRENCY_YEN (IconName): The :code:`currency-yen` icon.\n CURRENCY_YEN_OFF (IconName): The :code:`currency-yen-off` icon.\n CURRENCY_YUAN (IconName): The :code:`currency-yuan` icon.\n CURRENCY_ZLOTY (IconName): The :code:`currency-zloty` icon.\n CURRENT_LOCATION (IconName): The :code:`current-location` icon.\n CURRENT_LOCATION_OFF (IconName): The :code:`current-location-off` icon.\n CURSOR_OFF (IconName): The :code:`cursor-off` icon.\n CURSOR_TEXT (IconName): The :code:`cursor-text` icon.\n CUT (IconName): The :code:`cut` icon.\n CYLINDER (IconName): The :code:`cylinder` icon.\n CYLINDER_OFF (IconName): The :code:`cylinder-off` icon.\n CYLINDER_PLUS (IconName): The :code:`cylinder-plus` icon.\n DASHBOARD (IconName): The :code:`dashboard` icon.\n DASHBOARD_OFF (IconName): The :code:`dashboard-off` icon.\n DATABASE (IconName): The :code:`database` icon.\n DATABASE_COG (IconName): The :code:`database-cog` icon.\n DATABASE_DOLLAR (IconName): The :code:`database-dollar` icon.\n DATABASE_EDIT (IconName): The :code:`database-edit` icon.\n DATABASE_EXCLAMATION (IconName): The :code:`database-exclamation` icon.\n DATABASE_EXPORT (IconName): The :code:`database-export` icon.\n DATABASE_HEART (IconName): The :code:`database-heart` icon.\n DATABASE_IMPORT (IconName): The :code:`database-import` icon.\n DATABASE_LEAK (IconName): The :code:`database-leak` icon.\n DATABASE_MINUS (IconName): The :code:`database-minus` icon.\n DATABASE_OFF (IconName): The :code:`database-off` icon.\n DATABASE_PLUS (IconName): The :code:`database-plus` icon.\n DATABASE_SEARCH (IconName): The :code:`database-search` icon.\n DATABASE_SHARE (IconName): The :code:`database-share` icon.\n DATABASE_STAR (IconName): The :code:`database-star` icon.\n DATABASE_X (IconName): The :code:`database-x` icon.\n DECIMAL (IconName): The :code:`decimal` icon.\n DEER (IconName): The :code:`deer` icon.\n DELTA (IconName): The :code:`delta` icon.\n DENTAL (IconName): The :code:`dental` icon.\n DENTAL_BROKEN (IconName): The :code:`dental-broken` icon.\n DENTAL_OFF (IconName): The :code:`dental-off` icon.\n DESELECT (IconName): The :code:`deselect` icon.\n DETAILS (IconName): The :code:`details` icon.\n DETAILS_OFF (IconName): The :code:`details-off` icon.\n DEVICE_AIRPODS (IconName): The :code:`device-airpods` icon.\n DEVICE_AIRPODS_CASE (IconName): The :code:`device-airpods-case` icon.\n DEVICE_AIRTAG (IconName): The :code:`device-airtag` icon.\n DEVICE_ANALYTICS (IconName): The :code:`device-analytics` icon.\n DEVICE_AUDIO_TAPE (IconName): The :code:`device-audio-tape` icon.\n DEVICE_CAMERA_PHONE (IconName): The :code:`device-camera-phone` icon.\n DEVICE_CCTV (IconName): The :code:`device-cctv` icon.\n DEVICE_CCTV_OFF (IconName): The :code:`device-cctv-off` icon.\n DEVICE_COMPUTER_CAMERA (IconName): The :code:`device-computer-camera` icon.\n DEVICE_COMPUTER_CAMERA_OFF (IconName): The :code:`device-computer-camera-off` icon.\n DEVICE_DESKTOP (IconName): The :code:`device-desktop` icon.\n DEVICE_DESKTOP_ANALYTICS (IconName): The :code:`device-desktop-analytics` icon.\n DEVICE_DESKTOP_BOLT (IconName): The :code:`device-desktop-bolt` icon.\n DEVICE_DESKTOP_CANCEL (IconName): The :code:`device-desktop-cancel` icon.\n DEVICE_DESKTOP_CHECK (IconName): The :code:`device-desktop-check` icon.\n DEVICE_DESKTOP_CODE (IconName): The :code:`device-desktop-code` icon.\n DEVICE_DESKTOP_COG (IconName): The :code:`device-desktop-cog` icon.\n DEVICE_DESKTOP_DOLLAR (IconName): The :code:`device-desktop-dollar` icon.\n DEVICE_DESKTOP_DOWN (IconName): The :code:`device-desktop-down` icon.\n DEVICE_DESKTOP_EXCLAMATION (IconName): The :code:`device-desktop-exclamation` icon.\n DEVICE_DESKTOP_HEART (IconName): The :code:`device-desktop-heart` icon.\n DEVICE_DESKTOP_MINUS (IconName): The :code:`device-desktop-minus` icon.\n DEVICE_DESKTOP_OFF (IconName): The :code:`device-desktop-off` icon.\n DEVICE_DESKTOP_PAUSE (IconName): The :code:`device-desktop-pause` icon.\n DEVICE_DESKTOP_PIN (IconName): The :code:`device-desktop-pin` icon.\n DEVICE_DESKTOP_PLUS (IconName): The :code:`device-desktop-plus` icon.\n DEVICE_DESKTOP_QUESTION (IconName): The :code:`device-desktop-question` icon.\n DEVICE_DESKTOP_SEARCH (IconName): The :code:`device-desktop-search` icon.\n DEVICE_DESKTOP_SHARE (IconName): The :code:`device-desktop-share` icon.\n DEVICE_DESKTOP_STAR (IconName): The :code:`device-desktop-star` icon.\n DEVICE_DESKTOP_UP (IconName): The :code:`device-desktop-up` icon.\n DEVICE_DESKTOP_X (IconName): The :code:`device-desktop-x` icon.\n DEVICE_FLOPPY (IconName): The :code:`device-floppy` icon.\n DEVICE_GAMEPAD (IconName): The :code:`device-gamepad` icon.\n DEVICE_GAMEPAD_2 (IconName): The :code:`device-gamepad-2` icon.\n DEVICE_HEART_MONITOR (IconName): The :code:`device-heart-monitor` icon.\n DEVICE_HEART_MONITOR_FILLED (IconName): The :code:`device-heart-monitor-filled` icon.\n DEVICE_IMAC (IconName): The :code:`device-imac` icon.\n DEVICE_IMAC_BOLT (IconName): The :code:`device-imac-bolt` icon.\n DEVICE_IMAC_CANCEL (IconName): The :code:`device-imac-cancel` icon.\n DEVICE_IMAC_CHECK (IconName): The :code:`device-imac-check` icon.\n DEVICE_IMAC_CODE (IconName): The :code:`device-imac-code` icon.\n DEVICE_IMAC_COG (IconName): The :code:`device-imac-cog` icon.\n DEVICE_IMAC_DOLLAR (IconName): The :code:`device-imac-dollar` icon.\n DEVICE_IMAC_DOWN (IconName): The :code:`device-imac-down` icon.\n DEVICE_IMAC_EXCLAMATION (IconName): The :code:`device-imac-exclamation` icon.\n DEVICE_IMAC_HEART (IconName): The :code:`device-imac-heart` icon.\n DEVICE_IMAC_MINUS (IconName): The :code:`device-imac-minus` icon.\n DEVICE_IMAC_OFF (IconName): The :code:`device-imac-off` icon.\n DEVICE_IMAC_PAUSE (IconName): The :code:`device-imac-pause` icon.\n DEVICE_IMAC_PIN (IconName): The :code:`device-imac-pin` icon.\n DEVICE_IMAC_PLUS (IconName): The :code:`device-imac-plus` icon.\n DEVICE_IMAC_QUESTION (IconName): The :code:`device-imac-question` icon.\n DEVICE_IMAC_SEARCH (IconName): The :code:`device-imac-search` icon.\n DEVICE_IMAC_SHARE (IconName): The :code:`device-imac-share` icon.\n DEVICE_IMAC_STAR (IconName): The :code:`device-imac-star` icon.\n DEVICE_IMAC_UP (IconName): The :code:`device-imac-up` icon.\n DEVICE_IMAC_X (IconName): The :code:`device-imac-x` icon.\n DEVICE_IPAD (IconName): The :code:`device-ipad` icon.\n DEVICE_IPAD_BOLT (IconName): The :code:`device-ipad-bolt` icon.\n DEVICE_IPAD_CANCEL (IconName): The :code:`device-ipad-cancel` icon.\n DEVICE_IPAD_CHECK (IconName): The :code:`device-ipad-check` icon.\n DEVICE_IPAD_CODE (IconName): The :code:`device-ipad-code` icon.\n DEVICE_IPAD_COG (IconName): The :code:`device-ipad-cog` icon.\n DEVICE_IPAD_DOLLAR (IconName): The :code:`device-ipad-dollar` icon.\n DEVICE_IPAD_DOWN (IconName): The :code:`device-ipad-down` icon.\n DEVICE_IPAD_EXCLAMATION (IconName): The :code:`device-ipad-exclamation` icon.\n DEVICE_IPAD_HEART (IconName): The :code:`device-ipad-heart` icon.\n DEVICE_IPAD_HORIZONTAL (IconName): The :code:`device-ipad-horizontal` icon.\n DEVICE_IPAD_HORIZONTAL_BOLT (IconName): The :code:`device-ipad-horizontal-bolt` icon.\n DEVICE_IPAD_HORIZONTAL_CANCEL (IconName): The :code:`device-ipad-horizontal-cancel` icon.\n DEVICE_IPAD_HORIZONTAL_CHECK (IconName): The :code:`device-ipad-horizontal-check` icon.\n DEVICE_IPAD_HORIZONTAL_CODE (IconName): The :code:`device-ipad-horizontal-code` icon.\n DEVICE_IPAD_HORIZONTAL_COG (IconName): The :code:`device-ipad-horizontal-cog` icon.\n DEVICE_IPAD_HORIZONTAL_DOLLAR (IconName): The :code:`device-ipad-horizontal-dollar` icon.\n DEVICE_IPAD_HORIZONTAL_DOWN (IconName): The :code:`device-ipad-horizontal-down` icon.\n DEVICE_IPAD_HORIZONTAL_EXCLAMATION (IconName): The :code:`device-ipad-horizontal-exclamation` icon.\n DEVICE_IPAD_HORIZONTAL_HEART (IconName): The :code:`device-ipad-horizontal-heart` icon.\n DEVICE_IPAD_HORIZONTAL_MINUS (IconName): The :code:`device-ipad-horizontal-minus` icon.\n DEVICE_IPAD_HORIZONTAL_OFF (IconName): The :code:`device-ipad-horizontal-off` icon.\n DEVICE_IPAD_HORIZONTAL_PAUSE (IconName): The :code:`device-ipad-horizontal-pause` icon.\n DEVICE_IPAD_HORIZONTAL_PIN (IconName): The :code:`device-ipad-horizontal-pin` icon.\n DEVICE_IPAD_HORIZONTAL_PLUS (IconName): The :code:`device-ipad-horizontal-plus` icon.\n DEVICE_IPAD_HORIZONTAL_QUESTION (IconName): The :code:`device-ipad-horizontal-question` icon.\n DEVICE_IPAD_HORIZONTAL_SEARCH (IconName): The :code:`device-ipad-horizontal-search` icon.\n DEVICE_IPAD_HORIZONTAL_SHARE (IconName): The :code:`device-ipad-horizontal-share` icon.\n DEVICE_IPAD_HORIZONTAL_STAR (IconName): The :code:`device-ipad-horizontal-star` icon.\n DEVICE_IPAD_HORIZONTAL_UP (IconName): The :code:`device-ipad-horizontal-up` icon.\n DEVICE_IPAD_HORIZONTAL_X (IconName): The :code:`device-ipad-horizontal-x` icon.\n DEVICE_IPAD_MINUS (IconName): The :code:`device-ipad-minus` icon.\n DEVICE_IPAD_OFF (IconName): The :code:`device-ipad-off` icon.\n DEVICE_IPAD_PAUSE (IconName): The :code:`device-ipad-pause` icon.\n DEVICE_IPAD_PIN (IconName): The :code:`device-ipad-pin` icon.\n DEVICE_IPAD_PLUS (IconName): The :code:`device-ipad-plus` icon.\n DEVICE_IPAD_QUESTION (IconName): The :code:`device-ipad-question` icon.\n DEVICE_IPAD_SEARCH (IconName): The :code:`device-ipad-search` icon.\n DEVICE_IPAD_SHARE (IconName): The :code:`device-ipad-share` icon.\n DEVICE_IPAD_STAR (IconName): The :code:`device-ipad-star` icon.\n DEVICE_IPAD_UP (IconName): The :code:`device-ipad-up` icon.\n DEVICE_IPAD_X (IconName): The :code:`device-ipad-x` icon.\n DEVICE_LANDLINE_PHONE (IconName): The :code:`device-landline-phone` icon.\n DEVICE_LAPTOP (IconName): The :code:`device-laptop` icon.\n DEVICE_LAPTOP_OFF (IconName): The :code:`device-laptop-off` icon.\n DEVICE_MOBILE (IconName): The :code:`device-mobile` icon.\n DEVICE_MOBILE_BOLT (IconName): The :code:`device-mobile-bolt` icon.\n DEVICE_MOBILE_CANCEL (IconName): The :code:`device-mobile-cancel` icon.\n DEVICE_MOBILE_CHARGING (IconName): The :code:`device-mobile-charging` icon.\n DEVICE_MOBILE_CHECK (IconName): The :code:`device-mobile-check` icon.\n DEVICE_MOBILE_CODE (IconName): The :code:`device-mobile-code` icon.\n DEVICE_MOBILE_COG (IconName): The :code:`device-mobile-cog` icon.\n DEVICE_MOBILE_DOLLAR (IconName): The :code:`device-mobile-dollar` icon.\n DEVICE_MOBILE_DOWN (IconName): The :code:`device-mobile-down` icon.\n DEVICE_MOBILE_EXCLAMATION (IconName): The :code:`device-mobile-exclamation` icon.\n DEVICE_MOBILE_FILLED (IconName): The :code:`device-mobile-filled` icon.\n DEVICE_MOBILE_HEART (IconName): The :code:`device-mobile-heart` icon.\n DEVICE_MOBILE_MESSAGE (IconName): The :code:`device-mobile-message` icon.\n DEVICE_MOBILE_MINUS (IconName): The :code:`device-mobile-minus` icon.\n DEVICE_MOBILE_OFF (IconName): The :code:`device-mobile-off` icon.\n DEVICE_MOBILE_PAUSE (IconName): The :code:`device-mobile-pause` icon.\n DEVICE_MOBILE_PIN (IconName): The :code:`device-mobile-pin` icon.\n DEVICE_MOBILE_PLUS (IconName): The :code:`device-mobile-plus` icon.\n DEVICE_MOBILE_QUESTION (IconName): The :code:`device-mobile-question` icon.\n DEVICE_MOBILE_ROTATED (IconName): The :code:`device-mobile-rotated` icon.\n DEVICE_MOBILE_SEARCH (IconName): The :code:`device-mobile-search` icon.\n DEVICE_MOBILE_SHARE (IconName): The :code:`device-mobile-share` icon.\n DEVICE_MOBILE_STAR (IconName): The :code:`device-mobile-star` icon.\n DEVICE_MOBILE_UP (IconName): The :code:`device-mobile-up` icon.\n DEVICE_MOBILE_VIBRATION (IconName): The :code:`device-mobile-vibration` icon.\n DEVICE_MOBILE_X (IconName): The :code:`device-mobile-x` icon.\n DEVICE_NINTENDO (IconName): The :code:`device-nintendo` icon.\n DEVICE_NINTENDO_OFF (IconName): The :code:`device-nintendo-off` icon.\n DEVICE_REMOTE (IconName): The :code:`device-remote` icon.\n DEVICE_SD_CARD (IconName): The :code:`device-sd-card` icon.\n DEVICE_SIM (IconName): The :code:`device-sim` icon.\n DEVICE_SIM_1 (IconName): The :code:`device-sim-1` icon.\n DEVICE_SIM_2 (IconName): The :code:`device-sim-2` icon.\n DEVICE_SIM_3 (IconName): The :code:`device-sim-3` icon.\n DEVICE_SPEAKER (IconName): The :code:`device-speaker` icon.\n DEVICE_SPEAKER_OFF (IconName): The :code:`device-speaker-off` icon.\n DEVICE_TABLET (IconName): The :code:`device-tablet` icon.\n DEVICE_TABLET_BOLT (IconName): The :code:`device-tablet-bolt` icon.\n DEVICE_TABLET_CANCEL (IconName): The :code:`device-tablet-cancel` icon.\n DEVICE_TABLET_CHECK (IconName): The :code:`device-tablet-check` icon.\n DEVICE_TABLET_CODE (IconName): The :code:`device-tablet-code` icon.\n DEVICE_TABLET_COG (IconName): The :code:`device-tablet-cog` icon.\n DEVICE_TABLET_DOLLAR (IconName): The :code:`device-tablet-dollar` icon.\n DEVICE_TABLET_DOWN (IconName): The :code:`device-tablet-down` icon.\n DEVICE_TABLET_EXCLAMATION (IconName): The :code:`device-tablet-exclamation` icon.\n DEVICE_TABLET_FILLED (IconName): The :code:`device-tablet-filled` icon.\n DEVICE_TABLET_HEART (IconName): The :code:`device-tablet-heart` icon.\n DEVICE_TABLET_MINUS (IconName): The :code:`device-tablet-minus` icon.\n DEVICE_TABLET_OFF (IconName): The :code:`device-tablet-off` icon.\n DEVICE_TABLET_PAUSE (IconName): The :code:`device-tablet-pause` icon.\n DEVICE_TABLET_PIN (IconName): The :code:`device-tablet-pin` icon.\n DEVICE_TABLET_PLUS (IconName): The :code:`device-tablet-plus` icon.\n DEVICE_TABLET_QUESTION (IconName): The :code:`device-tablet-question` icon.\n DEVICE_TABLET_SEARCH (IconName): The :code:`device-tablet-search` icon.\n DEVICE_TABLET_SHARE (IconName): The :code:`device-tablet-share` icon.\n DEVICE_TABLET_STAR (IconName): The :code:`device-tablet-star` icon.\n DEVICE_TABLET_UP (IconName): The :code:`device-tablet-up` icon.\n DEVICE_TABLET_X (IconName): The :code:`device-tablet-x` icon.\n DEVICE_TV (IconName): The :code:`device-tv` icon.\n DEVICE_TV_OFF (IconName): The :code:`device-tv-off` icon.\n DEVICE_TV_OLD (IconName): The :code:`device-tv-old` icon.\n DEVICE_VISION_PRO (IconName): The :code:`device-vision-pro` icon.\n DEVICE_WATCH (IconName): The :code:`device-watch` icon.\n DEVICE_WATCH_BOLT (IconName): The :code:`device-watch-bolt` icon.\n DEVICE_WATCH_CANCEL (IconName): The :code:`device-watch-cancel` icon.\n DEVICE_WATCH_CHECK (IconName): The :code:`device-watch-check` icon.\n DEVICE_WATCH_CODE (IconName): The :code:`device-watch-code` icon.\n DEVICE_WATCH_COG (IconName): The :code:`device-watch-cog` icon.\n DEVICE_WATCH_DOLLAR (IconName): The :code:`device-watch-dollar` icon.\n DEVICE_WATCH_DOWN (IconName): The :code:`device-watch-down` icon.\n DEVICE_WATCH_EXCLAMATION (IconName): The :code:`device-watch-exclamation` icon.\n DEVICE_WATCH_HEART (IconName): The :code:`device-watch-heart` icon.\n DEVICE_WATCH_MINUS (IconName): The :code:`device-watch-minus` icon.\n DEVICE_WATCH_OFF (IconName): The :code:`device-watch-off` icon.\n DEVICE_WATCH_PAUSE (IconName): The :code:`device-watch-pause` icon.\n DEVICE_WATCH_PIN (IconName): The :code:`device-watch-pin` icon.\n DEVICE_WATCH_PLUS (IconName): The :code:`device-watch-plus` icon.\n DEVICE_WATCH_QUESTION (IconName): The :code:`device-watch-question` icon.\n DEVICE_WATCH_SEARCH (IconName): The :code:`device-watch-search` icon.\n DEVICE_WATCH_SHARE (IconName): The :code:`device-watch-share` icon.\n DEVICE_WATCH_STAR (IconName): The :code:`device-watch-star` icon.\n DEVICE_WATCH_STATS (IconName): The :code:`device-watch-stats` icon.\n DEVICE_WATCH_STATS_2 (IconName): The :code:`device-watch-stats-2` icon.\n DEVICE_WATCH_UP (IconName): The :code:`device-watch-up` icon.\n DEVICE_WATCH_X (IconName): The :code:`device-watch-x` icon.\n DEVICES (IconName): The :code:`devices` icon.\n DEVICES_2 (IconName): The :code:`devices-2` icon.\n DEVICES_BOLT (IconName): The :code:`devices-bolt` icon.\n DEVICES_CANCEL (IconName): The :code:`devices-cancel` icon.\n DEVICES_CHECK (IconName): The :code:`devices-check` icon.\n DEVICES_CODE (IconName): The :code:`devices-code` icon.\n DEVICES_COG (IconName): The :code:`devices-cog` icon.\n DEVICES_DOLLAR (IconName): The :code:`devices-dollar` icon.\n DEVICES_DOWN (IconName): The :code:`devices-down` icon.\n DEVICES_EXCLAMATION (IconName): The :code:`devices-exclamation` icon.\n DEVICES_HEART (IconName): The :code:`devices-heart` icon.\n DEVICES_MINUS (IconName): The :code:`devices-minus` icon.\n DEVICES_OFF (IconName): The :code:`devices-off` icon.\n DEVICES_PAUSE (IconName): The :code:`devices-pause` icon.\n DEVICES_PC (IconName): The :code:`devices-pc` icon.\n DEVICES_PC_OFF (IconName): The :code:`devices-pc-off` icon.\n DEVICES_PIN (IconName): The :code:`devices-pin` icon.\n DEVICES_PLUS (IconName): The :code:`devices-plus` icon.\n DEVICES_QUESTION (IconName): The :code:`devices-question` icon.\n DEVICES_SEARCH (IconName): The :code:`devices-search` icon.\n DEVICES_SHARE (IconName): The :code:`devices-share` icon.\n DEVICES_STAR (IconName): The :code:`devices-star` icon.\n DEVICES_UP (IconName): The :code:`devices-up` icon.\n DEVICES_X (IconName): The :code:`devices-x` icon.\n DIABOLO (IconName): The :code:`diabolo` icon.\n DIABOLO_OFF (IconName): The :code:`diabolo-off` icon.\n DIABOLO_PLUS (IconName): The :code:`diabolo-plus` icon.\n DIALPAD (IconName): The :code:`dialpad` icon.\n DIALPAD_FILLED (IconName): The :code:`dialpad-filled` icon.\n DIALPAD_OFF (IconName): The :code:`dialpad-off` icon.\n DIAMOND (IconName): The :code:`diamond` icon.\n DIAMOND_FILLED (IconName): The :code:`diamond-filled` icon.\n DIAMOND_OFF (IconName): The :code:`diamond-off` icon.\n DIAMONDS (IconName): The :code:`diamonds` icon.\n DIAMONDS_FILLED (IconName): The :code:`diamonds-filled` icon.\n DICE (IconName): The :code:`dice` icon.\n DICE_1 (IconName): The :code:`dice-1` icon.\n DICE_1_FILLED (IconName): The :code:`dice-1-filled` icon.\n DICE_2 (IconName): The :code:`dice-2` icon.\n DICE_2_FILLED (IconName): The :code:`dice-2-filled` icon.\n DICE_3 (IconName): The :code:`dice-3` icon.\n DICE_3_FILLED (IconName): The :code:`dice-3-filled` icon.\n DICE_4 (IconName): The :code:`dice-4` icon.\n DICE_4_FILLED (IconName): The :code:`dice-4-filled` icon.\n DICE_5 (IconName): The :code:`dice-5` icon.\n DICE_5_FILLED (IconName): The :code:`dice-5-filled` icon.\n DICE_6 (IconName): The :code:`dice-6` icon.\n DICE_6_FILLED (IconName): The :code:`dice-6-filled` icon.\n DICE_FILLED (IconName): The :code:`dice-filled` icon.\n DIMENSIONS (IconName): The :code:`dimensions` icon.\n DIRECTION (IconName): The :code:`direction` icon.\n DIRECTION_HORIZONTAL (IconName): The :code:`direction-horizontal` icon.\n DIRECTION_SIGN (IconName): The :code:`direction-sign` icon.\n DIRECTION_SIGN_FILLED (IconName): The :code:`direction-sign-filled` icon.\n DIRECTION_SIGN_OFF (IconName): The :code:`direction-sign-off` icon.\n DIRECTIONS (IconName): The :code:`directions` icon.\n DIRECTIONS_OFF (IconName): The :code:`directions-off` icon.\n DISABLED (IconName): The :code:`disabled` icon.\n DISABLED_2 (IconName): The :code:`disabled-2` icon.\n DISABLED_OFF (IconName): The :code:`disabled-off` icon.\n DISC (IconName): The :code:`disc` icon.\n DISC_GOLF (IconName): The :code:`disc-golf` icon.\n DISC_OFF (IconName): The :code:`disc-off` icon.\n DISCOUNT (IconName): The :code:`discount` icon.\n DISCOUNT_2 (IconName): The :code:`discount-2` icon.\n DISCOUNT_2_OFF (IconName): The :code:`discount-2-off` icon.\n DISCOUNT_CHECK (IconName): The :code:`discount-check` icon.\n DISCOUNT_CHECK_FILLED (IconName): The :code:`discount-check-filled` icon.\n DISCOUNT_OFF (IconName): The :code:`discount-off` icon.\n DIVIDE (IconName): The :code:`divide` icon.\n DNA (IconName): The :code:`dna` icon.\n DNA_2 (IconName): The :code:`dna-2` icon.\n DNA_2_OFF (IconName): The :code:`dna-2-off` icon.\n DNA_OFF (IconName): The :code:`dna-off` icon.\n DOG (IconName): The :code:`dog` icon.\n DOG_BOWL (IconName): The :code:`dog-bowl` icon.\n DOOR (IconName): The :code:`door` icon.\n DOOR_ENTER (IconName): The :code:`door-enter` icon.\n DOOR_EXIT (IconName): The :code:`door-exit` icon.\n DOOR_OFF (IconName): The :code:`door-off` icon.\n DOTS (IconName): The :code:`dots` icon.\n DOTS_CIRCLE_HORIZONTAL (IconName): The :code:`dots-circle-horizontal` icon.\n DOTS_DIAGONAL (IconName): The :code:`dots-diagonal` icon.\n DOTS_DIAGONAL_2 (IconName): The :code:`dots-diagonal-2` icon.\n DOTS_VERTICAL (IconName): The :code:`dots-vertical` icon.\n DOWNLOAD (IconName): The :code:`download` icon.\n DOWNLOAD_OFF (IconName): The :code:`download-off` icon.\n DRAG_DROP (IconName): The :code:`drag-drop` icon.\n DRAG_DROP_2 (IconName): The :code:`drag-drop-2` icon.\n DRONE (IconName): The :code:`drone` icon.\n DRONE_OFF (IconName): The :code:`drone-off` icon.\n DROP_CIRCLE (IconName): The :code:`drop-circle` icon.\n DROPLET (IconName): The :code:`droplet` icon.\n DROPLET_BOLT (IconName): The :code:`droplet-bolt` icon.\n DROPLET_CANCEL (IconName): The :code:`droplet-cancel` icon.\n DROPLET_CHECK (IconName): The :code:`droplet-check` icon.\n DROPLET_CODE (IconName): The :code:`droplet-code` icon.\n DROPLET_COG (IconName): The :code:`droplet-cog` icon.\n DROPLET_DOLLAR (IconName): The :code:`droplet-dollar` icon.\n DROPLET_DOWN (IconName): The :code:`droplet-down` icon.\n DROPLET_EXCLAMATION (IconName): The :code:`droplet-exclamation` icon.\n DROPLET_FILLED (IconName): The :code:`droplet-filled` icon.\n DROPLET_FILLED_2 (IconName): The :code:`droplet-filled-2` icon.\n DROPLET_HALF (IconName): The :code:`droplet-half` icon.\n DROPLET_HALF_2 (IconName): The :code:`droplet-half-2` icon.\n DROPLET_HALF_FILLED (IconName): The :code:`droplet-half-filled` icon.\n DROPLET_HEART (IconName): The :code:`droplet-heart` icon.\n DROPLET_MINUS (IconName): The :code:`droplet-minus` icon.\n DROPLET_OFF (IconName): The :code:`droplet-off` icon.\n DROPLET_PAUSE (IconName): The :code:`droplet-pause` icon.\n DROPLET_PIN (IconName): The :code:`droplet-pin` icon.\n DROPLET_PLUS (IconName): The :code:`droplet-plus` icon.\n DROPLET_QUESTION (IconName): The :code:`droplet-question` icon.\n DROPLET_SEARCH (IconName): The :code:`droplet-search` icon.\n DROPLET_SHARE (IconName): The :code:`droplet-share` icon.\n DROPLET_STAR (IconName): The :code:`droplet-star` icon.\n DROPLET_UP (IconName): The :code:`droplet-up` icon.\n DROPLET_X (IconName): The :code:`droplet-x` icon.\n DUAL_SCREEN (IconName): The :code:`dual-screen` icon.\n E_PASSPORT (IconName): The :code:`e-passport` icon.\n EAR (IconName): The :code:`ear` icon.\n EAR_OFF (IconName): The :code:`ear-off` icon.\n EASE_IN (IconName): The :code:`ease-in` icon.\n EASE_IN_CONTROL_POINT (IconName): The :code:`ease-in-control-point` icon.\n EASE_IN_OUT (IconName): The :code:`ease-in-out` icon.\n EASE_IN_OUT_CONTROL_POINTS (IconName): The :code:`ease-in-out-control-points` icon.\n EASE_OUT (IconName): The :code:`ease-out` icon.\n EASE_OUT_CONTROL_POINT (IconName): The :code:`ease-out-control-point` icon.\n EDIT (IconName): The :code:`edit` icon.\n EDIT_CIRCLE (IconName): The :code:`edit-circle` icon.\n EDIT_CIRCLE_OFF (IconName): The :code:`edit-circle-off` icon.\n EDIT_OFF (IconName): The :code:`edit-off` icon.\n EGG (IconName): The :code:`egg` icon.\n EGG_CRACKED (IconName): The :code:`egg-cracked` icon.\n EGG_FILLED (IconName): The :code:`egg-filled` icon.\n EGG_FRIED (IconName): The :code:`egg-fried` icon.\n EGG_OFF (IconName): The :code:`egg-off` icon.\n EGGS (IconName): The :code:`eggs` icon.\n ELEVATOR (IconName): The :code:`elevator` icon.\n ELEVATOR_OFF (IconName): The :code:`elevator-off` icon.\n EMERGENCY_BED (IconName): The :code:`emergency-bed` icon.\n EMPATHIZE (IconName): The :code:`empathize` icon.\n EMPATHIZE_OFF (IconName): The :code:`empathize-off` icon.\n EMPHASIS (IconName): The :code:`emphasis` icon.\n ENGINE (IconName): The :code:`engine` icon.\n ENGINE_OFF (IconName): The :code:`engine-off` icon.\n EQUAL (IconName): The :code:`equal` icon.\n EQUAL_DOUBLE (IconName): The :code:`equal-double` icon.\n EQUAL_NOT (IconName): The :code:`equal-not` icon.\n ERASER (IconName): The :code:`eraser` icon.\n ERASER_OFF (IconName): The :code:`eraser-off` icon.\n ERROR_404 (IconName): The :code:`error-404` icon.\n ERROR_404_OFF (IconName): The :code:`error-404-off` icon.\n EXCHANGE (IconName): The :code:`exchange` icon.\n EXCHANGE_OFF (IconName): The :code:`exchange-off` icon.\n EXCLAMATION_CIRCLE (IconName): The :code:`exclamation-circle` icon.\n EXCLAMATION_MARK (IconName): The :code:`exclamation-mark` icon.\n EXCLAMATION_MARK_OFF (IconName): The :code:`exclamation-mark-off` icon.\n EXPLICIT (IconName): The :code:`explicit` icon.\n EXPLICIT_OFF (IconName): The :code:`explicit-off` icon.\n EXPOSURE (IconName): The :code:`exposure` icon.\n EXPOSURE_0 (IconName): The :code:`exposure-0` icon.\n EXPOSURE_MINUS_1 (IconName): The :code:`exposure-minus-1` icon.\n EXPOSURE_MINUS_2 (IconName): The :code:`exposure-minus-2` icon.\n EXPOSURE_OFF (IconName): The :code:`exposure-off` icon.\n EXPOSURE_PLUS_1 (IconName): The :code:`exposure-plus-1` icon.\n EXPOSURE_PLUS_2 (IconName): The :code:`exposure-plus-2` icon.\n EXTERNAL_LINK (IconName): The :code:`external-link` icon.\n EXTERNAL_LINK_OFF (IconName): The :code:`external-link-off` icon.\n EYE (IconName): The :code:`eye` icon.\n EYE_CHECK (IconName): The :code:`eye-check` icon.\n EYE_CLOSED (IconName): The :code:`eye-closed` icon.\n EYE_COG (IconName): The :code:`eye-cog` icon.\n EYE_EDIT (IconName): The :code:`eye-edit` icon.\n EYE_EXCLAMATION (IconName): The :code:`eye-exclamation` icon.\n EYE_FILLED (IconName): The :code:`eye-filled` icon.\n EYE_HEART (IconName): The :code:`eye-heart` icon.\n EYE_OFF (IconName): The :code:`eye-off` icon.\n EYE_TABLE (IconName): The :code:`eye-table` icon.\n EYE_X (IconName): The :code:`eye-x` icon.\n EYEGLASS (IconName): The :code:`eyeglass` icon.\n EYEGLASS_2 (IconName): The :code:`eyeglass-2` icon.\n EYEGLASS_OFF (IconName): The :code:`eyeglass-off` icon.\n FACE_ID (IconName): The :code:`face-id` icon.\n FACE_ID_ERROR (IconName): The :code:`face-id-error` icon.\n FACE_MASK (IconName): The :code:`face-mask` icon.\n FACE_MASK_OFF (IconName): The :code:`face-mask-off` icon.\n FALL (IconName): The :code:`fall` icon.\n FEATHER (IconName): The :code:`feather` icon.\n FEATHER_OFF (IconName): The :code:`feather-off` icon.\n FENCE (IconName): The :code:`fence` icon.\n FENCE_OFF (IconName): The :code:`fence-off` icon.\n FIDGET_SPINNER (IconName): The :code:`fidget-spinner` icon.\n FILE (IconName): The :code:`file` icon.\n FILE_3D (IconName): The :code:`file-3d` icon.\n FILE_ALERT (IconName): The :code:`file-alert` icon.\n FILE_ANALYTICS (IconName): The :code:`file-analytics` icon.\n FILE_ARROW_LEFT (IconName): The :code:`file-arrow-left` icon.\n FILE_ARROW_RIGHT (IconName): The :code:`file-arrow-right` icon.\n FILE_BARCODE (IconName): The :code:`file-barcode` icon.\n FILE_BROKEN (IconName): The :code:`file-broken` icon.\n FILE_CERTIFICATE (IconName): The :code:`file-certificate` icon.\n FILE_CHART (IconName): The :code:`file-chart` icon.\n FILE_CHECK (IconName): The :code:`file-check` icon.\n FILE_CODE (IconName): The :code:`file-code` icon.\n FILE_CODE_2 (IconName): The :code:`file-code-2` icon.\n FILE_CV (IconName): The :code:`file-cv` icon.\n FILE_DATABASE (IconName): The :code:`file-database` icon.\n FILE_DELTA (IconName): The :code:`file-delta` icon.\n FILE_DESCRIPTION (IconName): The :code:`file-description` icon.\n FILE_DIFF (IconName): The :code:`file-diff` icon.\n FILE_DIGIT (IconName): The :code:`file-digit` icon.\n FILE_DISLIKE (IconName): The :code:`file-dislike` icon.\n FILE_DOLLAR (IconName): The :code:`file-dollar` icon.\n FILE_DOTS (IconName): The :code:`file-dots` icon.\n FILE_DOWNLOAD (IconName): The :code:`file-download` icon.\n FILE_EURO (IconName): The :code:`file-euro` icon.\n FILE_EXPORT (IconName): The :code:`file-export` icon.\n FILE_FILLED (IconName): The :code:`file-filled` icon.\n FILE_FUNCTION (IconName): The :code:`file-function` icon.\n FILE_HORIZONTAL (IconName): The :code:`file-horizontal` icon.\n FILE_IMPORT (IconName): The :code:`file-import` icon.\n FILE_INFINITY (IconName): The :code:`file-infinity` icon.\n FILE_INFO (IconName): The :code:`file-info` icon.\n FILE_INVOICE (IconName): The :code:`file-invoice` icon.\n FILE_LAMBDA (IconName): The :code:`file-lambda` icon.\n FILE_LIKE (IconName): The :code:`file-like` icon.\n FILE_MINUS (IconName): The :code:`file-minus` icon.\n FILE_MUSIC (IconName): The :code:`file-music` icon.\n FILE_OFF (IconName): The :code:`file-off` icon.\n FILE_ORIENTATION (IconName): The :code:`file-orientation` icon.\n FILE_PENCIL (IconName): The :code:`file-pencil` icon.\n FILE_PERCENT (IconName): The :code:`file-percent` icon.\n FILE_PHONE (IconName): The :code:`file-phone` icon.\n FILE_PLUS (IconName): The :code:`file-plus` icon.\n FILE_POWER (IconName): The :code:`file-power` icon.\n FILE_REPORT (IconName): The :code:`file-report` icon.\n FILE_RSS (IconName): The :code:`file-rss` icon.\n FILE_SCISSORS (IconName): The :code:`file-scissors` icon.\n FILE_SEARCH (IconName): The :code:`file-search` icon.\n FILE_SETTINGS (IconName): The :code:`file-settings` icon.\n FILE_SHREDDER (IconName): The :code:`file-shredder` icon.\n FILE_SIGNAL (IconName): The :code:`file-signal` icon.\n FILE_SPREADSHEET (IconName): The :code:`file-spreadsheet` icon.\n FILE_STACK (IconName): The :code:`file-stack` icon.\n FILE_STAR (IconName): The :code:`file-star` icon.\n FILE_SYMLINK (IconName): The :code:`file-symlink` icon.\n FILE_TEXT (IconName): The :code:`file-text` icon.\n FILE_TEXT_AI (IconName): The :code:`file-text-ai` icon.\n FILE_TIME (IconName): The :code:`file-time` icon.\n FILE_TYPOGRAPHY (IconName): The :code:`file-typography` icon.\n FILE_UNKNOWN (IconName): The :code:`file-unknown` icon.\n FILE_UPLOAD (IconName): The :code:`file-upload` icon.\n FILE_VECTOR (IconName): The :code:`file-vector` icon.\n FILE_X (IconName): The :code:`file-x` icon.\n FILE_X_FILLED (IconName): The :code:`file-x-filled` icon.\n FILE_ZIP (IconName): The :code:`file-zip` icon.\n FILES (IconName): The :code:`files` icon.\n FILES_OFF (IconName): The :code:`files-off` icon.\n FILTER (IconName): The :code:`filter` icon.\n FILTER_COG (IconName): The :code:`filter-cog` icon.\n FILTER_DOLLAR (IconName): The :code:`filter-dollar` icon.\n FILTER_EDIT (IconName): The :code:`filter-edit` icon.\n FILTER_MINUS (IconName): The :code:`filter-minus` icon.\n FILTER_OFF (IconName): The :code:`filter-off` icon.\n FILTER_PLUS (IconName): The :code:`filter-plus` icon.\n FILTER_STAR (IconName): The :code:`filter-star` icon.\n FILTER_X (IconName): The :code:`filter-x` icon.\n FILTERS (IconName): The :code:`filters` icon.\n FINGERPRINT (IconName): The :code:`fingerprint` icon.\n FINGERPRINT_OFF (IconName): The :code:`fingerprint-off` icon.\n FIRE_EXTINGUISHER (IconName): The :code:`fire-extinguisher` icon.\n FIRE_HYDRANT (IconName): The :code:`fire-hydrant` icon.\n FIRE_HYDRANT_OFF (IconName): The :code:`fire-hydrant-off` icon.\n FIRETRUCK (IconName): The :code:`firetruck` icon.\n FIRST_AID_KIT (IconName): The :code:`first-aid-kit` icon.\n FIRST_AID_KIT_OFF (IconName): The :code:`first-aid-kit-off` icon.\n FISH (IconName): The :code:`fish` icon.\n FISH_BONE (IconName): The :code:`fish-bone` icon.\n FISH_CHRISTIANITY (IconName): The :code:`fish-christianity` icon.\n FISH_HOOK (IconName): The :code:`fish-hook` icon.\n FISH_HOOK_OFF (IconName): The :code:`fish-hook-off` icon.\n FISH_OFF (IconName): The :code:`fish-off` icon.\n FLAG (IconName): The :code:`flag` icon.\n FLAG_2 (IconName): The :code:`flag-2` icon.\n FLAG_2_FILLED (IconName): The :code:`flag-2-filled` icon.\n FLAG_2_OFF (IconName): The :code:`flag-2-off` icon.\n FLAG_3 (IconName): The :code:`flag-3` icon.\n FLAG_3_FILLED (IconName): The :code:`flag-3-filled` icon.\n FLAG_FILLED (IconName): The :code:`flag-filled` icon.\n FLAG_OFF (IconName): The :code:`flag-off` icon.\n FLAME (IconName): The :code:`flame` icon.\n FLAME_OFF (IconName): The :code:`flame-off` icon.\n FLARE (IconName): The :code:`flare` icon.\n FLASK (IconName): The :code:`flask` icon.\n FLASK_2 (IconName): The :code:`flask-2` icon.\n FLASK_2_OFF (IconName): The :code:`flask-2-off` icon.\n FLASK_OFF (IconName): The :code:`flask-off` icon.\n FLIP_FLOPS (IconName): The :code:`flip-flops` icon.\n FLIP_HORIZONTAL (IconName): The :code:`flip-horizontal` icon.\n FLIP_VERTICAL (IconName): The :code:`flip-vertical` icon.\n FLOAT_CENTER (IconName): The :code:`float-center` icon.\n FLOAT_LEFT (IconName): The :code:`float-left` icon.\n FLOAT_NONE (IconName): The :code:`float-none` icon.\n FLOAT_RIGHT (IconName): The :code:`float-right` icon.\n FLOWER (IconName): The :code:`flower` icon.\n FLOWER_OFF (IconName): The :code:`flower-off` icon.\n FOCUS (IconName): The :code:`focus` icon.\n FOCUS_2 (IconName): The :code:`focus-2` icon.\n FOCUS_AUTO (IconName): The :code:`focus-auto` icon.\n FOCUS_CENTERED (IconName): The :code:`focus-centered` icon.\n FOLD (IconName): The :code:`fold` icon.\n FOLD_DOWN (IconName): The :code:`fold-down` icon.\n FOLD_UP (IconName): The :code:`fold-up` icon.\n FOLDER (IconName): The :code:`folder` icon.\n FOLDER_BOLT (IconName): The :code:`folder-bolt` icon.\n FOLDER_CANCEL (IconName): The :code:`folder-cancel` icon.\n FOLDER_CHECK (IconName): The :code:`folder-check` icon.\n FOLDER_CODE (IconName): The :code:`folder-code` icon.\n FOLDER_COG (IconName): The :code:`folder-cog` icon.\n FOLDER_DOLLAR (IconName): The :code:`folder-dollar` icon.\n FOLDER_DOWN (IconName): The :code:`folder-down` icon.\n FOLDER_EXCLAMATION (IconName): The :code:`folder-exclamation` icon.\n FOLDER_FILLED (IconName): The :code:`folder-filled` icon.\n FOLDER_HEART (IconName): The :code:`folder-heart` icon.\n FOLDER_MINUS (IconName): The :code:`folder-minus` icon.\n FOLDER_OFF (IconName): The :code:`folder-off` icon.\n FOLDER_OPEN (IconName): The :code:`folder-open` icon.\n FOLDER_PAUSE (IconName): The :code:`folder-pause` icon.\n FOLDER_PIN (IconName): The :code:`folder-pin` icon.\n FOLDER_PLUS (IconName): The :code:`folder-plus` icon.\n FOLDER_QUESTION (IconName): The :code:`folder-question` icon.\n FOLDER_SEARCH (IconName): The :code:`folder-search` icon.\n FOLDER_SHARE (IconName): The :code:`folder-share` icon.\n FOLDER_STAR (IconName): The :code:`folder-star` icon.\n FOLDER_SYMLINK (IconName): The :code:`folder-symlink` icon.\n FOLDER_UP (IconName): The :code:`folder-up` icon.\n FOLDER_X (IconName): The :code:`folder-x` icon.\n FOLDERS (IconName): The :code:`folders` icon.\n FOLDERS_OFF (IconName): The :code:`folders-off` icon.\n FORBID (IconName): The :code:`forbid` icon.\n FORBID_2 (IconName): The :code:`forbid-2` icon.\n FORKLIFT (IconName): The :code:`forklift` icon.\n FORMS (IconName): The :code:`forms` icon.\n FOUNTAIN (IconName): The :code:`fountain` icon.\n FOUNTAIN_OFF (IconName): The :code:`fountain-off` icon.\n FRAME (IconName): The :code:`frame` icon.\n FRAME_OFF (IconName): The :code:`frame-off` icon.\n FREE_RIGHTS (IconName): The :code:`free-rights` icon.\n FREEZE_COLUMN (IconName): The :code:`freeze-column` icon.\n FREEZE_ROW (IconName): The :code:`freeze-row` icon.\n FREEZE_ROW_COLUMN (IconName): The :code:`freeze-row-column` icon.\n FRIDGE (IconName): The :code:`fridge` icon.\n FRIDGE_OFF (IconName): The :code:`fridge-off` icon.\n FRIENDS (IconName): The :code:`friends` icon.\n FRIENDS_OFF (IconName): The :code:`friends-off` icon.\n FRUSTUM (IconName): The :code:`frustum` icon.\n FRUSTUM_OFF (IconName): The :code:`frustum-off` icon.\n FRUSTUM_PLUS (IconName): The :code:`frustum-plus` icon.\n FUNCTION (IconName): The :code:`function` icon.\n FUNCTION_OFF (IconName): The :code:`function-off` icon.\n GARDEN_CART (IconName): The :code:`garden-cart` icon.\n GARDEN_CART_OFF (IconName): The :code:`garden-cart-off` icon.\n GAS_STATION (IconName): The :code:`gas-station` icon.\n GAS_STATION_OFF (IconName): The :code:`gas-station-off` icon.\n GAUGE (IconName): The :code:`gauge` icon.\n GAUGE_OFF (IconName): The :code:`gauge-off` icon.\n GAVEL (IconName): The :code:`gavel` icon.\n GENDER_AGENDER (IconName): The :code:`gender-agender` icon.\n GENDER_ANDROGYNE (IconName): The :code:`gender-androgyne` icon.\n GENDER_BIGENDER (IconName): The :code:`gender-bigender` icon.\n GENDER_DEMIBOY (IconName): The :code:`gender-demiboy` icon.\n GENDER_DEMIGIRL (IconName): The :code:`gender-demigirl` icon.\n GENDER_EPICENE (IconName): The :code:`gender-epicene` icon.\n GENDER_FEMALE (IconName): The :code:`gender-female` icon.\n GENDER_FEMME (IconName): The :code:`gender-femme` icon.\n GENDER_GENDERFLUID (IconName): The :code:`gender-genderfluid` icon.\n GENDER_GENDERLESS (IconName): The :code:`gender-genderless` icon.\n GENDER_GENDERQUEER (IconName): The :code:`gender-genderqueer` icon.\n GENDER_HERMAPHRODITE (IconName): The :code:`gender-hermaphrodite` icon.\n GENDER_INTERGENDER (IconName): The :code:`gender-intergender` icon.\n GENDER_MALE (IconName): The :code:`gender-male` icon.\n GENDER_NEUTROIS (IconName): The :code:`gender-neutrois` icon.\n GENDER_THIRD (IconName): The :code:`gender-third` icon.\n GENDER_TRANSGENDER (IconName): The :code:`gender-transgender` icon.\n GENDER_TRASVESTI (IconName): The :code:`gender-trasvesti` icon.\n GEOMETRY (IconName): The :code:`geometry` icon.\n GHOST (IconName): The :code:`ghost` icon.\n GHOST_2 (IconName): The :code:`ghost-2` icon.\n GHOST_2_FILLED (IconName): The :code:`ghost-2-filled` icon.\n GHOST_FILLED (IconName): The :code:`ghost-filled` icon.\n GHOST_OFF (IconName): The :code:`ghost-off` icon.\n GIF (IconName): The :code:`gif` icon.\n GIFT (IconName): The :code:`gift` icon.\n GIFT_CARD (IconName): The :code:`gift-card` icon.\n GIFT_OFF (IconName): The :code:`gift-off` icon.\n GIT_BRANCH (IconName): The :code:`git-branch` icon.\n GIT_BRANCH_DELETED (IconName): The :code:`git-branch-deleted` icon.\n GIT_CHERRY_PICK (IconName): The :code:`git-cherry-pick` icon.\n GIT_COMMIT (IconName): The :code:`git-commit` icon.\n GIT_COMPARE (IconName): The :code:`git-compare` icon.\n GIT_FORK (IconName): The :code:`git-fork` icon.\n GIT_MERGE (IconName): The :code:`git-merge` icon.\n GIT_PULL_REQUEST (IconName): The :code:`git-pull-request` icon.\n GIT_PULL_REQUEST_CLOSED (IconName): The :code:`git-pull-request-closed` icon.\n GIT_PULL_REQUEST_DRAFT (IconName): The :code:`git-pull-request-draft` icon.\n GIZMO (IconName): The :code:`gizmo` icon.\n GLASS (IconName): The :code:`glass` icon.\n GLASS_FULL (IconName): The :code:`glass-full` icon.\n GLASS_OFF (IconName): The :code:`glass-off` icon.\n GLOBE (IconName): The :code:`globe` icon.\n GLOBE_OFF (IconName): The :code:`globe-off` icon.\n GO_GAME (IconName): The :code:`go-game` icon.\n GOLF (IconName): The :code:`golf` icon.\n GOLF_OFF (IconName): The :code:`golf-off` icon.\n GPS (IconName): The :code:`gps` icon.\n GRADIENTER (IconName): The :code:`gradienter` icon.\n GRAIN (IconName): The :code:`grain` icon.\n GRAPH (IconName): The :code:`graph` icon.\n GRAPH_OFF (IconName): The :code:`graph-off` icon.\n GRAVE (IconName): The :code:`grave` icon.\n GRAVE_2 (IconName): The :code:`grave-2` icon.\n GRID_DOTS (IconName): The :code:`grid-dots` icon.\n GRID_PATTERN (IconName): The :code:`grid-pattern` icon.\n GRILL (IconName): The :code:`grill` icon.\n GRILL_FORK (IconName): The :code:`grill-fork` icon.\n GRILL_OFF (IconName): The :code:`grill-off` icon.\n GRILL_SPATULA (IconName): The :code:`grill-spatula` icon.\n GRIP_HORIZONTAL (IconName): The :code:`grip-horizontal` icon.\n GRIP_VERTICAL (IconName): The :code:`grip-vertical` icon.\n GROWTH (IconName): The :code:`growth` icon.\n GUITAR_PICK (IconName): The :code:`guitar-pick` icon.\n GUITAR_PICK_FILLED (IconName): The :code:`guitar-pick-filled` icon.\n H_1 (IconName): The :code:`h-1` icon.\n H_2 (IconName): The :code:`h-2` icon.\n H_3 (IconName): The :code:`h-3` icon.\n H_4 (IconName): The :code:`h-4` icon.\n H_5 (IconName): The :code:`h-5` icon.\n H_6 (IconName): The :code:`h-6` icon.\n HAMMER (IconName): The :code:`hammer` icon.\n HAMMER_OFF (IconName): The :code:`hammer-off` icon.\n HAND_CLICK (IconName): The :code:`hand-click` icon.\n HAND_FINGER (IconName): The :code:`hand-finger` icon.\n HAND_FINGER_OFF (IconName): The :code:`hand-finger-off` icon.\n HAND_GRAB (IconName): The :code:`hand-grab` icon.\n HAND_LITTLE_FINGER (IconName): The :code:`hand-little-finger` icon.\n HAND_MIDDLE_FINGER (IconName): The :code:`hand-middle-finger` icon.\n HAND_MOVE (IconName): The :code:`hand-move` icon.\n HAND_OFF (IconName): The :code:`hand-off` icon.\n HAND_RING_FINGER (IconName): The :code:`hand-ring-finger` icon.\n HAND_ROCK (IconName): The :code:`hand-rock` icon.\n HAND_SANITIZER (IconName): The :code:`hand-sanitizer` icon.\n HAND_STOP (IconName): The :code:`hand-stop` icon.\n HAND_THREE_FINGERS (IconName): The :code:`hand-three-fingers` icon.\n HAND_TWO_FINGERS (IconName): The :code:`hand-two-fingers` icon.\n HANGER (IconName): The :code:`hanger` icon.\n HANGER_2 (IconName): The :code:`hanger-2` icon.\n HANGER_OFF (IconName): The :code:`hanger-off` icon.\n HASH (IconName): The :code:`hash` icon.\n HAZE (IconName): The :code:`haze` icon.\n HAZE_MOON (IconName): The :code:`haze-moon` icon.\n HDR (IconName): The :code:`hdr` icon.\n HEADING (IconName): The :code:`heading` icon.\n HEADING_OFF (IconName): The :code:`heading-off` icon.\n HEADPHONES (IconName): The :code:`headphones` icon.\n HEADPHONES_FILLED (IconName): The :code:`headphones-filled` icon.\n HEADPHONES_OFF (IconName): The :code:`headphones-off` icon.\n HEADSET (IconName): The :code:`headset` icon.\n HEADSET_OFF (IconName): The :code:`headset-off` icon.\n HEALTH_RECOGNITION (IconName): The :code:`health-recognition` icon.\n HEART (IconName): The :code:`heart` icon.\n HEART_BROKEN (IconName): The :code:`heart-broken` icon.\n HEART_FILLED (IconName): The :code:`heart-filled` icon.\n HEART_HANDSHAKE (IconName): The :code:`heart-handshake` icon.\n HEART_MINUS (IconName): The :code:`heart-minus` icon.\n HEART_OFF (IconName): The :code:`heart-off` icon.\n HEART_PLUS (IconName): The :code:`heart-plus` icon.\n HEART_RATE_MONITOR (IconName): The :code:`heart-rate-monitor` icon.\n HEARTBEAT (IconName): The :code:`heartbeat` icon.\n HEARTS (IconName): The :code:`hearts` icon.\n HEARTS_OFF (IconName): The :code:`hearts-off` icon.\n HELICOPTER (IconName): The :code:`helicopter` icon.\n HELICOPTER_LANDING (IconName): The :code:`helicopter-landing` icon.\n HELMET (IconName): The :code:`helmet` icon.\n HELMET_OFF (IconName): The :code:`helmet-off` icon.\n HELP (IconName): The :code:`help` icon.\n HELP_CIRCLE (IconName): The :code:`help-circle` icon.\n HELP_CIRCLE_FILLED (IconName): The :code:`help-circle-filled` icon.\n HELP_HEXAGON (IconName): The :code:`help-hexagon` icon.\n HELP_HEXAGON_FILLED (IconName): The :code:`help-hexagon-filled` icon.\n HELP_OCTAGON (IconName): The :code:`help-octagon` icon.\n HELP_OCTAGON_FILLED (IconName): The :code:`help-octagon-filled` icon.\n HELP_OFF (IconName): The :code:`help-off` icon.\n HELP_SMALL (IconName): The :code:`help-small` icon.\n HELP_SQUARE (IconName): The :code:`help-square` icon.\n HELP_SQUARE_FILLED (IconName): The :code:`help-square-filled` icon.\n HELP_SQUARE_ROUNDED (IconName): The :code:`help-square-rounded` icon.\n HELP_SQUARE_ROUNDED_FILLED (IconName): The :code:`help-square-rounded-filled` icon.\n HELP_TRIANGLE (IconName): The :code:`help-triangle` icon.\n HELP_TRIANGLE_FILLED (IconName): The :code:`help-triangle-filled` icon.\n HEMISPHERE (IconName): The :code:`hemisphere` icon.\n HEMISPHERE_OFF (IconName): The :code:`hemisphere-off` icon.\n HEMISPHERE_PLUS (IconName): The :code:`hemisphere-plus` icon.\n HEXAGON (IconName): The :code:`hexagon` icon.\n HEXAGON_0_FILLED (IconName): The :code:`hexagon-0-filled` icon.\n HEXAGON_1_FILLED (IconName): The :code:`hexagon-1-filled` icon.\n HEXAGON_2_FILLED (IconName): The :code:`hexagon-2-filled` icon.\n HEXAGON_3_FILLED (IconName): The :code:`hexagon-3-filled` icon.\n HEXAGON_3D (IconName): The :code:`hexagon-3d` icon.\n HEXAGON_4_FILLED (IconName): The :code:`hexagon-4-filled` icon.\n HEXAGON_5_FILLED (IconName): The :code:`hexagon-5-filled` icon.\n HEXAGON_6_FILLED (IconName): The :code:`hexagon-6-filled` icon.\n HEXAGON_7_FILLED (IconName): The :code:`hexagon-7-filled` icon.\n HEXAGON_8_FILLED (IconName): The :code:`hexagon-8-filled` icon.\n HEXAGON_9_FILLED (IconName): The :code:`hexagon-9-filled` icon.\n HEXAGON_FILLED (IconName): The :code:`hexagon-filled` icon.\n HEXAGON_LETTER_A (IconName): The :code:`hexagon-letter-a` icon.\n HEXAGON_LETTER_B (IconName): The :code:`hexagon-letter-b` icon.\n HEXAGON_LETTER_C (IconName): The :code:`hexagon-letter-c` icon.\n HEXAGON_LETTER_D (IconName): The :code:`hexagon-letter-d` icon.\n HEXAGON_LETTER_E (IconName): The :code:`hexagon-letter-e` icon.\n HEXAGON_LETTER_F (IconName): The :code:`hexagon-letter-f` icon.\n HEXAGON_LETTER_G (IconName): The :code:`hexagon-letter-g` icon.\n HEXAGON_LETTER_H (IconName): The :code:`hexagon-letter-h` icon.\n HEXAGON_LETTER_I (IconName): The :code:`hexagon-letter-i` icon.\n HEXAGON_LETTER_J (IconName): The :code:`hexagon-letter-j` icon.\n HEXAGON_LETTER_K (IconName): The :code:`hexagon-letter-k` icon.\n HEXAGON_LETTER_L (IconName): The :code:`hexagon-letter-l` icon.\n HEXAGON_LETTER_M (IconName): The :code:`hexagon-letter-m` icon.\n HEXAGON_LETTER_N (IconName): The :code:`hexagon-letter-n` icon.\n HEXAGON_LETTER_O (IconName): The :code:`hexagon-letter-o` icon.\n HEXAGON_LETTER_P (IconName): The :code:`hexagon-letter-p` icon.\n HEXAGON_LETTER_Q (IconName): The :code:`hexagon-letter-q` icon.\n HEXAGON_LETTER_R (IconName): The :code:`hexagon-letter-r` icon.\n HEXAGON_LETTER_S (IconName): The :code:`hexagon-letter-s` icon.\n HEXAGON_LETTER_T (IconName): The :code:`hexagon-letter-t` icon.\n HEXAGON_LETTER_U (IconName): The :code:`hexagon-letter-u` icon.\n HEXAGON_LETTER_V (IconName): The :code:`hexagon-letter-v` icon.\n HEXAGON_LETTER_W (IconName): The :code:`hexagon-letter-w` icon.\n HEXAGON_LETTER_X (IconName): The :code:`hexagon-letter-x` icon.\n HEXAGON_LETTER_Y (IconName): The :code:`hexagon-letter-y` icon.\n HEXAGON_LETTER_Z (IconName): The :code:`hexagon-letter-z` icon.\n HEXAGON_NUMBER_0 (IconName): The :code:`hexagon-number-0` icon.\n HEXAGON_NUMBER_1 (IconName): The :code:`hexagon-number-1` icon.\n HEXAGON_NUMBER_2 (IconName): The :code:`hexagon-number-2` icon.\n HEXAGON_NUMBER_3 (IconName): The :code:`hexagon-number-3` icon.\n HEXAGON_NUMBER_4 (IconName): The :code:`hexagon-number-4` icon.\n HEXAGON_NUMBER_5 (IconName): The :code:`hexagon-number-5` icon.\n HEXAGON_NUMBER_6 (IconName): The :code:`hexagon-number-6` icon.\n HEXAGON_NUMBER_7 (IconName): The :code:`hexagon-number-7` icon.\n HEXAGON_NUMBER_8 (IconName): The :code:`hexagon-number-8` icon.\n HEXAGON_NUMBER_9 (IconName): The :code:`hexagon-number-9` icon.\n HEXAGON_OFF (IconName): The :code:`hexagon-off` icon.\n HEXAGONAL_PRISM (IconName): The :code:`hexagonal-prism` icon.\n HEXAGONAL_PRISM_OFF (IconName): The :code:`hexagonal-prism-off` icon.\n HEXAGONAL_PRISM_PLUS (IconName): The :code:`hexagonal-prism-plus` icon.\n HEXAGONAL_PYRAMID (IconName): The :code:`hexagonal-pyramid` icon.\n HEXAGONAL_PYRAMID_OFF (IconName): The :code:`hexagonal-pyramid-off` icon.\n HEXAGONAL_PYRAMID_PLUS (IconName): The :code:`hexagonal-pyramid-plus` icon.\n HEXAGONS (IconName): The :code:`hexagons` icon.\n HEXAGONS_OFF (IconName): The :code:`hexagons-off` icon.\n HIERARCHY (IconName): The :code:`hierarchy` icon.\n HIERARCHY_2 (IconName): The :code:`hierarchy-2` icon.\n HIERARCHY_3 (IconName): The :code:`hierarchy-3` icon.\n HIERARCHY_OFF (IconName): The :code:`hierarchy-off` icon.\n HIGHLIGHT (IconName): The :code:`highlight` icon.\n HIGHLIGHT_OFF (IconName): The :code:`highlight-off` icon.\n HISTORY (IconName): The :code:`history` icon.\n HISTORY_OFF (IconName): The :code:`history-off` icon.\n HISTORY_TOGGLE (IconName): The :code:`history-toggle` icon.\n HOME (IconName): The :code:`home` icon.\n HOME_2 (IconName): The :code:`home-2` icon.\n HOME_BOLT (IconName): The :code:`home-bolt` icon.\n HOME_CANCEL (IconName): The :code:`home-cancel` icon.\n HOME_CHECK (IconName): The :code:`home-check` icon.\n HOME_COG (IconName): The :code:`home-cog` icon.\n HOME_DOLLAR (IconName): The :code:`home-dollar` icon.\n HOME_DOT (IconName): The :code:`home-dot` icon.\n HOME_DOWN (IconName): The :code:`home-down` icon.\n HOME_ECO (IconName): The :code:`home-eco` icon.\n HOME_EDIT (IconName): The :code:`home-edit` icon.\n HOME_EXCLAMATION (IconName): The :code:`home-exclamation` icon.\n HOME_HAND (IconName): The :code:`home-hand` icon.\n HOME_HEART (IconName): The :code:`home-heart` icon.\n HOME_INFINITY (IconName): The :code:`home-infinity` icon.\n HOME_LINK (IconName): The :code:`home-link` icon.\n HOME_MINUS (IconName): The :code:`home-minus` icon.\n HOME_MOVE (IconName): The :code:`home-move` icon.\n HOME_OFF (IconName): The :code:`home-off` icon.\n HOME_PLUS (IconName): The :code:`home-plus` icon.\n HOME_QUESTION (IconName): The :code:`home-question` icon.\n HOME_RIBBON (IconName): The :code:`home-ribbon` icon.\n HOME_SEARCH (IconName): The :code:`home-search` icon.\n HOME_SHARE (IconName): The :code:`home-share` icon.\n HOME_SHIELD (IconName): The :code:`home-shield` icon.\n HOME_SIGNAL (IconName): The :code:`home-signal` icon.\n HOME_STAR (IconName): The :code:`home-star` icon.\n HOME_STATS (IconName): The :code:`home-stats` icon.\n HOME_UP (IconName): The :code:`home-up` icon.\n HOME_X (IconName): The :code:`home-x` icon.\n HORSE_TOY (IconName): The :code:`horse-toy` icon.\n HOTEL_SERVICE (IconName): The :code:`hotel-service` icon.\n HOURGLASS (IconName): The :code:`hourglass` icon.\n HOURGLASS_EMPTY (IconName): The :code:`hourglass-empty` icon.\n HOURGLASS_FILLED (IconName): The :code:`hourglass-filled` icon.\n HOURGLASS_HIGH (IconName): The :code:`hourglass-high` icon.\n HOURGLASS_LOW (IconName): The :code:`hourglass-low` icon.\n HOURGLASS_OFF (IconName): The :code:`hourglass-off` icon.\n HTML (IconName): The :code:`html` icon.\n HTTP_CONNECT (IconName): The :code:`http-connect` icon.\n HTTP_DELETE (IconName): The :code:`http-delete` icon.\n HTTP_GET (IconName): The :code:`http-get` icon.\n HTTP_HEAD (IconName): The :code:`http-head` icon.\n HTTP_OPTIONS (IconName): The :code:`http-options` icon.\n HTTP_PATCH (IconName): The :code:`http-patch` icon.\n HTTP_POST (IconName): The :code:`http-post` icon.\n HTTP_PUT (IconName): The :code:`http-put` icon.\n HTTP_QUE (IconName): The :code:`http-que` icon.\n HTTP_TRACE (IconName): The :code:`http-trace` icon.\n ICE_CREAM (IconName): The :code:`ice-cream` icon.\n ICE_CREAM_2 (IconName): The :code:`ice-cream-2` icon.\n ICE_CREAM_OFF (IconName): The :code:`ice-cream-off` icon.\n ICE_SKATING (IconName): The :code:`ice-skating` icon.\n ICONS (IconName): The :code:`icons` icon.\n ICONS_OFF (IconName): The :code:`icons-off` icon.\n ID (IconName): The :code:`id` icon.\n ID_BADGE (IconName): The :code:`id-badge` icon.\n ID_BADGE_2 (IconName): The :code:`id-badge-2` icon.\n ID_BADGE_OFF (IconName): The :code:`id-badge-off` icon.\n ID_OFF (IconName): The :code:`id-off` icon.\n INBOX (IconName): The :code:`inbox` icon.\n INBOX_OFF (IconName): The :code:`inbox-off` icon.\n INDENT_DECREASE (IconName): The :code:`indent-decrease` icon.\n INDENT_INCREASE (IconName): The :code:`indent-increase` icon.\n INFINITY (IconName): The :code:`infinity` icon.\n INFINITY_OFF (IconName): The :code:`infinity-off` icon.\n INFO_CIRCLE (IconName): The :code:`info-circle` icon.\n INFO_CIRCLE_FILLED (IconName): The :code:`info-circle-filled` icon.\n INFO_HEXAGON (IconName): The :code:`info-hexagon` icon.\n INFO_HEXAGON_FILLED (IconName): The :code:`info-hexagon-filled` icon.\n INFO_OCTAGON (IconName): The :code:`info-octagon` icon.\n INFO_OCTAGON_FILLED (IconName): The :code:`info-octagon-filled` icon.\n INFO_SMALL (IconName): The :code:`info-small` icon.\n INFO_SQUARE (IconName): The :code:`info-square` icon.\n INFO_SQUARE_FILLED (IconName): The :code:`info-square-filled` icon.\n INFO_SQUARE_ROUNDED (IconName): The :code:`info-square-rounded` icon.\n INFO_SQUARE_ROUNDED_FILLED (IconName): The :code:`info-square-rounded-filled` icon.\n INFO_TRIANGLE (IconName): The :code:`info-triangle` icon.\n INFO_TRIANGLE_FILLED (IconName): The :code:`info-triangle-filled` icon.\n INNER_SHADOW_BOTTOM (IconName): The :code:`inner-shadow-bottom` icon.\n INNER_SHADOW_BOTTOM_FILLED (IconName): The :code:`inner-shadow-bottom-filled` icon.\n INNER_SHADOW_BOTTOM_LEFT (IconName): The :code:`inner-shadow-bottom-left` icon.\n INNER_SHADOW_BOTTOM_LEFT_FILLED (IconName): The :code:`inner-shadow-bottom-left-filled` icon.\n INNER_SHADOW_BOTTOM_RIGHT (IconName): The :code:`inner-shadow-bottom-right` icon.\n INNER_SHADOW_BOTTOM_RIGHT_FILLED (IconName): The :code:`inner-shadow-bottom-right-filled` icon.\n INNER_SHADOW_LEFT (IconName): The :code:`inner-shadow-left` icon.\n INNER_SHADOW_LEFT_FILLED (IconName): The :code:`inner-shadow-left-filled` icon.\n INNER_SHADOW_RIGHT (IconName): The :code:`inner-shadow-right` icon.\n INNER_SHADOW_RIGHT_FILLED (IconName): The :code:`inner-shadow-right-filled` icon.\n INNER_SHADOW_TOP (IconName): The :code:`inner-shadow-top` icon.\n INNER_SHADOW_TOP_FILLED (IconName): The :code:`inner-shadow-top-filled` icon.\n INNER_SHADOW_TOP_LEFT (IconName): The :code:`inner-shadow-top-left` icon.\n INNER_SHADOW_TOP_LEFT_FILLED (IconName): The :code:`inner-shadow-top-left-filled` icon.\n INNER_SHADOW_TOP_RIGHT (IconName): The :code:`inner-shadow-top-right` icon.\n INNER_SHADOW_TOP_RIGHT_FILLED (IconName): The :code:`inner-shadow-top-right-filled` icon.\n INPUT_SEARCH (IconName): The :code:`input-search` icon.\n IRONING (IconName): The :code:`ironing` icon.\n IRONING_1 (IconName): The :code:`ironing-1` icon.\n IRONING_2 (IconName): The :code:`ironing-2` icon.\n IRONING_3 (IconName): The :code:`ironing-3` icon.\n IRONING_OFF (IconName): The :code:`ironing-off` icon.\n IRONING_STEAM (IconName): The :code:`ironing-steam` icon.\n IRONING_STEAM_OFF (IconName): The :code:`ironing-steam-off` icon.\n IRREGULAR_POLYHEDRON (IconName): The :code:`irregular-polyhedron` icon.\n IRREGULAR_POLYHEDRON_OFF (IconName): The :code:`irregular-polyhedron-off` icon.\n IRREGULAR_POLYHEDRON_PLUS (IconName): The :code:`irregular-polyhedron-plus` icon.\n ITALIC (IconName): The :code:`italic` icon.\n JACKET (IconName): The :code:`jacket` icon.\n JETPACK (IconName): The :code:`jetpack` icon.\n JEWISH_STAR (IconName): The :code:`jewish-star` icon.\n JEWISH_STAR_FILLED (IconName): The :code:`jewish-star-filled` icon.\n JPG (IconName): The :code:`jpg` icon.\n JSON (IconName): The :code:`json` icon.\n JUMP_ROPE (IconName): The :code:`jump-rope` icon.\n KARATE (IconName): The :code:`karate` icon.\n KAYAK (IconName): The :code:`kayak` icon.\n KERING (IconName): The :code:`kering` icon.\n KEY (IconName): The :code:`key` icon.\n KEY_OFF (IconName): The :code:`key-off` icon.\n KEYBOARD (IconName): The :code:`keyboard` icon.\n KEYBOARD_HIDE (IconName): The :code:`keyboard-hide` icon.\n KEYBOARD_OFF (IconName): The :code:`keyboard-off` icon.\n KEYBOARD_SHOW (IconName): The :code:`keyboard-show` icon.\n KEYFRAME (IconName): The :code:`keyframe` icon.\n KEYFRAME_ALIGN_CENTER (IconName): The :code:`keyframe-align-center` icon.\n KEYFRAME_ALIGN_HORIZONTAL (IconName): The :code:`keyframe-align-horizontal` icon.\n KEYFRAME_ALIGN_VERTICAL (IconName): The :code:`keyframe-align-vertical` icon.\n KEYFRAMES (IconName): The :code:`keyframes` icon.\n LADDER (IconName): The :code:`ladder` icon.\n LADDER_OFF (IconName): The :code:`ladder-off` icon.\n LAMBDA (IconName): The :code:`lambda` icon.\n LAMP (IconName): The :code:`lamp` icon.\n LAMP_2 (IconName): The :code:`lamp-2` icon.\n LAMP_OFF (IconName): The :code:`lamp-off` icon.\n LANE (IconName): The :code:`lane` icon.\n LANGUAGE (IconName): The :code:`language` icon.\n LANGUAGE_HIRAGANA (IconName): The :code:`language-hiragana` icon.\n LANGUAGE_KATAKANA (IconName): The :code:`language-katakana` icon.\n LANGUAGE_OFF (IconName): The :code:`language-off` icon.\n LASSO (IconName): The :code:`lasso` icon.\n LASSO_OFF (IconName): The :code:`lasso-off` icon.\n LASSO_POLYGON (IconName): The :code:`lasso-polygon` icon.\n LAYERS_DIFFERENCE (IconName): The :code:`layers-difference` icon.\n LAYERS_INTERSECT (IconName): The :code:`layers-intersect` icon.\n LAYERS_INTERSECT_2 (IconName): The :code:`layers-intersect-2` icon.\n LAYERS_LINKED (IconName): The :code:`layers-linked` icon.\n LAYERS_OFF (IconName): The :code:`layers-off` icon.\n LAYERS_SUBTRACT (IconName): The :code:`layers-subtract` icon.\n LAYERS_UNION (IconName): The :code:`layers-union` icon.\n LAYOUT (IconName): The :code:`layout` icon.\n LAYOUT_2 (IconName): The :code:`layout-2` icon.\n LAYOUT_ALIGN_BOTTOM (IconName): The :code:`layout-align-bottom` icon.\n LAYOUT_ALIGN_CENTER (IconName): The :code:`layout-align-center` icon.\n LAYOUT_ALIGN_LEFT (IconName): The :code:`layout-align-left` icon.\n LAYOUT_ALIGN_MIDDLE (IconName): The :code:`layout-align-middle` icon.\n LAYOUT_ALIGN_RIGHT (IconName): The :code:`layout-align-right` icon.\n LAYOUT_ALIGN_TOP (IconName): The :code:`layout-align-top` icon.\n LAYOUT_BOARD (IconName): The :code:`layout-board` icon.\n LAYOUT_BOARD_SPLIT (IconName): The :code:`layout-board-split` icon.\n LAYOUT_BOTTOMBAR (IconName): The :code:`layout-bottombar` icon.\n LAYOUT_BOTTOMBAR_COLLAPSE (IconName): The :code:`layout-bottombar-collapse` icon.\n LAYOUT_BOTTOMBAR_EXPAND (IconName): The :code:`layout-bottombar-expand` icon.\n LAYOUT_CARDS (IconName): The :code:`layout-cards` icon.\n LAYOUT_COLLAGE (IconName): The :code:`layout-collage` icon.\n LAYOUT_COLUMNS (IconName): The :code:`layout-columns` icon.\n LAYOUT_DASHBOARD (IconName): The :code:`layout-dashboard` icon.\n LAYOUT_DISTRIBUTE_HORIZONTAL (IconName): The :code:`layout-distribute-horizontal` icon.\n LAYOUT_DISTRIBUTE_VERTICAL (IconName): The :code:`layout-distribute-vertical` icon.\n LAYOUT_GRID (IconName): The :code:`layout-grid` icon.\n LAYOUT_GRID_ADD (IconName): The :code:`layout-grid-add` icon.\n LAYOUT_GRID_REMOVE (IconName): The :code:`layout-grid-remove` icon.\n LAYOUT_KANBAN (IconName): The :code:`layout-kanban` icon.\n LAYOUT_LIST (IconName): The :code:`layout-list` icon.\n LAYOUT_NAVBAR (IconName): The :code:`layout-navbar` icon.\n LAYOUT_NAVBAR_COLLAPSE (IconName): The :code:`layout-navbar-collapse` icon.\n LAYOUT_NAVBAR_EXPAND (IconName): The :code:`layout-navbar-expand` icon.\n LAYOUT_OFF (IconName): The :code:`layout-off` icon.\n LAYOUT_ROWS (IconName): The :code:`layout-rows` icon.\n LAYOUT_SIDEBAR (IconName): The :code:`layout-sidebar` icon.\n LAYOUT_SIDEBAR_LEFT_COLLAPSE (IconName): The :code:`layout-sidebar-left-collapse` icon.\n LAYOUT_SIDEBAR_LEFT_EXPAND (IconName): The :code:`layout-sidebar-left-expand` icon.\n LAYOUT_SIDEBAR_RIGHT (IconName): The :code:`layout-sidebar-right` icon.\n LAYOUT_SIDEBAR_RIGHT_COLLAPSE (IconName): The :code:`layout-sidebar-right-collapse` icon.\n LAYOUT_SIDEBAR_RIGHT_EXPAND (IconName): The :code:`layout-sidebar-right-expand` icon.\n LEAF (IconName): The :code:`leaf` icon.\n LEAF_OFF (IconName): The :code:`leaf-off` icon.\n LEGO (IconName): The :code:`lego` icon.\n LEGO_OFF (IconName): The :code:`lego-off` icon.\n LEMON (IconName): The :code:`lemon` icon.\n LEMON_2 (IconName): The :code:`lemon-2` icon.\n LETTER_A (IconName): The :code:`letter-a` icon.\n LETTER_B (IconName): The :code:`letter-b` icon.\n LETTER_C (IconName): The :code:`letter-c` icon.\n LETTER_CASE (IconName): The :code:`letter-case` icon.\n LETTER_CASE_LOWER (IconName): The :code:`letter-case-lower` icon.\n LETTER_CASE_TOGGLE (IconName): The :code:`letter-case-toggle` icon.\n LETTER_CASE_UPPER (IconName): The :code:`letter-case-upper` icon.\n LETTER_D (IconName): The :code:`letter-d` icon.\n LETTER_E (IconName): The :code:`letter-e` icon.\n LETTER_F (IconName): The :code:`letter-f` icon.\n LETTER_G (IconName): The :code:`letter-g` icon.\n LETTER_H (IconName): The :code:`letter-h` icon.\n LETTER_I (IconName): The :code:`letter-i` icon.\n LETTER_J (IconName): The :code:`letter-j` icon.\n LETTER_K (IconName): The :code:`letter-k` icon.\n LETTER_L (IconName): The :code:`letter-l` icon.\n LETTER_M (IconName): The :code:`letter-m` icon.\n LETTER_N (IconName): The :code:`letter-n` icon.\n LETTER_O (IconName): The :code:`letter-o` icon.\n LETTER_P (IconName): The :code:`letter-p` icon.\n LETTER_Q (IconName): The :code:`letter-q` icon.\n LETTER_R (IconName): The :code:`letter-r` icon.\n LETTER_S (IconName): The :code:`letter-s` icon.\n LETTER_SPACING (IconName): The :code:`letter-spacing` icon.\n LETTER_T (IconName): The :code:`letter-t` icon.\n LETTER_U (IconName): The :code:`letter-u` icon.\n LETTER_V (IconName): The :code:`letter-v` icon.\n LETTER_W (IconName): The :code:`letter-w` icon.\n LETTER_X (IconName): The :code:`letter-x` icon.\n LETTER_Y (IconName): The :code:`letter-y` icon.\n LETTER_Z (IconName): The :code:`letter-z` icon.\n LICENSE (IconName): The :code:`license` icon.\n LICENSE_OFF (IconName): The :code:`license-off` icon.\n LIFEBUOY (IconName): The :code:`lifebuoy` icon.\n LIFEBUOY_OFF (IconName): The :code:`lifebuoy-off` icon.\n LIGHTER (IconName): The :code:`lighter` icon.\n LINE (IconName): The :code:`line` icon.\n LINE_DASHED (IconName): The :code:`line-dashed` icon.\n LINE_DOTTED (IconName): The :code:`line-dotted` icon.\n LINE_HEIGHT (IconName): The :code:`line-height` icon.\n LINK (IconName): The :code:`link` icon.\n LINK_OFF (IconName): The :code:`link-off` icon.\n LIST (IconName): The :code:`list` icon.\n LIST_CHECK (IconName): The :code:`list-check` icon.\n LIST_DETAILS (IconName): The :code:`list-details` icon.\n LIST_NUMBERS (IconName): The :code:`list-numbers` icon.\n LIST_SEARCH (IconName): The :code:`list-search` icon.\n LIST_TREE (IconName): The :code:`list-tree` icon.\n LIVE_PHOTO (IconName): The :code:`live-photo` icon.\n LIVE_PHOTO_OFF (IconName): The :code:`live-photo-off` icon.\n LIVE_VIEW (IconName): The :code:`live-view` icon.\n LOAD_BALANCER (IconName): The :code:`load-balancer` icon.\n LOADER (IconName): The :code:`loader` icon.\n LOADER_2 (IconName): The :code:`loader-2` icon.\n LOADER_3 (IconName): The :code:`loader-3` icon.\n LOADER_QUARTER (IconName): The :code:`loader-quarter` icon.\n LOCATION (IconName): The :code:`location` icon.\n LOCATION_BROKEN (IconName): The :code:`location-broken` icon.\n LOCATION_FILLED (IconName): The :code:`location-filled` icon.\n LOCATION_OFF (IconName): The :code:`location-off` icon.\n LOCK (IconName): The :code:`lock` icon.\n LOCK_ACCESS (IconName): The :code:`lock-access` icon.\n LOCK_ACCESS_OFF (IconName): The :code:`lock-access-off` icon.\n LOCK_BOLT (IconName): The :code:`lock-bolt` icon.\n LOCK_CANCEL (IconName): The :code:`lock-cancel` icon.\n LOCK_CHECK (IconName): The :code:`lock-check` icon.\n LOCK_CODE (IconName): The :code:`lock-code` icon.\n LOCK_COG (IconName): The :code:`lock-cog` icon.\n LOCK_DOLLAR (IconName): The :code:`lock-dollar` icon.\n LOCK_DOWN (IconName): The :code:`lock-down` icon.\n LOCK_EXCLAMATION (IconName): The :code:`lock-exclamation` icon.\n LOCK_HEART (IconName): The :code:`lock-heart` icon.\n LOCK_MINUS (IconName): The :code:`lock-minus` icon.\n LOCK_OFF (IconName): The :code:`lock-off` icon.\n LOCK_OPEN (IconName): The :code:`lock-open` icon.\n LOCK_OPEN_OFF (IconName): The :code:`lock-open-off` icon.\n LOCK_PAUSE (IconName): The :code:`lock-pause` icon.\n LOCK_PIN (IconName): The :code:`lock-pin` icon.\n LOCK_PLUS (IconName): The :code:`lock-plus` icon.\n LOCK_QUESTION (IconName): The :code:`lock-question` icon.\n LOCK_SEARCH (IconName): The :code:`lock-search` icon.\n LOCK_SHARE (IconName): The :code:`lock-share` icon.\n LOCK_SQUARE (IconName): The :code:`lock-square` icon.\n LOCK_SQUARE_ROUNDED (IconName): The :code:`lock-square-rounded` icon.\n LOCK_SQUARE_ROUNDED_FILLED (IconName): The :code:`lock-square-rounded-filled` icon.\n LOCK_STAR (IconName): The :code:`lock-star` icon.\n LOCK_UP (IconName): The :code:`lock-up` icon.\n LOCK_X (IconName): The :code:`lock-x` icon.\n LOGIC_AND (IconName): The :code:`logic-and` icon.\n LOGIC_BUFFER (IconName): The :code:`logic-buffer` icon.\n LOGIC_NAND (IconName): The :code:`logic-nand` icon.\n LOGIC_NOR (IconName): The :code:`logic-nor` icon.\n LOGIC_NOT (IconName): The :code:`logic-not` icon.\n LOGIC_OR (IconName): The :code:`logic-or` icon.\n LOGIC_XNOR (IconName): The :code:`logic-xnor` icon.\n LOGIC_XOR (IconName): The :code:`logic-xor` icon.\n LOGIN (IconName): The :code:`login` icon.\n LOGOUT (IconName): The :code:`logout` icon.\n LOGOUT_2 (IconName): The :code:`logout-2` icon.\n LOLLIPOP (IconName): The :code:`lollipop` icon.\n LOLLIPOP_OFF (IconName): The :code:`lollipop-off` icon.\n LUGGAGE (IconName): The :code:`luggage` icon.\n LUGGAGE_OFF (IconName): The :code:`luggage-off` icon.\n LUNGS (IconName): The :code:`lungs` icon.\n LUNGS_OFF (IconName): The :code:`lungs-off` icon.\n MACRO (IconName): The :code:`macro` icon.\n MACRO_OFF (IconName): The :code:`macro-off` icon.\n MAGNET (IconName): The :code:`magnet` icon.\n MAGNET_OFF (IconName): The :code:`magnet-off` icon.\n MAIL (IconName): The :code:`mail` icon.\n MAIL_AI (IconName): The :code:`mail-ai` icon.\n MAIL_BOLT (IconName): The :code:`mail-bolt` icon.\n MAIL_CANCEL (IconName): The :code:`mail-cancel` icon.\n MAIL_CHECK (IconName): The :code:`mail-check` icon.\n MAIL_CODE (IconName): The :code:`mail-code` icon.\n MAIL_COG (IconName): The :code:`mail-cog` icon.\n MAIL_DOLLAR (IconName): The :code:`mail-dollar` icon.\n MAIL_DOWN (IconName): The :code:`mail-down` icon.\n MAIL_EXCLAMATION (IconName): The :code:`mail-exclamation` icon.\n MAIL_FAST (IconName): The :code:`mail-fast` icon.\n MAIL_FILLED (IconName): The :code:`mail-filled` icon.\n MAIL_FORWARD (IconName): The :code:`mail-forward` icon.\n MAIL_HEART (IconName): The :code:`mail-heart` icon.\n MAIL_MINUS (IconName): The :code:`mail-minus` icon.\n MAIL_OFF (IconName): The :code:`mail-off` icon.\n MAIL_OPENED (IconName): The :code:`mail-opened` icon.\n MAIL_OPENED_FILLED (IconName): The :code:`mail-opened-filled` icon.\n MAIL_PAUSE (IconName): The :code:`mail-pause` icon.\n MAIL_PIN (IconName): The :code:`mail-pin` icon.\n MAIL_PLUS (IconName): The :code:`mail-plus` icon.\n MAIL_QUESTION (IconName): The :code:`mail-question` icon.\n MAIL_SEARCH (IconName): The :code:`mail-search` icon.\n MAIL_SHARE (IconName): The :code:`mail-share` icon.\n MAIL_STAR (IconName): The :code:`mail-star` icon.\n MAIL_UP (IconName): The :code:`mail-up` icon.\n MAIL_X (IconName): The :code:`mail-x` icon.\n MAILBOX (IconName): The :code:`mailbox` icon.\n MAILBOX_OFF (IconName): The :code:`mailbox-off` icon.\n MAN (IconName): The :code:`man` icon.\n MANUAL_GEARBOX (IconName): The :code:`manual-gearbox` icon.\n MAP (IconName): The :code:`map` icon.\n MAP_2 (IconName): The :code:`map-2` icon.\n MAP_OFF (IconName): The :code:`map-off` icon.\n MAP_PIN (IconName): The :code:`map-pin` icon.\n MAP_PIN_BOLT (IconName): The :code:`map-pin-bolt` icon.\n MAP_PIN_CANCEL (IconName): The :code:`map-pin-cancel` icon.\n MAP_PIN_CHECK (IconName): The :code:`map-pin-check` icon.\n MAP_PIN_CODE (IconName): The :code:`map-pin-code` icon.\n MAP_PIN_COG (IconName): The :code:`map-pin-cog` icon.\n MAP_PIN_DOLLAR (IconName): The :code:`map-pin-dollar` icon.\n MAP_PIN_DOWN (IconName): The :code:`map-pin-down` icon.\n MAP_PIN_EXCLAMATION (IconName): The :code:`map-pin-exclamation` icon.\n MAP_PIN_FILLED (IconName): The :code:`map-pin-filled` icon.\n MAP_PIN_HEART (IconName): The :code:`map-pin-heart` icon.\n MAP_PIN_MINUS (IconName): The :code:`map-pin-minus` icon.\n MAP_PIN_OFF (IconName): The :code:`map-pin-off` icon.\n MAP_PIN_PAUSE (IconName): The :code:`map-pin-pause` icon.\n MAP_PIN_PIN (IconName): The :code:`map-pin-pin` icon.\n MAP_PIN_PLUS (IconName): The :code:`map-pin-plus` icon.\n MAP_PIN_QUESTION (IconName): The :code:`map-pin-question` icon.\n MAP_PIN_SEARCH (IconName): The :code:`map-pin-search` icon.\n MAP_PIN_SHARE (IconName): The :code:`map-pin-share` icon.\n MAP_PIN_STAR (IconName): The :code:`map-pin-star` icon.\n MAP_PIN_UP (IconName): The :code:`map-pin-up` icon.\n MAP_PIN_X (IconName): The :code:`map-pin-x` icon.\n MAP_PINS (IconName): The :code:`map-pins` icon.\n MAP_SEARCH (IconName): The :code:`map-search` icon.\n MARKDOWN (IconName): The :code:`markdown` icon.\n MARKDOWN_OFF (IconName): The :code:`markdown-off` icon.\n MARQUEE (IconName): The :code:`marquee` icon.\n MARQUEE_2 (IconName): The :code:`marquee-2` icon.\n MARQUEE_OFF (IconName): The :code:`marquee-off` icon.\n MARS (IconName): The :code:`mars` icon.\n MASK (IconName): The :code:`mask` icon.\n MASK_OFF (IconName): The :code:`mask-off` icon.\n MASKS_THEATER (IconName): The :code:`masks-theater` icon.\n MASKS_THEATER_OFF (IconName): The :code:`masks-theater-off` icon.\n MASSAGE (IconName): The :code:`massage` icon.\n MATCHSTICK (IconName): The :code:`matchstick` icon.\n MATH (IconName): The :code:`math` icon.\n MATH_1_DIVIDE_2 (IconName): The :code:`math-1-divide-2` icon.\n MATH_1_DIVIDE_3 (IconName): The :code:`math-1-divide-3` icon.\n MATH_AVG (IconName): The :code:`math-avg` icon.\n MATH_EQUAL_GREATER (IconName): The :code:`math-equal-greater` icon.\n MATH_EQUAL_LOWER (IconName): The :code:`math-equal-lower` icon.\n MATH_FUNCTION (IconName): The :code:`math-function` icon.\n MATH_FUNCTION_OFF (IconName): The :code:`math-function-off` icon.\n MATH_FUNCTION_Y (IconName): The :code:`math-function-y` icon.\n MATH_GREATER (IconName): The :code:`math-greater` icon.\n MATH_INTEGRAL (IconName): The :code:`math-integral` icon.\n MATH_INTEGRAL_X (IconName): The :code:`math-integral-x` icon.\n MATH_INTEGRALS (IconName): The :code:`math-integrals` icon.\n MATH_LOWER (IconName): The :code:`math-lower` icon.\n MATH_MAX (IconName): The :code:`math-max` icon.\n MATH_MIN (IconName): The :code:`math-min` icon.\n MATH_NOT (IconName): The :code:`math-not` icon.\n MATH_OFF (IconName): The :code:`math-off` icon.\n MATH_PI (IconName): The :code:`math-pi` icon.\n MATH_PI_DIVIDE_2 (IconName): The :code:`math-pi-divide-2` icon.\n MATH_SYMBOLS (IconName): The :code:`math-symbols` icon.\n MATH_X_DIVIDE_2 (IconName): The :code:`math-x-divide-2` icon.\n MATH_X_DIVIDE_Y (IconName): The :code:`math-x-divide-y` icon.\n MATH_X_DIVIDE_Y_2 (IconName): The :code:`math-x-divide-y-2` icon.\n MATH_X_MINUS_X (IconName): The :code:`math-x-minus-x` icon.\n MATH_X_MINUS_Y (IconName): The :code:`math-x-minus-y` icon.\n MATH_X_PLUS_X (IconName): The :code:`math-x-plus-x` icon.\n MATH_X_PLUS_Y (IconName): The :code:`math-x-plus-y` icon.\n MATH_XY (IconName): The :code:`math-xy` icon.\n MATH_Y_MINUS_Y (IconName): The :code:`math-y-minus-y` icon.\n MATH_Y_PLUS_Y (IconName): The :code:`math-y-plus-y` icon.\n MAXIMIZE (IconName): The :code:`maximize` icon.\n MAXIMIZE_OFF (IconName): The :code:`maximize-off` icon.\n MEAT (IconName): The :code:`meat` icon.\n MEAT_OFF (IconName): The :code:`meat-off` icon.\n MEDAL (IconName): The :code:`medal` icon.\n MEDAL_2 (IconName): The :code:`medal-2` icon.\n MEDICAL_CROSS (IconName): The :code:`medical-cross` icon.\n MEDICAL_CROSS_CIRCLE (IconName): The :code:`medical-cross-circle` icon.\n MEDICAL_CROSS_FILLED (IconName): The :code:`medical-cross-filled` icon.\n MEDICAL_CROSS_OFF (IconName): The :code:`medical-cross-off` icon.\n MEDICINE_SYRUP (IconName): The :code:`medicine-syrup` icon.\n MEEPLE (IconName): The :code:`meeple` icon.\n MENORAH (IconName): The :code:`menorah` icon.\n MENU (IconName): The :code:`menu` icon.\n MENU_2 (IconName): The :code:`menu-2` icon.\n MENU_DEEP (IconName): The :code:`menu-deep` icon.\n MENU_ORDER (IconName): The :code:`menu-order` icon.\n MESSAGE (IconName): The :code:`message` icon.\n MESSAGE_2 (IconName): The :code:`message-2` icon.\n MESSAGE_2_BOLT (IconName): The :code:`message-2-bolt` icon.\n MESSAGE_2_CANCEL (IconName): The :code:`message-2-cancel` icon.\n MESSAGE_2_CHECK (IconName): The :code:`message-2-check` icon.\n MESSAGE_2_CODE (IconName): The :code:`message-2-code` icon.\n MESSAGE_2_COG (IconName): The :code:`message-2-cog` icon.\n MESSAGE_2_DOLLAR (IconName): The :code:`message-2-dollar` icon.\n MESSAGE_2_DOWN (IconName): The :code:`message-2-down` icon.\n MESSAGE_2_EXCLAMATION (IconName): The :code:`message-2-exclamation` icon.\n MESSAGE_2_HEART (IconName): The :code:`message-2-heart` icon.\n MESSAGE_2_MINUS (IconName): The :code:`message-2-minus` icon.\n MESSAGE_2_OFF (IconName): The :code:`message-2-off` icon.\n MESSAGE_2_PAUSE (IconName): The :code:`message-2-pause` icon.\n MESSAGE_2_PIN (IconName): The :code:`message-2-pin` icon.\n MESSAGE_2_PLUS (IconName): The :code:`message-2-plus` icon.\n MESSAGE_2_QUESTION (IconName): The :code:`message-2-question` icon.\n MESSAGE_2_SEARCH (IconName): The :code:`message-2-search` icon.\n MESSAGE_2_SHARE (IconName): The :code:`message-2-share` icon.\n MESSAGE_2_STAR (IconName): The :code:`message-2-star` icon.\n MESSAGE_2_UP (IconName): The :code:`message-2-up` icon.\n MESSAGE_2_X (IconName): The :code:`message-2-x` icon.\n MESSAGE_BOLT (IconName): The :code:`message-bolt` icon.\n MESSAGE_CANCEL (IconName): The :code:`message-cancel` icon.\n MESSAGE_CHATBOT (IconName): The :code:`message-chatbot` icon.\n MESSAGE_CHECK (IconName): The :code:`message-check` icon.\n MESSAGE_CIRCLE (IconName): The :code:`message-circle` icon.\n MESSAGE_CIRCLE_2 (IconName): The :code:`message-circle-2` icon.\n MESSAGE_CIRCLE_2_FILLED (IconName): The :code:`message-circle-2-filled` icon.\n MESSAGE_CIRCLE_BOLT (IconName): The :code:`message-circle-bolt` icon.\n MESSAGE_CIRCLE_CANCEL (IconName): The :code:`message-circle-cancel` icon.\n MESSAGE_CIRCLE_CHECK (IconName): The :code:`message-circle-check` icon.\n MESSAGE_CIRCLE_CODE (IconName): The :code:`message-circle-code` icon.\n MESSAGE_CIRCLE_COG (IconName): The :code:`message-circle-cog` icon.\n MESSAGE_CIRCLE_DOLLAR (IconName): The :code:`message-circle-dollar` icon.\n MESSAGE_CIRCLE_DOWN (IconName): The :code:`message-circle-down` icon.\n MESSAGE_CIRCLE_EXCLAMATION (IconName): The :code:`message-circle-exclamation` icon.\n MESSAGE_CIRCLE_HEART (IconName): The :code:`message-circle-heart` icon.\n MESSAGE_CIRCLE_MINUS (IconName): The :code:`message-circle-minus` icon.\n MESSAGE_CIRCLE_OFF (IconName): The :code:`message-circle-off` icon.\n MESSAGE_CIRCLE_PAUSE (IconName): The :code:`message-circle-pause` icon.\n MESSAGE_CIRCLE_PIN (IconName): The :code:`message-circle-pin` icon.\n MESSAGE_CIRCLE_PLUS (IconName): The :code:`message-circle-plus` icon.\n MESSAGE_CIRCLE_QUESTION (IconName): The :code:`message-circle-question` icon.\n MESSAGE_CIRCLE_SEARCH (IconName): The :code:`message-circle-search` icon.\n MESSAGE_CIRCLE_SHARE (IconName): The :code:`message-circle-share` icon.\n MESSAGE_CIRCLE_STAR (IconName): The :code:`message-circle-star` icon.\n MESSAGE_CIRCLE_UP (IconName): The :code:`message-circle-up` icon.\n MESSAGE_CIRCLE_X (IconName): The :code:`message-circle-x` icon.\n MESSAGE_CODE (IconName): The :code:`message-code` icon.\n MESSAGE_COG (IconName): The :code:`message-cog` icon.\n MESSAGE_DOLLAR (IconName): The :code:`message-dollar` icon.\n MESSAGE_DOTS (IconName): The :code:`message-dots` icon.\n MESSAGE_DOWN (IconName): The :code:`message-down` icon.\n MESSAGE_EXCLAMATION (IconName): The :code:`message-exclamation` icon.\n MESSAGE_FORWARD (IconName): The :code:`message-forward` icon.\n MESSAGE_HEART (IconName): The :code:`message-heart` icon.\n MESSAGE_LANGUAGE (IconName): The :code:`message-language` icon.\n MESSAGE_MINUS (IconName): The :code:`message-minus` icon.\n MESSAGE_OFF (IconName): The :code:`message-off` icon.\n MESSAGE_PAUSE (IconName): The :code:`message-pause` icon.\n MESSAGE_PIN (IconName): The :code:`message-pin` icon.\n MESSAGE_PLUS (IconName): The :code:`message-plus` icon.\n MESSAGE_QUESTION (IconName): The :code:`message-question` icon.\n MESSAGE_REPORT (IconName): The :code:`message-report` icon.\n MESSAGE_SEARCH (IconName): The :code:`message-search` icon.\n MESSAGE_SHARE (IconName): The :code:`message-share` icon.\n MESSAGE_STAR (IconName): The :code:`message-star` icon.\n MESSAGE_UP (IconName): The :code:`message-up` icon.\n MESSAGE_X (IconName): The :code:`message-x` icon.\n MESSAGES (IconName): The :code:`messages` icon.\n MESSAGES_OFF (IconName): The :code:`messages-off` icon.\n METEOR (IconName): The :code:`meteor` icon.\n METEOR_OFF (IconName): The :code:`meteor-off` icon.\n MICHELIN_BIB_GOURMAND (IconName): The :code:`michelin-bib-gourmand` icon.\n MICHELIN_STAR (IconName): The :code:`michelin-star` icon.\n MICHELIN_STAR_GREEN (IconName): The :code:`michelin-star-green` icon.\n MICKEY (IconName): The :code:`mickey` icon.\n MICKEY_FILLED (IconName): The :code:`mickey-filled` icon.\n MICROPHONE (IconName): The :code:`microphone` icon.\n MICROPHONE_2 (IconName): The :code:`microphone-2` icon.\n MICROPHONE_2_OFF (IconName): The :code:`microphone-2-off` icon.\n MICROPHONE_OFF (IconName): The :code:`microphone-off` icon.\n MICROSCOPE (IconName): The :code:`microscope` icon.\n MICROSCOPE_OFF (IconName): The :code:`microscope-off` icon.\n MICROWAVE (IconName): The :code:`microwave` icon.\n MICROWAVE_OFF (IconName): The :code:`microwave-off` icon.\n MILITARY_AWARD (IconName): The :code:`military-award` icon.\n MILITARY_RANK (IconName): The :code:`military-rank` icon.\n MILK (IconName): The :code:`milk` icon.\n MILK_OFF (IconName): The :code:`milk-off` icon.\n MILKSHAKE (IconName): The :code:`milkshake` icon.\n MINIMIZE (IconName): The :code:`minimize` icon.\n MINUS (IconName): The :code:`minus` icon.\n MINUS_VERTICAL (IconName): The :code:`minus-vertical` icon.\n MIST (IconName): The :code:`mist` icon.\n MIST_OFF (IconName): The :code:`mist-off` icon.\n MOBILEDATA (IconName): The :code:`mobiledata` icon.\n MOBILEDATA_OFF (IconName): The :code:`mobiledata-off` icon.\n MONEYBAG (IconName): The :code:`moneybag` icon.\n MOOD_ANGRY (IconName): The :code:`mood-angry` icon.\n MOOD_ANNOYED (IconName): The :code:`mood-annoyed` icon.\n MOOD_ANNOYED_2 (IconName): The :code:`mood-annoyed-2` icon.\n MOOD_BOY (IconName): The :code:`mood-boy` icon.\n MOOD_CHECK (IconName): The :code:`mood-check` icon.\n MOOD_COG (IconName): The :code:`mood-cog` icon.\n MOOD_CONFUZED (IconName): The :code:`mood-confuzed` icon.\n MOOD_CONFUZED_FILLED (IconName): The :code:`mood-confuzed-filled` icon.\n MOOD_CRAZY_HAPPY (IconName): The :code:`mood-crazy-happy` icon.\n MOOD_CRY (IconName): The :code:`mood-cry` icon.\n MOOD_DOLLAR (IconName): The :code:`mood-dollar` icon.\n MOOD_EDIT (IconName): The :code:`mood-edit` icon.\n MOOD_EMPTY (IconName): The :code:`mood-empty` icon.\n MOOD_EMPTY_FILLED (IconName): The :code:`mood-empty-filled` icon.\n MOOD_HAPPY (IconName): The :code:`mood-happy` icon.\n MOOD_HAPPY_FILLED (IconName): The :code:`mood-happy-filled` icon.\n MOOD_HEART (IconName): The :code:`mood-heart` icon.\n MOOD_KID (IconName): The :code:`mood-kid` icon.\n MOOD_KID_FILLED (IconName): The :code:`mood-kid-filled` icon.\n MOOD_LOOK_LEFT (IconName): The :code:`mood-look-left` icon.\n MOOD_LOOK_RIGHT (IconName): The :code:`mood-look-right` icon.\n MOOD_MINUS (IconName): The :code:`mood-minus` icon.\n MOOD_NERD (IconName): The :code:`mood-nerd` icon.\n MOOD_NERVOUS (IconName): The :code:`mood-nervous` icon.\n MOOD_NEUTRAL (IconName): The :code:`mood-neutral` icon.\n MOOD_NEUTRAL_FILLED (IconName): The :code:`mood-neutral-filled` icon.\n MOOD_OFF (IconName): The :code:`mood-off` icon.\n MOOD_PIN (IconName): The :code:`mood-pin` icon.\n MOOD_PLUS (IconName): The :code:`mood-plus` icon.\n MOOD_SAD (IconName): The :code:`mood-sad` icon.\n MOOD_SAD_2 (IconName): The :code:`mood-sad-2` icon.\n MOOD_SAD_DIZZY (IconName): The :code:`mood-sad-dizzy` icon.\n MOOD_SAD_FILLED (IconName): The :code:`mood-sad-filled` icon.\n MOOD_SAD_SQUINT (IconName): The :code:`mood-sad-squint` icon.\n MOOD_SEARCH (IconName): The :code:`mood-search` icon.\n MOOD_SHARE (IconName): The :code:`mood-share` icon.\n MOOD_SICK (IconName): The :code:`mood-sick` icon.\n MOOD_SILENCE (IconName): The :code:`mood-silence` icon.\n MOOD_SING (IconName): The :code:`mood-sing` icon.\n MOOD_SMILE (IconName): The :code:`mood-smile` icon.\n MOOD_SMILE_BEAM (IconName): The :code:`mood-smile-beam` icon.\n MOOD_SMILE_DIZZY (IconName): The :code:`mood-smile-dizzy` icon.\n MOOD_SMILE_FILLED (IconName): The :code:`mood-smile-filled` icon.\n MOOD_SUPRISED (IconName): The :code:`mood-suprised` icon.\n MOOD_TONGUE (IconName): The :code:`mood-tongue` icon.\n MOOD_TONGUE_WINK (IconName): The :code:`mood-tongue-wink` icon.\n MOOD_TONGUE_WINK_2 (IconName): The :code:`mood-tongue-wink-2` icon.\n MOOD_UNAMUSED (IconName): The :code:`mood-unamused` icon.\n MOOD_UP (IconName): The :code:`mood-up` icon.\n MOOD_WINK (IconName): The :code:`mood-wink` icon.\n MOOD_WINK_2 (IconName): The :code:`mood-wink-2` icon.\n MOOD_WRRR (IconName): The :code:`mood-wrrr` icon.\n MOOD_X (IconName): The :code:`mood-x` icon.\n MOOD_XD (IconName): The :code:`mood-xd` icon.\n MOON (IconName): The :code:`moon` icon.\n MOON_2 (IconName): The :code:`moon-2` icon.\n MOON_FILLED (IconName): The :code:`moon-filled` icon.\n MOON_OFF (IconName): The :code:`moon-off` icon.\n MOON_STARS (IconName): The :code:`moon-stars` icon.\n MOPED (IconName): The :code:`moped` icon.\n MOTORBIKE (IconName): The :code:`motorbike` icon.\n MOUNTAIN (IconName): The :code:`mountain` icon.\n MOUNTAIN_OFF (IconName): The :code:`mountain-off` icon.\n MOUSE (IconName): The :code:`mouse` icon.\n MOUSE_2 (IconName): The :code:`mouse-2` icon.\n MOUSE_OFF (IconName): The :code:`mouse-off` icon.\n MOUSTACHE (IconName): The :code:`moustache` icon.\n MOVIE (IconName): The :code:`movie` icon.\n MOVIE_OFF (IconName): The :code:`movie-off` icon.\n MUG (IconName): The :code:`mug` icon.\n MUG_OFF (IconName): The :code:`mug-off` icon.\n MULTIPLIER_0_5X (IconName): The :code:`multiplier-0-5x` icon.\n MULTIPLIER_1_5X (IconName): The :code:`multiplier-1-5x` icon.\n MULTIPLIER_1X (IconName): The :code:`multiplier-1x` icon.\n MULTIPLIER_2X (IconName): The :code:`multiplier-2x` icon.\n MUSHROOM (IconName): The :code:`mushroom` icon.\n MUSHROOM_FILLED (IconName): The :code:`mushroom-filled` icon.\n MUSHROOM_OFF (IconName): The :code:`mushroom-off` icon.\n MUSIC (IconName): The :code:`music` icon.\n MUSIC_OFF (IconName): The :code:`music-off` icon.\n NAVIGATION (IconName): The :code:`navigation` icon.\n NAVIGATION_FILLED (IconName): The :code:`navigation-filled` icon.\n NAVIGATION_NORTH (IconName): The :code:`navigation-north` icon.\n NAVIGATION_OFF (IconName): The :code:`navigation-off` icon.\n NEEDLE (IconName): The :code:`needle` icon.\n NEEDLE_THREAD (IconName): The :code:`needle-thread` icon.\n NETWORK (IconName): The :code:`network` icon.\n NETWORK_OFF (IconName): The :code:`network-off` icon.\n NEW_SECTION (IconName): The :code:`new-section` icon.\n NEWS (IconName): The :code:`news` icon.\n NEWS_OFF (IconName): The :code:`news-off` icon.\n NFC (IconName): The :code:`nfc` icon.\n NFC_OFF (IconName): The :code:`nfc-off` icon.\n NO_COPYRIGHT (IconName): The :code:`no-copyright` icon.\n NO_CREATIVE_COMMONS (IconName): The :code:`no-creative-commons` icon.\n NO_DERIVATIVES (IconName): The :code:`no-derivatives` icon.\n NORTH_STAR (IconName): The :code:`north-star` icon.\n NOTE (IconName): The :code:`note` icon.\n NOTE_OFF (IconName): The :code:`note-off` icon.\n NOTEBOOK (IconName): The :code:`notebook` icon.\n NOTEBOOK_OFF (IconName): The :code:`notebook-off` icon.\n NOTES (IconName): The :code:`notes` icon.\n NOTES_OFF (IconName): The :code:`notes-off` icon.\n NOTIFICATION (IconName): The :code:`notification` icon.\n NOTIFICATION_OFF (IconName): The :code:`notification-off` icon.\n NUMBER (IconName): The :code:`number` icon.\n NUMBER_0 (IconName): The :code:`number-0` icon.\n NUMBER_1 (IconName): The :code:`number-1` icon.\n NUMBER_2 (IconName): The :code:`number-2` icon.\n NUMBER_3 (IconName): The :code:`number-3` icon.\n NUMBER_4 (IconName): The :code:`number-4` icon.\n NUMBER_5 (IconName): The :code:`number-5` icon.\n NUMBER_6 (IconName): The :code:`number-6` icon.\n NUMBER_7 (IconName): The :code:`number-7` icon.\n NUMBER_8 (IconName): The :code:`number-8` icon.\n NUMBER_9 (IconName): The :code:`number-9` icon.\n NUMBERS (IconName): The :code:`numbers` icon.\n NURSE (IconName): The :code:`nurse` icon.\n OCTAGON (IconName): The :code:`octagon` icon.\n OCTAGON_FILLED (IconName): The :code:`octagon-filled` icon.\n OCTAGON_OFF (IconName): The :code:`octagon-off` icon.\n OCTAHEDRON (IconName): The :code:`octahedron` icon.\n OCTAHEDRON_OFF (IconName): The :code:`octahedron-off` icon.\n OCTAHEDRON_PLUS (IconName): The :code:`octahedron-plus` icon.\n OLD (IconName): The :code:`old` icon.\n OLYMPICS (IconName): The :code:`olympics` icon.\n OLYMPICS_OFF (IconName): The :code:`olympics-off` icon.\n OM (IconName): The :code:`om` icon.\n OMEGA (IconName): The :code:`omega` icon.\n OUTBOUND (IconName): The :code:`outbound` icon.\n OUTLET (IconName): The :code:`outlet` icon.\n OVAL (IconName): The :code:`oval` icon.\n OVAL_FILLED (IconName): The :code:`oval-filled` icon.\n OVAL_VERTICAL (IconName): The :code:`oval-vertical` icon.\n OVAL_VERTICAL_FILLED (IconName): The :code:`oval-vertical-filled` icon.\n OVERLINE (IconName): The :code:`overline` icon.\n PACKAGE (IconName): The :code:`package` icon.\n PACKAGE_EXPORT (IconName): The :code:`package-export` icon.\n PACKAGE_IMPORT (IconName): The :code:`package-import` icon.\n PACKAGE_OFF (IconName): The :code:`package-off` icon.\n PACKAGES (IconName): The :code:`packages` icon.\n PACMAN (IconName): The :code:`pacman` icon.\n PAGE_BREAK (IconName): The :code:`page-break` icon.\n PAINT (IconName): The :code:`paint` icon.\n PAINT_FILLED (IconName): The :code:`paint-filled` icon.\n PAINT_OFF (IconName): The :code:`paint-off` icon.\n PALETTE (IconName): The :code:`palette` icon.\n PALETTE_OFF (IconName): The :code:`palette-off` icon.\n PANORAMA_HORIZONTAL (IconName): The :code:`panorama-horizontal` icon.\n PANORAMA_HORIZONTAL_OFF (IconName): The :code:`panorama-horizontal-off` icon.\n PANORAMA_VERTICAL (IconName): The :code:`panorama-vertical` icon.\n PANORAMA_VERTICAL_OFF (IconName): The :code:`panorama-vertical-off` icon.\n PAPER_BAG (IconName): The :code:`paper-bag` icon.\n PAPER_BAG_OFF (IconName): The :code:`paper-bag-off` icon.\n PAPERCLIP (IconName): The :code:`paperclip` icon.\n PARACHUTE (IconName): The :code:`parachute` icon.\n PARACHUTE_OFF (IconName): The :code:`parachute-off` icon.\n PARENTHESES (IconName): The :code:`parentheses` icon.\n PARENTHESES_OFF (IconName): The :code:`parentheses-off` icon.\n PARKING (IconName): The :code:`parking` icon.\n PARKING_OFF (IconName): The :code:`parking-off` icon.\n PASSWORD (IconName): The :code:`password` icon.\n PAW (IconName): The :code:`paw` icon.\n PAW_FILLED (IconName): The :code:`paw-filled` icon.\n PAW_OFF (IconName): The :code:`paw-off` icon.\n PDF (IconName): The :code:`pdf` icon.\n PEACE (IconName): The :code:`peace` icon.\n PENCIL (IconName): The :code:`pencil` icon.\n PENCIL_MINUS (IconName): The :code:`pencil-minus` icon.\n PENCIL_OFF (IconName): The :code:`pencil-off` icon.\n PENCIL_PLUS (IconName): The :code:`pencil-plus` icon.\n PENNANT (IconName): The :code:`pennant` icon.\n PENNANT_2 (IconName): The :code:`pennant-2` icon.\n PENNANT_2_FILLED (IconName): The :code:`pennant-2-filled` icon.\n PENNANT_FILLED (IconName): The :code:`pennant-filled` icon.\n PENNANT_OFF (IconName): The :code:`pennant-off` icon.\n PENTAGON (IconName): The :code:`pentagon` icon.\n PENTAGON_FILLED (IconName): The :code:`pentagon-filled` icon.\n PENTAGON_OFF (IconName): The :code:`pentagon-off` icon.\n PENTAGRAM (IconName): The :code:`pentagram` icon.\n PEPPER (IconName): The :code:`pepper` icon.\n PEPPER_OFF (IconName): The :code:`pepper-off` icon.\n PERCENTAGE (IconName): The :code:`percentage` icon.\n PERFUME (IconName): The :code:`perfume` icon.\n PERSPECTIVE (IconName): The :code:`perspective` icon.\n PERSPECTIVE_OFF (IconName): The :code:`perspective-off` icon.\n PHONE (IconName): The :code:`phone` icon.\n PHONE_CALL (IconName): The :code:`phone-call` icon.\n PHONE_CALLING (IconName): The :code:`phone-calling` icon.\n PHONE_CHECK (IconName): The :code:`phone-check` icon.\n PHONE_FILLED (IconName): The :code:`phone-filled` icon.\n PHONE_INCOMING (IconName): The :code:`phone-incoming` icon.\n PHONE_OFF (IconName): The :code:`phone-off` icon.\n PHONE_OUTGOING (IconName): The :code:`phone-outgoing` icon.\n PHONE_PAUSE (IconName): The :code:`phone-pause` icon.\n PHONE_PLUS (IconName): The :code:`phone-plus` icon.\n PHONE_X (IconName): The :code:`phone-x` icon.\n PHOTO (IconName): The :code:`photo` icon.\n PHOTO_AI (IconName): The :code:`photo-ai` icon.\n PHOTO_BOLT (IconName): The :code:`photo-bolt` icon.\n PHOTO_CANCEL (IconName): The :code:`photo-cancel` icon.\n PHOTO_CHECK (IconName): The :code:`photo-check` icon.\n PHOTO_CODE (IconName): The :code:`photo-code` icon.\n PHOTO_COG (IconName): The :code:`photo-cog` icon.\n PHOTO_DOLLAR (IconName): The :code:`photo-dollar` icon.\n PHOTO_DOWN (IconName): The :code:`photo-down` icon.\n PHOTO_EDIT (IconName): The :code:`photo-edit` icon.\n PHOTO_EXCLAMATION (IconName): The :code:`photo-exclamation` icon.\n PHOTO_FILLED (IconName): The :code:`photo-filled` icon.\n PHOTO_HEART (IconName): The :code:`photo-heart` icon.\n PHOTO_MINUS (IconName): The :code:`photo-minus` icon.\n PHOTO_OFF (IconName): The :code:`photo-off` icon.\n PHOTO_PAUSE (IconName): The :code:`photo-pause` icon.\n PHOTO_PIN (IconName): The :code:`photo-pin` icon.\n PHOTO_PLUS (IconName): The :code:`photo-plus` icon.\n PHOTO_QUESTION (IconName): The :code:`photo-question` icon.\n PHOTO_SEARCH (IconName): The :code:`photo-search` icon.\n PHOTO_SENSOR (IconName): The :code:`photo-sensor` icon.\n PHOTO_SENSOR_2 (IconName): The :code:`photo-sensor-2` icon.\n PHOTO_SENSOR_3 (IconName): The :code:`photo-sensor-3` icon.\n PHOTO_SHARE (IconName): The :code:`photo-share` icon.\n PHOTO_SHIELD (IconName): The :code:`photo-shield` icon.\n PHOTO_STAR (IconName): The :code:`photo-star` icon.\n PHOTO_UP (IconName): The :code:`photo-up` icon.\n PHOTO_X (IconName): The :code:`photo-x` icon.\n PHYSOTHERAPIST (IconName): The :code:`physotherapist` icon.\n PIANO (IconName): The :code:`piano` icon.\n PICK (IconName): The :code:`pick` icon.\n PICTURE_IN_PICTURE (IconName): The :code:`picture-in-picture` icon.\n PICTURE_IN_PICTURE_OFF (IconName): The :code:`picture-in-picture-off` icon.\n PICTURE_IN_PICTURE_ON (IconName): The :code:`picture-in-picture-on` icon.\n PICTURE_IN_PICTURE_TOP (IconName): The :code:`picture-in-picture-top` icon.\n PIG (IconName): The :code:`pig` icon.\n PIG_MONEY (IconName): The :code:`pig-money` icon.\n PIG_OFF (IconName): The :code:`pig-off` icon.\n PILCROW (IconName): The :code:`pilcrow` icon.\n PILL (IconName): The :code:`pill` icon.\n PILL_OFF (IconName): The :code:`pill-off` icon.\n PILLS (IconName): The :code:`pills` icon.\n PIN (IconName): The :code:`pin` icon.\n PIN_FILLED (IconName): The :code:`pin-filled` icon.\n PING_PONG (IconName): The :code:`ping-pong` icon.\n PINNED (IconName): The :code:`pinned` icon.\n PINNED_FILLED (IconName): The :code:`pinned-filled` icon.\n PINNED_OFF (IconName): The :code:`pinned-off` icon.\n PIZZA (IconName): The :code:`pizza` icon.\n PIZZA_OFF (IconName): The :code:`pizza-off` icon.\n PLACEHOLDER (IconName): The :code:`placeholder` icon.\n PLANE (IconName): The :code:`plane` icon.\n PLANE_ARRIVAL (IconName): The :code:`plane-arrival` icon.\n PLANE_DEPARTURE (IconName): The :code:`plane-departure` icon.\n PLANE_INFLIGHT (IconName): The :code:`plane-inflight` icon.\n PLANE_OFF (IconName): The :code:`plane-off` icon.\n PLANE_TILT (IconName): The :code:`plane-tilt` icon.\n PLANET (IconName): The :code:`planet` icon.\n PLANET_OFF (IconName): The :code:`planet-off` icon.\n PLANT (IconName): The :code:`plant` icon.\n PLANT_2 (IconName): The :code:`plant-2` icon.\n PLANT_2_OFF (IconName): The :code:`plant-2-off` icon.\n PLANT_OFF (IconName): The :code:`plant-off` icon.\n PLAY_BASKETBALL (IconName): The :code:`play-basketball` icon.\n PLAY_CARD (IconName): The :code:`play-card` icon.\n PLAY_CARD_OFF (IconName): The :code:`play-card-off` icon.\n PLAY_FOOTBALL (IconName): The :code:`play-football` icon.\n PLAY_HANDBALL (IconName): The :code:`play-handball` icon.\n PLAY_VOLLEYBALL (IconName): The :code:`play-volleyball` icon.\n PLAYER_EJECT (IconName): The :code:`player-eject` icon.\n PLAYER_EJECT_FILLED (IconName): The :code:`player-eject-filled` icon.\n PLAYER_PAUSE (IconName): The :code:`player-pause` icon.\n PLAYER_PAUSE_FILLED (IconName): The :code:`player-pause-filled` icon.\n PLAYER_PLAY (IconName): The :code:`player-play` icon.\n PLAYER_PLAY_FILLED (IconName): The :code:`player-play-filled` icon.\n PLAYER_RECORD (IconName): The :code:`player-record` icon.\n PLAYER_RECORD_FILLED (IconName): The :code:`player-record-filled` icon.\n PLAYER_SKIP_BACK (IconName): The :code:`player-skip-back` icon.\n PLAYER_SKIP_BACK_FILLED (IconName): The :code:`player-skip-back-filled` icon.\n PLAYER_SKIP_FORWARD (IconName): The :code:`player-skip-forward` icon.\n PLAYER_SKIP_FORWARD_FILLED (IconName): The :code:`player-skip-forward-filled` icon.\n PLAYER_STOP (IconName): The :code:`player-stop` icon.\n PLAYER_STOP_FILLED (IconName): The :code:`player-stop-filled` icon.\n PLAYER_TRACK_NEXT (IconName): The :code:`player-track-next` icon.\n PLAYER_TRACK_NEXT_FILLED (IconName): The :code:`player-track-next-filled` icon.\n PLAYER_TRACK_PREV (IconName): The :code:`player-track-prev` icon.\n PLAYER_TRACK_PREV_FILLED (IconName): The :code:`player-track-prev-filled` icon.\n PLAYLIST (IconName): The :code:`playlist` icon.\n PLAYLIST_ADD (IconName): The :code:`playlist-add` icon.\n PLAYLIST_OFF (IconName): The :code:`playlist-off` icon.\n PLAYLIST_X (IconName): The :code:`playlist-x` icon.\n PLAYSTATION_CIRCLE (IconName): The :code:`playstation-circle` icon.\n PLAYSTATION_SQUARE (IconName): The :code:`playstation-square` icon.\n PLAYSTATION_TRIANGLE (IconName): The :code:`playstation-triangle` icon.\n PLAYSTATION_X (IconName): The :code:`playstation-x` icon.\n PLUG (IconName): The :code:`plug` icon.\n PLUG_CONNECTED (IconName): The :code:`plug-connected` icon.\n PLUG_CONNECTED_X (IconName): The :code:`plug-connected-x` icon.\n PLUG_OFF (IconName): The :code:`plug-off` icon.\n PLUG_X (IconName): The :code:`plug-x` icon.\n PLUS (IconName): The :code:`plus` icon.\n PLUS_EQUAL (IconName): The :code:`plus-equal` icon.\n PLUS_MINUS (IconName): The :code:`plus-minus` icon.\n PNG (IconName): The :code:`png` icon.\n PODIUM (IconName): The :code:`podium` icon.\n PODIUM_OFF (IconName): The :code:`podium-off` icon.\n POINT (IconName): The :code:`point` icon.\n POINT_FILLED (IconName): The :code:`point-filled` icon.\n POINT_OFF (IconName): The :code:`point-off` icon.\n POINTER (IconName): The :code:`pointer` icon.\n POINTER_BOLT (IconName): The :code:`pointer-bolt` icon.\n POINTER_CANCEL (IconName): The :code:`pointer-cancel` icon.\n POINTER_CHECK (IconName): The :code:`pointer-check` icon.\n POINTER_CODE (IconName): The :code:`pointer-code` icon.\n POINTER_COG (IconName): The :code:`pointer-cog` icon.\n POINTER_DOLLAR (IconName): The :code:`pointer-dollar` icon.\n POINTER_DOWN (IconName): The :code:`pointer-down` icon.\n POINTER_EXCLAMATION (IconName): The :code:`pointer-exclamation` icon.\n POINTER_HEART (IconName): The :code:`pointer-heart` icon.\n POINTER_MINUS (IconName): The :code:`pointer-minus` icon.\n POINTER_OFF (IconName): The :code:`pointer-off` icon.\n POINTER_PAUSE (IconName): The :code:`pointer-pause` icon.\n POINTER_PIN (IconName): The :code:`pointer-pin` icon.\n POINTER_PLUS (IconName): The :code:`pointer-plus` icon.\n POINTER_QUESTION (IconName): The :code:`pointer-question` icon.\n POINTER_SEARCH (IconName): The :code:`pointer-search` icon.\n POINTER_SHARE (IconName): The :code:`pointer-share` icon.\n POINTER_STAR (IconName): The :code:`pointer-star` icon.\n POINTER_UP (IconName): The :code:`pointer-up` icon.\n POINTER_X (IconName): The :code:`pointer-x` icon.\n POKEBALL (IconName): The :code:`pokeball` icon.\n POKEBALL_OFF (IconName): The :code:`pokeball-off` icon.\n POKER_CHIP (IconName): The :code:`poker-chip` icon.\n POLAROID (IconName): The :code:`polaroid` icon.\n POLAROID_FILLED (IconName): The :code:`polaroid-filled` icon.\n POLYGON (IconName): The :code:`polygon` icon.\n POLYGON_OFF (IconName): The :code:`polygon-off` icon.\n POO (IconName): The :code:`poo` icon.\n POOL (IconName): The :code:`pool` icon.\n POOL_OFF (IconName): The :code:`pool-off` icon.\n POWER (IconName): The :code:`power` icon.\n PRAY (IconName): The :code:`pray` icon.\n PREMIUM_RIGHTS (IconName): The :code:`premium-rights` icon.\n PRESCRIPTION (IconName): The :code:`prescription` icon.\n PRESENTATION (IconName): The :code:`presentation` icon.\n PRESENTATION_ANALYTICS (IconName): The :code:`presentation-analytics` icon.\n PRESENTATION_OFF (IconName): The :code:`presentation-off` icon.\n PRINTER (IconName): The :code:`printer` icon.\n PRINTER_OFF (IconName): The :code:`printer-off` icon.\n PRISM (IconName): The :code:`prism` icon.\n PRISM_OFF (IconName): The :code:`prism-off` icon.\n PRISM_PLUS (IconName): The :code:`prism-plus` icon.\n PRISON (IconName): The :code:`prison` icon.\n PROGRESS (IconName): The :code:`progress` icon.\n PROGRESS_ALERT (IconName): The :code:`progress-alert` icon.\n PROGRESS_BOLT (IconName): The :code:`progress-bolt` icon.\n PROGRESS_CHECK (IconName): The :code:`progress-check` icon.\n PROGRESS_DOWN (IconName): The :code:`progress-down` icon.\n PROGRESS_HELP (IconName): The :code:`progress-help` icon.\n PROGRESS_X (IconName): The :code:`progress-x` icon.\n PROMPT (IconName): The :code:`prompt` icon.\n PROPELLER (IconName): The :code:`propeller` icon.\n PROPELLER_OFF (IconName): The :code:`propeller-off` icon.\n PUMPKIN_SCARY (IconName): The :code:`pumpkin-scary` icon.\n PUZZLE (IconName): The :code:`puzzle` icon.\n PUZZLE_2 (IconName): The :code:`puzzle-2` icon.\n PUZZLE_FILLED (IconName): The :code:`puzzle-filled` icon.\n PUZZLE_OFF (IconName): The :code:`puzzle-off` icon.\n PYRAMID (IconName): The :code:`pyramid` icon.\n PYRAMID_OFF (IconName): The :code:`pyramid-off` icon.\n PYRAMID_PLUS (IconName): The :code:`pyramid-plus` icon.\n QRCODE (IconName): The :code:`qrcode` icon.\n QRCODE_OFF (IconName): The :code:`qrcode-off` icon.\n QUESTION_MARK (IconName): The :code:`question-mark` icon.\n QUOTE (IconName): The :code:`quote` icon.\n QUOTE_OFF (IconName): The :code:`quote-off` icon.\n RADAR (IconName): The :code:`radar` icon.\n RADAR_2 (IconName): The :code:`radar-2` icon.\n RADAR_OFF (IconName): The :code:`radar-off` icon.\n RADIO (IconName): The :code:`radio` icon.\n RADIO_OFF (IconName): The :code:`radio-off` icon.\n RADIOACTIVE (IconName): The :code:`radioactive` icon.\n RADIOACTIVE_FILLED (IconName): The :code:`radioactive-filled` icon.\n RADIOACTIVE_OFF (IconName): The :code:`radioactive-off` icon.\n RADIUS_BOTTOM_LEFT (IconName): The :code:`radius-bottom-left` icon.\n RADIUS_BOTTOM_RIGHT (IconName): The :code:`radius-bottom-right` icon.\n RADIUS_TOP_LEFT (IconName): The :code:`radius-top-left` icon.\n RADIUS_TOP_RIGHT (IconName): The :code:`radius-top-right` icon.\n RAINBOW (IconName): The :code:`rainbow` icon.\n RAINBOW_OFF (IconName): The :code:`rainbow-off` icon.\n RATING_12_PLUS (IconName): The :code:`rating-12-plus` icon.\n RATING_14_PLUS (IconName): The :code:`rating-14-plus` icon.\n RATING_16_PLUS (IconName): The :code:`rating-16-plus` icon.\n RATING_18_PLUS (IconName): The :code:`rating-18-plus` icon.\n RATING_21_PLUS (IconName): The :code:`rating-21-plus` icon.\n RAZOR (IconName): The :code:`razor` icon.\n RAZOR_ELECTRIC (IconName): The :code:`razor-electric` icon.\n RECEIPT (IconName): The :code:`receipt` icon.\n RECEIPT_2 (IconName): The :code:`receipt-2` icon.\n RECEIPT_OFF (IconName): The :code:`receipt-off` icon.\n RECEIPT_REFUND (IconName): The :code:`receipt-refund` icon.\n RECEIPT_TAX (IconName): The :code:`receipt-tax` icon.\n RECHARGING (IconName): The :code:`recharging` icon.\n RECORD_MAIL (IconName): The :code:`record-mail` icon.\n RECORD_MAIL_OFF (IconName): The :code:`record-mail-off` icon.\n RECTANGLE (IconName): The :code:`rectangle` icon.\n RECTANGLE_FILLED (IconName): The :code:`rectangle-filled` icon.\n RECTANGLE_ROUNDED_BOTTOM (IconName): The :code:`rectangle-rounded-bottom` icon.\n RECTANGLE_ROUNDED_TOP (IconName): The :code:`rectangle-rounded-top` icon.\n RECTANGLE_VERTICAL (IconName): The :code:`rectangle-vertical` icon.\n RECTANGLE_VERTICAL_FILLED (IconName): The :code:`rectangle-vertical-filled` icon.\n RECTANGULAR_PRISM (IconName): The :code:`rectangular-prism` icon.\n RECTANGULAR_PRISM_OFF (IconName): The :code:`rectangular-prism-off` icon.\n RECTANGULAR_PRISM_PLUS (IconName): The :code:`rectangular-prism-plus` icon.\n RECYCLE (IconName): The :code:`recycle` icon.\n RECYCLE_OFF (IconName): The :code:`recycle-off` icon.\n REFRESH (IconName): The :code:`refresh` icon.\n REFRESH_ALERT (IconName): The :code:`refresh-alert` icon.\n REFRESH_DOT (IconName): The :code:`refresh-dot` icon.\n REFRESH_OFF (IconName): The :code:`refresh-off` icon.\n REGEX (IconName): The :code:`regex` icon.\n REGEX_OFF (IconName): The :code:`regex-off` icon.\n REGISTERED (IconName): The :code:`registered` icon.\n RELATION_MANY_TO_MANY (IconName): The :code:`relation-many-to-many` icon.\n RELATION_ONE_TO_MANY (IconName): The :code:`relation-one-to-many` icon.\n RELATION_ONE_TO_ONE (IconName): The :code:`relation-one-to-one` icon.\n RELOAD (IconName): The :code:`reload` icon.\n REPEAT (IconName): The :code:`repeat` icon.\n REPEAT_OFF (IconName): The :code:`repeat-off` icon.\n REPEAT_ONCE (IconName): The :code:`repeat-once` icon.\n REPLACE (IconName): The :code:`replace` icon.\n REPLACE_FILLED (IconName): The :code:`replace-filled` icon.\n REPLACE_OFF (IconName): The :code:`replace-off` icon.\n REPORT (IconName): The :code:`report` icon.\n REPORT_ANALYTICS (IconName): The :code:`report-analytics` icon.\n REPORT_MEDICAL (IconName): The :code:`report-medical` icon.\n REPORT_MONEY (IconName): The :code:`report-money` icon.\n REPORT_OFF (IconName): The :code:`report-off` icon.\n REPORT_SEARCH (IconName): The :code:`report-search` icon.\n RESERVED_LINE (IconName): The :code:`reserved-line` icon.\n RESIZE (IconName): The :code:`resize` icon.\n RESTORE (IconName): The :code:`restore` icon.\n REWIND_BACKWARD_10 (IconName): The :code:`rewind-backward-10` icon.\n REWIND_BACKWARD_15 (IconName): The :code:`rewind-backward-15` icon.\n REWIND_BACKWARD_20 (IconName): The :code:`rewind-backward-20` icon.\n REWIND_BACKWARD_30 (IconName): The :code:`rewind-backward-30` icon.\n REWIND_BACKWARD_40 (IconName): The :code:`rewind-backward-40` icon.\n REWIND_BACKWARD_5 (IconName): The :code:`rewind-backward-5` icon.\n REWIND_BACKWARD_50 (IconName): The :code:`rewind-backward-50` icon.\n REWIND_BACKWARD_60 (IconName): The :code:`rewind-backward-60` icon.\n REWIND_FORWARD_10 (IconName): The :code:`rewind-forward-10` icon.\n REWIND_FORWARD_15 (IconName): The :code:`rewind-forward-15` icon.\n REWIND_FORWARD_20 (IconName): The :code:`rewind-forward-20` icon.\n REWIND_FORWARD_30 (IconName): The :code:`rewind-forward-30` icon.\n REWIND_FORWARD_40 (IconName): The :code:`rewind-forward-40` icon.\n REWIND_FORWARD_5 (IconName): The :code:`rewind-forward-5` icon.\n REWIND_FORWARD_50 (IconName): The :code:`rewind-forward-50` icon.\n REWIND_FORWARD_60 (IconName): The :code:`rewind-forward-60` icon.\n RIBBON_HEALTH (IconName): The :code:`ribbon-health` icon.\n RINGS (IconName): The :code:`rings` icon.\n RIPPLE (IconName): The :code:`ripple` icon.\n RIPPLE_OFF (IconName): The :code:`ripple-off` icon.\n ROAD (IconName): The :code:`road` icon.\n ROAD_OFF (IconName): The :code:`road-off` icon.\n ROAD_SIGN (IconName): The :code:`road-sign` icon.\n ROBOT (IconName): The :code:`robot` icon.\n ROBOT_OFF (IconName): The :code:`robot-off` icon.\n ROCKET (IconName): The :code:`rocket` icon.\n ROCKET_OFF (IconName): The :code:`rocket-off` icon.\n ROLLER_SKATING (IconName): The :code:`roller-skating` icon.\n ROLLERCOASTER (IconName): The :code:`rollercoaster` icon.\n ROLLERCOASTER_OFF (IconName): The :code:`rollercoaster-off` icon.\n ROSETTE (IconName): The :code:`rosette` icon.\n ROSETTE_FILLED (IconName): The :code:`rosette-filled` icon.\n ROSETTE_NUMBER_0 (IconName): The :code:`rosette-number-0` icon.\n ROSETTE_NUMBER_1 (IconName): The :code:`rosette-number-1` icon.\n ROSETTE_NUMBER_2 (IconName): The :code:`rosette-number-2` icon.\n ROSETTE_NUMBER_3 (IconName): The :code:`rosette-number-3` icon.\n ROSETTE_NUMBER_4 (IconName): The :code:`rosette-number-4` icon.\n ROSETTE_NUMBER_5 (IconName): The :code:`rosette-number-5` icon.\n ROSETTE_NUMBER_6 (IconName): The :code:`rosette-number-6` icon.\n ROSETTE_NUMBER_7 (IconName): The :code:`rosette-number-7` icon.\n ROSETTE_NUMBER_8 (IconName): The :code:`rosette-number-8` icon.\n ROSETTE_NUMBER_9 (IconName): The :code:`rosette-number-9` icon.\n ROTATE (IconName): The :code:`rotate` icon.\n ROTATE_2 (IconName): The :code:`rotate-2` icon.\n ROTATE_360 (IconName): The :code:`rotate-360` icon.\n ROTATE_CLOCKWISE (IconName): The :code:`rotate-clockwise` icon.\n ROTATE_CLOCKWISE_2 (IconName): The :code:`rotate-clockwise-2` icon.\n ROTATE_DOT (IconName): The :code:`rotate-dot` icon.\n ROTATE_RECTANGLE (IconName): The :code:`rotate-rectangle` icon.\n ROUTE (IconName): The :code:`route` icon.\n ROUTE_2 (IconName): The :code:`route-2` icon.\n ROUTE_OFF (IconName): The :code:`route-off` icon.\n ROUTER (IconName): The :code:`router` icon.\n ROUTER_OFF (IconName): The :code:`router-off` icon.\n ROW_INSERT_BOTTOM (IconName): The :code:`row-insert-bottom` icon.\n ROW_INSERT_TOP (IconName): The :code:`row-insert-top` icon.\n ROW_REMOVE (IconName): The :code:`row-remove` icon.\n RSS (IconName): The :code:`rss` icon.\n RUBBER_STAMP (IconName): The :code:`rubber-stamp` icon.\n RUBBER_STAMP_OFF (IconName): The :code:`rubber-stamp-off` icon.\n RULER (IconName): The :code:`ruler` icon.\n RULER_2 (IconName): The :code:`ruler-2` icon.\n RULER_2_OFF (IconName): The :code:`ruler-2-off` icon.\n RULER_3 (IconName): The :code:`ruler-3` icon.\n RULER_MEASURE (IconName): The :code:`ruler-measure` icon.\n RULER_OFF (IconName): The :code:`ruler-off` icon.\n RUN (IconName): The :code:`run` icon.\n S_TURN_DOWN (IconName): The :code:`s-turn-down` icon.\n S_TURN_LEFT (IconName): The :code:`s-turn-left` icon.\n S_TURN_RIGHT (IconName): The :code:`s-turn-right` icon.\n S_TURN_UP (IconName): The :code:`s-turn-up` icon.\n SAILBOAT (IconName): The :code:`sailboat` icon.\n SAILBOAT_2 (IconName): The :code:`sailboat-2` icon.\n SAILBOAT_OFF (IconName): The :code:`sailboat-off` icon.\n SALAD (IconName): The :code:`salad` icon.\n SALT (IconName): The :code:`salt` icon.\n SATELLITE (IconName): The :code:`satellite` icon.\n SATELLITE_OFF (IconName): The :code:`satellite-off` icon.\n SAUSAGE (IconName): The :code:`sausage` icon.\n SCALE (IconName): The :code:`scale` icon.\n SCALE_OFF (IconName): The :code:`scale-off` icon.\n SCALE_OUTLINE (IconName): The :code:`scale-outline` icon.\n SCALE_OUTLINE_OFF (IconName): The :code:`scale-outline-off` icon.\n SCAN (IconName): The :code:`scan` icon.\n SCAN_EYE (IconName): The :code:`scan-eye` icon.\n SCHEMA (IconName): The :code:`schema` icon.\n SCHEMA_OFF (IconName): The :code:`schema-off` icon.\n SCHOOL (IconName): The :code:`school` icon.\n SCHOOL_BELL (IconName): The :code:`school-bell` icon.\n SCHOOL_OFF (IconName): The :code:`school-off` icon.\n SCISSORS (IconName): The :code:`scissors` icon.\n SCISSORS_OFF (IconName): The :code:`scissors-off` icon.\n SCOOTER (IconName): The :code:`scooter` icon.\n SCOOTER_ELECTRIC (IconName): The :code:`scooter-electric` icon.\n SCOREBOARD (IconName): The :code:`scoreboard` icon.\n SCREEN_SHARE (IconName): The :code:`screen-share` icon.\n SCREEN_SHARE_OFF (IconName): The :code:`screen-share-off` icon.\n SCREENSHOT (IconName): The :code:`screenshot` icon.\n SCRIBBLE (IconName): The :code:`scribble` icon.\n SCRIBBLE_OFF (IconName): The :code:`scribble-off` icon.\n SCRIPT (IconName): The :code:`script` icon.\n SCRIPT_MINUS (IconName): The :code:`script-minus` icon.\n SCRIPT_PLUS (IconName): The :code:`script-plus` icon.\n SCRIPT_X (IconName): The :code:`script-x` icon.\n SCUBA_MASK (IconName): The :code:`scuba-mask` icon.\n SCUBA_MASK_OFF (IconName): The :code:`scuba-mask-off` icon.\n SDK (IconName): The :code:`sdk` icon.\n SEARCH (IconName): The :code:`search` icon.\n SEARCH_OFF (IconName): The :code:`search-off` icon.\n SECTION (IconName): The :code:`section` icon.\n SECTION_SIGN (IconName): The :code:`section-sign` icon.\n SEEDING (IconName): The :code:`seeding` icon.\n SEEDING_OFF (IconName): The :code:`seeding-off` icon.\n SELECT (IconName): The :code:`select` icon.\n SELECT_ALL (IconName): The :code:`select-all` icon.\n SELECTOR (IconName): The :code:`selector` icon.\n SEND (IconName): The :code:`send` icon.\n SEND_OFF (IconName): The :code:`send-off` icon.\n SEO (IconName): The :code:`seo` icon.\n SEPARATOR (IconName): The :code:`separator` icon.\n SEPARATOR_HORIZONTAL (IconName): The :code:`separator-horizontal` icon.\n SEPARATOR_VERTICAL (IconName): The :code:`separator-vertical` icon.\n SERVER (IconName): The :code:`server` icon.\n SERVER_2 (IconName): The :code:`server-2` icon.\n SERVER_BOLT (IconName): The :code:`server-bolt` icon.\n SERVER_COG (IconName): The :code:`server-cog` icon.\n SERVER_OFF (IconName): The :code:`server-off` icon.\n SERVICEMARK (IconName): The :code:`servicemark` icon.\n SETTINGS (IconName): The :code:`settings` icon.\n SETTINGS_2 (IconName): The :code:`settings-2` icon.\n SETTINGS_AUTOMATION (IconName): The :code:`settings-automation` icon.\n SETTINGS_BOLT (IconName): The :code:`settings-bolt` icon.\n SETTINGS_CANCEL (IconName): The :code:`settings-cancel` icon.\n SETTINGS_CHECK (IconName): The :code:`settings-check` icon.\n SETTINGS_CODE (IconName): The :code:`settings-code` icon.\n SETTINGS_COG (IconName): The :code:`settings-cog` icon.\n SETTINGS_DOLLAR (IconName): The :code:`settings-dollar` icon.\n SETTINGS_DOWN (IconName): The :code:`settings-down` icon.\n SETTINGS_EXCLAMATION (IconName): The :code:`settings-exclamation` icon.\n SETTINGS_FILLED (IconName): The :code:`settings-filled` icon.\n SETTINGS_HEART (IconName): The :code:`settings-heart` icon.\n SETTINGS_MINUS (IconName): The :code:`settings-minus` icon.\n SETTINGS_OFF (IconName): The :code:`settings-off` icon.\n SETTINGS_PAUSE (IconName): The :code:`settings-pause` icon.\n SETTINGS_PIN (IconName): The :code:`settings-pin` icon.\n SETTINGS_PLUS (IconName): The :code:`settings-plus` icon.\n SETTINGS_QUESTION (IconName): The :code:`settings-question` icon.\n SETTINGS_SEARCH (IconName): The :code:`settings-search` icon.\n SETTINGS_SHARE (IconName): The :code:`settings-share` icon.\n SETTINGS_STAR (IconName): The :code:`settings-star` icon.\n SETTINGS_UP (IconName): The :code:`settings-up` icon.\n SETTINGS_X (IconName): The :code:`settings-x` icon.\n SHADOW (IconName): The :code:`shadow` icon.\n SHADOW_OFF (IconName): The :code:`shadow-off` icon.\n SHAPE (IconName): The :code:`shape` icon.\n SHAPE_2 (IconName): The :code:`shape-2` icon.\n SHAPE_3 (IconName): The :code:`shape-3` icon.\n SHAPE_OFF (IconName): The :code:`shape-off` icon.\n SHARE (IconName): The :code:`share` icon.\n SHARE_2 (IconName): The :code:`share-2` icon.\n SHARE_3 (IconName): The :code:`share-3` icon.\n SHARE_OFF (IconName): The :code:`share-off` icon.\n SHI_JUMPING (IconName): The :code:`shi-jumping` icon.\n SHIELD (IconName): The :code:`shield` icon.\n SHIELD_BOLT (IconName): The :code:`shield-bolt` icon.\n SHIELD_CANCEL (IconName): The :code:`shield-cancel` icon.\n SHIELD_CHECK (IconName): The :code:`shield-check` icon.\n SHIELD_CHECK_FILLED (IconName): The :code:`shield-check-filled` icon.\n SHIELD_CHECKERED (IconName): The :code:`shield-checkered` icon.\n SHIELD_CHECKERED_FILLED (IconName): The :code:`shield-checkered-filled` icon.\n SHIELD_CHEVRON (IconName): The :code:`shield-chevron` icon.\n SHIELD_CODE (IconName): The :code:`shield-code` icon.\n SHIELD_COG (IconName): The :code:`shield-cog` icon.\n SHIELD_DOLLAR (IconName): The :code:`shield-dollar` icon.\n SHIELD_DOWN (IconName): The :code:`shield-down` icon.\n SHIELD_EXCLAMATION (IconName): The :code:`shield-exclamation` icon.\n SHIELD_FILLED (IconName): The :code:`shield-filled` icon.\n SHIELD_HALF (IconName): The :code:`shield-half` icon.\n SHIELD_HALF_FILLED (IconName): The :code:`shield-half-filled` icon.\n SHIELD_HEART (IconName): The :code:`shield-heart` icon.\n SHIELD_LOCK (IconName): The :code:`shield-lock` icon.\n SHIELD_LOCK_FILLED (IconName): The :code:`shield-lock-filled` icon.\n SHIELD_MINUS (IconName): The :code:`shield-minus` icon.\n SHIELD_OFF (IconName): The :code:`shield-off` icon.\n SHIELD_PAUSE (IconName): The :code:`shield-pause` icon.\n SHIELD_PIN (IconName): The :code:`shield-pin` icon.\n SHIELD_PLUS (IconName): The :code:`shield-plus` icon.\n SHIELD_QUESTION (IconName): The :code:`shield-question` icon.\n SHIELD_SEARCH (IconName): The :code:`shield-search` icon.\n SHIELD_SHARE (IconName): The :code:`shield-share` icon.\n SHIELD_STAR (IconName): The :code:`shield-star` icon.\n SHIELD_UP (IconName): The :code:`shield-up` icon.\n SHIELD_X (IconName): The :code:`shield-x` icon.\n SHIP (IconName): The :code:`ship` icon.\n SHIP_OFF (IconName): The :code:`ship-off` icon.\n SHIRT (IconName): The :code:`shirt` icon.\n SHIRT_FILLED (IconName): The :code:`shirt-filled` icon.\n SHIRT_OFF (IconName): The :code:`shirt-off` icon.\n SHIRT_SPORT (IconName): The :code:`shirt-sport` icon.\n SHOE (IconName): The :code:`shoe` icon.\n SHOE_OFF (IconName): The :code:`shoe-off` icon.\n SHOPPING_BAG (IconName): The :code:`shopping-bag` icon.\n SHOPPING_CART (IconName): The :code:`shopping-cart` icon.\n SHOPPING_CART_DISCOUNT (IconName): The :code:`shopping-cart-discount` icon.\n SHOPPING_CART_OFF (IconName): The :code:`shopping-cart-off` icon.\n SHOPPING_CART_PLUS (IconName): The :code:`shopping-cart-plus` icon.\n SHOPPING_CART_X (IconName): The :code:`shopping-cart-x` icon.\n SHOVEL (IconName): The :code:`shovel` icon.\n SHREDDER (IconName): The :code:`shredder` icon.\n SIGN_LEFT (IconName): The :code:`sign-left` icon.\n SIGN_LEFT_FILLED (IconName): The :code:`sign-left-filled` icon.\n SIGN_RIGHT (IconName): The :code:`sign-right` icon.\n SIGN_RIGHT_FILLED (IconName): The :code:`sign-right-filled` icon.\n SIGNAL_2G (IconName): The :code:`signal-2g` icon.\n SIGNAL_3G (IconName): The :code:`signal-3g` icon.\n SIGNAL_4G (IconName): The :code:`signal-4g` icon.\n SIGNAL_4G_PLUS (IconName): The :code:`signal-4g-plus` icon.\n SIGNAL_5G (IconName): The :code:`signal-5g` icon.\n SIGNAL_6G (IconName): The :code:`signal-6g` icon.\n SIGNAL_E (IconName): The :code:`signal-e` icon.\n SIGNAL_G (IconName): The :code:`signal-g` icon.\n SIGNAL_H (IconName): The :code:`signal-h` icon.\n SIGNAL_H_PLUS (IconName): The :code:`signal-h-plus` icon.\n SIGNAL_LTE (IconName): The :code:`signal-lte` icon.\n SIGNATURE (IconName): The :code:`signature` icon.\n SIGNATURE_OFF (IconName): The :code:`signature-off` icon.\n SITEMAP (IconName): The :code:`sitemap` icon.\n SITEMAP_OFF (IconName): The :code:`sitemap-off` icon.\n SKATEBOARD (IconName): The :code:`skateboard` icon.\n SKATEBOARD_OFF (IconName): The :code:`skateboard-off` icon.\n SKATEBOARDING (IconName): The :code:`skateboarding` icon.\n SKULL (IconName): The :code:`skull` icon.\n SLASH (IconName): The :code:`slash` icon.\n SLASHES (IconName): The :code:`slashes` icon.\n SLEIGH (IconName): The :code:`sleigh` icon.\n SLICE (IconName): The :code:`slice` icon.\n SLIDESHOW (IconName): The :code:`slideshow` icon.\n SMART_HOME (IconName): The :code:`smart-home` icon.\n SMART_HOME_OFF (IconName): The :code:`smart-home-off` icon.\n SMOKING (IconName): The :code:`smoking` icon.\n SMOKING_NO (IconName): The :code:`smoking-no` icon.\n SNOWFLAKE (IconName): The :code:`snowflake` icon.\n SNOWFLAKE_OFF (IconName): The :code:`snowflake-off` icon.\n SNOWMAN (IconName): The :code:`snowman` icon.\n SOCCER_FIELD (IconName): The :code:`soccer-field` icon.\n SOCIAL (IconName): The :code:`social` icon.\n SOCIAL_OFF (IconName): The :code:`social-off` icon.\n SOCK (IconName): The :code:`sock` icon.\n SOFA (IconName): The :code:`sofa` icon.\n SOFA_OFF (IconName): The :code:`sofa-off` icon.\n SOLAR_PANEL (IconName): The :code:`solar-panel` icon.\n SOLAR_PANEL_2 (IconName): The :code:`solar-panel-2` icon.\n SORT_0_9 (IconName): The :code:`sort-0-9` icon.\n SORT_9_0 (IconName): The :code:`sort-9-0` icon.\n SORT_A_Z (IconName): The :code:`sort-a-z` icon.\n SORT_ASCENDING (IconName): The :code:`sort-ascending` icon.\n SORT_ASCENDING_2 (IconName): The :code:`sort-ascending-2` icon.\n SORT_ASCENDING_LETTERS (IconName): The :code:`sort-ascending-letters` icon.\n SORT_ASCENDING_NUMBERS (IconName): The :code:`sort-ascending-numbers` icon.\n SORT_DESCENDING (IconName): The :code:`sort-descending` icon.\n SORT_DESCENDING_2 (IconName): The :code:`sort-descending-2` icon.\n SORT_DESCENDING_LETTERS (IconName): The :code:`sort-descending-letters` icon.\n SORT_DESCENDING_NUMBERS (IconName): The :code:`sort-descending-numbers` icon.\n SORT_Z_A (IconName): The :code:`sort-z-a` icon.\n SOS (IconName): The :code:`sos` icon.\n SOUP (IconName): The :code:`soup` icon.\n SOUP_OFF (IconName): The :code:`soup-off` icon.\n SOURCE_CODE (IconName): The :code:`source-code` icon.\n SPACE (IconName): The :code:`space` icon.\n SPACE_OFF (IconName): The :code:`space-off` icon.\n SPACING_HORIZONTAL (IconName): The :code:`spacing-horizontal` icon.\n SPACING_VERTICAL (IconName): The :code:`spacing-vertical` icon.\n SPADE (IconName): The :code:`spade` icon.\n SPADE_FILLED (IconName): The :code:`spade-filled` icon.\n SPARKLES (IconName): The :code:`sparkles` icon.\n SPEAKERPHONE (IconName): The :code:`speakerphone` icon.\n SPEEDBOAT (IconName): The :code:`speedboat` icon.\n SPHERE (IconName): The :code:`sphere` icon.\n SPHERE_OFF (IconName): The :code:`sphere-off` icon.\n SPHERE_PLUS (IconName): The :code:`sphere-plus` icon.\n SPIDER (IconName): The :code:`spider` icon.\n SPIRAL (IconName): The :code:`spiral` icon.\n SPIRAL_OFF (IconName): The :code:`spiral-off` icon.\n SPORT_BILLARD (IconName): The :code:`sport-billard` icon.\n SPRAY (IconName): The :code:`spray` icon.\n SPY (IconName): The :code:`spy` icon.\n SPY_OFF (IconName): The :code:`spy-off` icon.\n SQL (IconName): The :code:`sql` icon.\n SQUARE (IconName): The :code:`square` icon.\n SQUARE_0_FILLED (IconName): The :code:`square-0-filled` icon.\n SQUARE_1_FILLED (IconName): The :code:`square-1-filled` icon.\n SQUARE_2_FILLED (IconName): The :code:`square-2-filled` icon.\n SQUARE_3_FILLED (IconName): The :code:`square-3-filled` icon.\n SQUARE_4_FILLED (IconName): The :code:`square-4-filled` icon.\n SQUARE_5_FILLED (IconName): The :code:`square-5-filled` icon.\n SQUARE_6_FILLED (IconName): The :code:`square-6-filled` icon.\n SQUARE_7_FILLED (IconName): The :code:`square-7-filled` icon.\n SQUARE_8_FILLED (IconName): The :code:`square-8-filled` icon.\n SQUARE_9_FILLED (IconName): The :code:`square-9-filled` icon.\n SQUARE_ARROW_DOWN (IconName): The :code:`square-arrow-down` icon.\n SQUARE_ARROW_LEFT (IconName): The :code:`square-arrow-left` icon.\n SQUARE_ARROW_RIGHT (IconName): The :code:`square-arrow-right` icon.\n SQUARE_ARROW_UP (IconName): The :code:`square-arrow-up` icon.\n SQUARE_ASTERISK (IconName): The :code:`square-asterisk` icon.\n SQUARE_CHECK (IconName): The :code:`square-check` icon.\n SQUARE_CHECK_FILLED (IconName): The :code:`square-check-filled` icon.\n SQUARE_CHEVRON_DOWN (IconName): The :code:`square-chevron-down` icon.\n SQUARE_CHEVRON_LEFT (IconName): The :code:`square-chevron-left` icon.\n SQUARE_CHEVRON_RIGHT (IconName): The :code:`square-chevron-right` icon.\n SQUARE_CHEVRON_UP (IconName): The :code:`square-chevron-up` icon.\n SQUARE_CHEVRONS_DOWN (IconName): The :code:`square-chevrons-down` icon.\n SQUARE_CHEVRONS_LEFT (IconName): The :code:`square-chevrons-left` icon.\n SQUARE_CHEVRONS_RIGHT (IconName): The :code:`square-chevrons-right` icon.\n SQUARE_CHEVRONS_UP (IconName): The :code:`square-chevrons-up` icon.\n SQUARE_DOT (IconName): The :code:`square-dot` icon.\n SQUARE_F0 (IconName): The :code:`square-f0` icon.\n SQUARE_F0_FILLED (IconName): The :code:`square-f0-filled` icon.\n SQUARE_F1 (IconName): The :code:`square-f1` icon.\n SQUARE_F1_FILLED (IconName): The :code:`square-f1-filled` icon.\n SQUARE_F2 (IconName): The :code:`square-f2` icon.\n SQUARE_F2_FILLED (IconName): The :code:`square-f2-filled` icon.\n SQUARE_F3 (IconName): The :code:`square-f3` icon.\n SQUARE_F3_FILLED (IconName): The :code:`square-f3-filled` icon.\n SQUARE_F4 (IconName): The :code:`square-f4` icon.\n SQUARE_F4_FILLED (IconName): The :code:`square-f4-filled` icon.\n SQUARE_F5 (IconName): The :code:`square-f5` icon.\n SQUARE_F5_FILLED (IconName): The :code:`square-f5-filled` icon.\n SQUARE_F6 (IconName): The :code:`square-f6` icon.\n SQUARE_F6_FILLED (IconName): The :code:`square-f6-filled` icon.\n SQUARE_F7 (IconName): The :code:`square-f7` icon.\n SQUARE_F7_FILLED (IconName): The :code:`square-f7-filled` icon.\n SQUARE_F8 (IconName): The :code:`square-f8` icon.\n SQUARE_F8_FILLED (IconName): The :code:`square-f8-filled` icon.\n SQUARE_F9 (IconName): The :code:`square-f9` icon.\n SQUARE_F9_FILLED (IconName): The :code:`square-f9-filled` icon.\n SQUARE_FORBID (IconName): The :code:`square-forbid` icon.\n SQUARE_FORBID_2 (IconName): The :code:`square-forbid-2` icon.\n SQUARE_HALF (IconName): The :code:`square-half` icon.\n SQUARE_KEY (IconName): The :code:`square-key` icon.\n SQUARE_LETTER_A (IconName): The :code:`square-letter-a` icon.\n SQUARE_LETTER_B (IconName): The :code:`square-letter-b` icon.\n SQUARE_LETTER_C (IconName): The :code:`square-letter-c` icon.\n SQUARE_LETTER_D (IconName): The :code:`square-letter-d` icon.\n SQUARE_LETTER_E (IconName): The :code:`square-letter-e` icon.\n SQUARE_LETTER_F (IconName): The :code:`square-letter-f` icon.\n SQUARE_LETTER_G (IconName): The :code:`square-letter-g` icon.\n SQUARE_LETTER_H (IconName): The :code:`square-letter-h` icon.\n SQUARE_LETTER_I (IconName): The :code:`square-letter-i` icon.\n SQUARE_LETTER_J (IconName): The :code:`square-letter-j` icon.\n SQUARE_LETTER_K (IconName): The :code:`square-letter-k` icon.\n SQUARE_LETTER_L (IconName): The :code:`square-letter-l` icon.\n SQUARE_LETTER_M (IconName): The :code:`square-letter-m` icon.\n SQUARE_LETTER_N (IconName): The :code:`square-letter-n` icon.\n SQUARE_LETTER_O (IconName): The :code:`square-letter-o` icon.\n SQUARE_LETTER_P (IconName): The :code:`square-letter-p` icon.\n SQUARE_LETTER_Q (IconName): The :code:`square-letter-q` icon.\n SQUARE_LETTER_R (IconName): The :code:`square-letter-r` icon.\n SQUARE_LETTER_S (IconName): The :code:`square-letter-s` icon.\n SQUARE_LETTER_T (IconName): The :code:`square-letter-t` icon.\n SQUARE_LETTER_U (IconName): The :code:`square-letter-u` icon.\n SQUARE_LETTER_V (IconName): The :code:`square-letter-v` icon.\n SQUARE_LETTER_W (IconName): The :code:`square-letter-w` icon.\n SQUARE_LETTER_X (IconName): The :code:`square-letter-x` icon.\n SQUARE_LETTER_Y (IconName): The :code:`square-letter-y` icon.\n SQUARE_LETTER_Z (IconName): The :code:`square-letter-z` icon.\n SQUARE_MINUS (IconName): The :code:`square-minus` icon.\n SQUARE_NUMBER_0 (IconName): The :code:`square-number-0` icon.\n SQUARE_NUMBER_1 (IconName): The :code:`square-number-1` icon.\n SQUARE_NUMBER_2 (IconName): The :code:`square-number-2` icon.\n SQUARE_NUMBER_3 (IconName): The :code:`square-number-3` icon.\n SQUARE_NUMBER_4 (IconName): The :code:`square-number-4` icon.\n SQUARE_NUMBER_5 (IconName): The :code:`square-number-5` icon.\n SQUARE_NUMBER_6 (IconName): The :code:`square-number-6` icon.\n SQUARE_NUMBER_7 (IconName): The :code:`square-number-7` icon.\n SQUARE_NUMBER_8 (IconName): The :code:`square-number-8` icon.\n SQUARE_NUMBER_9 (IconName): The :code:`square-number-9` icon.\n SQUARE_OFF (IconName): The :code:`square-off` icon.\n SQUARE_PLUS (IconName): The :code:`square-plus` icon.\n SQUARE_ROOT (IconName): The :code:`square-root` icon.\n SQUARE_ROOT_2 (IconName): The :code:`square-root-2` icon.\n SQUARE_ROTATED (IconName): The :code:`square-rotated` icon.\n SQUARE_ROTATED_FILLED (IconName): The :code:`square-rotated-filled` icon.\n SQUARE_ROTATED_FORBID (IconName): The :code:`square-rotated-forbid` icon.\n SQUARE_ROTATED_FORBID_2 (IconName): The :code:`square-rotated-forbid-2` icon.\n SQUARE_ROTATED_OFF (IconName): The :code:`square-rotated-off` icon.\n SQUARE_ROUNDED (IconName): The :code:`square-rounded` icon.\n SQUARE_ROUNDED_ARROW_DOWN (IconName): The :code:`square-rounded-arrow-down` icon.\n SQUARE_ROUNDED_ARROW_DOWN_FILLED (IconName): The :code:`square-rounded-arrow-down-filled` icon.\n SQUARE_ROUNDED_ARROW_LEFT (IconName): The :code:`square-rounded-arrow-left` icon.\n SQUARE_ROUNDED_ARROW_LEFT_FILLED (IconName): The :code:`square-rounded-arrow-left-filled` icon.\n SQUARE_ROUNDED_ARROW_RIGHT (IconName): The :code:`square-rounded-arrow-right` icon.\n SQUARE_ROUNDED_ARROW_RIGHT_FILLED (IconName): The :code:`square-rounded-arrow-right-filled` icon.\n SQUARE_ROUNDED_ARROW_UP (IconName): The :code:`square-rounded-arrow-up` icon.\n SQUARE_ROUNDED_ARROW_UP_FILLED (IconName): The :code:`square-rounded-arrow-up-filled` icon.\n SQUARE_ROUNDED_CHECK (IconName): The :code:`square-rounded-check` icon.\n SQUARE_ROUNDED_CHECK_FILLED (IconName): The :code:`square-rounded-check-filled` icon.\n SQUARE_ROUNDED_CHEVRON_DOWN (IconName): The :code:`square-rounded-chevron-down` icon.\n SQUARE_ROUNDED_CHEVRON_DOWN_FILLED (IconName): The :code:`square-rounded-chevron-down-filled` icon.\n SQUARE_ROUNDED_CHEVRON_LEFT (IconName): The :code:`square-rounded-chevron-left` icon.\n SQUARE_ROUNDED_CHEVRON_LEFT_FILLED (IconName): The :code:`square-rounded-chevron-left-filled` icon.\n SQUARE_ROUNDED_CHEVRON_RIGHT (IconName): The :code:`square-rounded-chevron-right` icon.\n SQUARE_ROUNDED_CHEVRON_RIGHT_FILLED (IconName): The :code:`square-rounded-chevron-right-filled` icon.\n SQUARE_ROUNDED_CHEVRON_UP (IconName): The :code:`square-rounded-chevron-up` icon.\n SQUARE_ROUNDED_CHEVRON_UP_FILLED (IconName): The :code:`square-rounded-chevron-up-filled` icon.\n SQUARE_ROUNDED_CHEVRONS_DOWN (IconName): The :code:`square-rounded-chevrons-down` icon.\n SQUARE_ROUNDED_CHEVRONS_DOWN_FILLED (IconName): The :code:`square-rounded-chevrons-down-filled` icon.\n SQUARE_ROUNDED_CHEVRONS_LEFT (IconName): The :code:`square-rounded-chevrons-left` icon.\n SQUARE_ROUNDED_CHEVRONS_LEFT_FILLED (IconName): The :code:`square-rounded-chevrons-left-filled` icon.\n SQUARE_ROUNDED_CHEVRONS_RIGHT (IconName): The :code:`square-rounded-chevrons-right` icon.\n SQUARE_ROUNDED_CHEVRONS_RIGHT_FILLED (IconName): The :code:`square-rounded-chevrons-right-filled` icon.\n SQUARE_ROUNDED_CHEVRONS_UP (IconName): The :code:`square-rounded-chevrons-up` icon.\n SQUARE_ROUNDED_CHEVRONS_UP_FILLED (IconName): The :code:`square-rounded-chevrons-up-filled` icon.\n SQUARE_ROUNDED_FILLED (IconName): The :code:`square-rounded-filled` icon.\n SQUARE_ROUNDED_LETTER_A (IconName): The :code:`square-rounded-letter-a` icon.\n SQUARE_ROUNDED_LETTER_B (IconName): The :code:`square-rounded-letter-b` icon.\n SQUARE_ROUNDED_LETTER_C (IconName): The :code:`square-rounded-letter-c` icon.\n SQUARE_ROUNDED_LETTER_D (IconName): The :code:`square-rounded-letter-d` icon.\n SQUARE_ROUNDED_LETTER_E (IconName): The :code:`square-rounded-letter-e` icon.\n SQUARE_ROUNDED_LETTER_F (IconName): The :code:`square-rounded-letter-f` icon.\n SQUARE_ROUNDED_LETTER_G (IconName): The :code:`square-rounded-letter-g` icon.\n SQUARE_ROUNDED_LETTER_H (IconName): The :code:`square-rounded-letter-h` icon.\n SQUARE_ROUNDED_LETTER_I (IconName): The :code:`square-rounded-letter-i` icon.\n SQUARE_ROUNDED_LETTER_J (IconName): The :code:`square-rounded-letter-j` icon.\n SQUARE_ROUNDED_LETTER_K (IconName): The :code:`square-rounded-letter-k` icon.\n SQUARE_ROUNDED_LETTER_L (IconName): The :code:`square-rounded-letter-l` icon.\n SQUARE_ROUNDED_LETTER_M (IconName): The :code:`square-rounded-letter-m` icon.\n SQUARE_ROUNDED_LETTER_N (IconName): The :code:`square-rounded-letter-n` icon.\n SQUARE_ROUNDED_LETTER_O (IconName): The :code:`square-rounded-letter-o` icon.\n SQUARE_ROUNDED_LETTER_P (IconName): The :code:`square-rounded-letter-p` icon.\n SQUARE_ROUNDED_LETTER_Q (IconName): The :code:`square-rounded-letter-q` icon.\n SQUARE_ROUNDED_LETTER_R (IconName): The :code:`square-rounded-letter-r` icon.\n SQUARE_ROUNDED_LETTER_S (IconName): The :code:`square-rounded-letter-s` icon.\n SQUARE_ROUNDED_LETTER_T (IconName): The :code:`square-rounded-letter-t` icon.\n SQUARE_ROUNDED_LETTER_U (IconName): The :code:`square-rounded-letter-u` icon.\n SQUARE_ROUNDED_LETTER_V (IconName): The :code:`square-rounded-letter-v` icon.\n SQUARE_ROUNDED_LETTER_W (IconName): The :code:`square-rounded-letter-w` icon.\n SQUARE_ROUNDED_LETTER_X (IconName): The :code:`square-rounded-letter-x` icon.\n SQUARE_ROUNDED_LETTER_Y (IconName): The :code:`square-rounded-letter-y` icon.\n SQUARE_ROUNDED_LETTER_Z (IconName): The :code:`square-rounded-letter-z` icon.\n SQUARE_ROUNDED_MINUS (IconName): The :code:`square-rounded-minus` icon.\n SQUARE_ROUNDED_NUMBER_0 (IconName): The :code:`square-rounded-number-0` icon.\n SQUARE_ROUNDED_NUMBER_0_FILLED (IconName): The :code:`square-rounded-number-0-filled` icon.\n SQUARE_ROUNDED_NUMBER_1 (IconName): The :code:`square-rounded-number-1` icon.\n SQUARE_ROUNDED_NUMBER_1_FILLED (IconName): The :code:`square-rounded-number-1-filled` icon.\n SQUARE_ROUNDED_NUMBER_2 (IconName): The :code:`square-rounded-number-2` icon.\n SQUARE_ROUNDED_NUMBER_2_FILLED (IconName): The :code:`square-rounded-number-2-filled` icon.\n SQUARE_ROUNDED_NUMBER_3 (IconName): The :code:`square-rounded-number-3` icon.\n SQUARE_ROUNDED_NUMBER_3_FILLED (IconName): The :code:`square-rounded-number-3-filled` icon.\n SQUARE_ROUNDED_NUMBER_4 (IconName): The :code:`square-rounded-number-4` icon.\n SQUARE_ROUNDED_NUMBER_4_FILLED (IconName): The :code:`square-rounded-number-4-filled` icon.\n SQUARE_ROUNDED_NUMBER_5 (IconName): The :code:`square-rounded-number-5` icon.\n SQUARE_ROUNDED_NUMBER_5_FILLED (IconName): The :code:`square-rounded-number-5-filled` icon.\n SQUARE_ROUNDED_NUMBER_6 (IconName): The :code:`square-rounded-number-6` icon.\n SQUARE_ROUNDED_NUMBER_6_FILLED (IconName): The :code:`square-rounded-number-6-filled` icon.\n SQUARE_ROUNDED_NUMBER_7 (IconName): The :code:`square-rounded-number-7` icon.\n SQUARE_ROUNDED_NUMBER_7_FILLED (IconName): The :code:`square-rounded-number-7-filled` icon.\n SQUARE_ROUNDED_NUMBER_8 (IconName): The :code:`square-rounded-number-8` icon.\n SQUARE_ROUNDED_NUMBER_8_FILLED (IconName): The :code:`square-rounded-number-8-filled` icon.\n SQUARE_ROUNDED_NUMBER_9 (IconName): The :code:`square-rounded-number-9` icon.\n SQUARE_ROUNDED_NUMBER_9_FILLED (IconName): The :code:`square-rounded-number-9-filled` icon.\n SQUARE_ROUNDED_PLUS (IconName): The :code:`square-rounded-plus` icon.\n SQUARE_ROUNDED_PLUS_FILLED (IconName): The :code:`square-rounded-plus-filled` icon.\n SQUARE_ROUNDED_X (IconName): The :code:`square-rounded-x` icon.\n SQUARE_ROUNDED_X_FILLED (IconName): The :code:`square-rounded-x-filled` icon.\n SQUARE_TOGGLE (IconName): The :code:`square-toggle` icon.\n SQUARE_TOGGLE_HORIZONTAL (IconName): The :code:`square-toggle-horizontal` icon.\n SQUARE_X (IconName): The :code:`square-x` icon.\n SQUARES_DIAGONAL (IconName): The :code:`squares-diagonal` icon.\n SQUARES_FILLED (IconName): The :code:`squares-filled` icon.\n STACK (IconName): The :code:`stack` icon.\n STACK_2 (IconName): The :code:`stack-2` icon.\n STACK_3 (IconName): The :code:`stack-3` icon.\n STACK_POP (IconName): The :code:`stack-pop` icon.\n STACK_PUSH (IconName): The :code:`stack-push` icon.\n STAIRS (IconName): The :code:`stairs` icon.\n STAIRS_DOWN (IconName): The :code:`stairs-down` icon.\n STAIRS_UP (IconName): The :code:`stairs-up` icon.\n STAR (IconName): The :code:`star` icon.\n STAR_FILLED (IconName): The :code:`star-filled` icon.\n STAR_HALF (IconName): The :code:`star-half` icon.\n STAR_HALF_FILLED (IconName): The :code:`star-half-filled` icon.\n STAR_OFF (IconName): The :code:`star-off` icon.\n STARS (IconName): The :code:`stars` icon.\n STARS_FILLED (IconName): The :code:`stars-filled` icon.\n STARS_OFF (IconName): The :code:`stars-off` icon.\n STATUS_CHANGE (IconName): The :code:`status-change` icon.\n STEAM (IconName): The :code:`steam` icon.\n STEERING_WHEEL (IconName): The :code:`steering-wheel` icon.\n STEERING_WHEEL_OFF (IconName): The :code:`steering-wheel-off` icon.\n STEP_INTO (IconName): The :code:`step-into` icon.\n STEP_OUT (IconName): The :code:`step-out` icon.\n STEREO_GLASSES (IconName): The :code:`stereo-glasses` icon.\n STETHOSCOPE (IconName): The :code:`stethoscope` icon.\n STETHOSCOPE_OFF (IconName): The :code:`stethoscope-off` icon.\n STICKER (IconName): The :code:`sticker` icon.\n STORM (IconName): The :code:`storm` icon.\n STORM_OFF (IconName): The :code:`storm-off` icon.\n STRETCHING (IconName): The :code:`stretching` icon.\n STRETCHING_2 (IconName): The :code:`stretching-2` icon.\n STRIKETHROUGH (IconName): The :code:`strikethrough` icon.\n SUBMARINE (IconName): The :code:`submarine` icon.\n SUBSCRIPT (IconName): The :code:`subscript` icon.\n SUBTASK (IconName): The :code:`subtask` icon.\n SUM (IconName): The :code:`sum` icon.\n SUM_OFF (IconName): The :code:`sum-off` icon.\n SUN (IconName): The :code:`sun` icon.\n SUN_FILLED (IconName): The :code:`sun-filled` icon.\n SUN_HIGH (IconName): The :code:`sun-high` icon.\n SUN_LOW (IconName): The :code:`sun-low` icon.\n SUN_MOON (IconName): The :code:`sun-moon` icon.\n SUN_OFF (IconName): The :code:`sun-off` icon.\n SUN_WIND (IconName): The :code:`sun-wind` icon.\n SUNGLASSES (IconName): The :code:`sunglasses` icon.\n SUNRISE (IconName): The :code:`sunrise` icon.\n SUNSET (IconName): The :code:`sunset` icon.\n SUNSET_2 (IconName): The :code:`sunset-2` icon.\n SUPERSCRIPT (IconName): The :code:`superscript` icon.\n SVG (IconName): The :code:`svg` icon.\n SWIMMING (IconName): The :code:`swimming` icon.\n SWIPE (IconName): The :code:`swipe` icon.\n SWITCH (IconName): The :code:`switch` icon.\n SWITCH_2 (IconName): The :code:`switch-2` icon.\n SWITCH_3 (IconName): The :code:`switch-3` icon.\n SWITCH_HORIZONTAL (IconName): The :code:`switch-horizontal` icon.\n SWITCH_VERTICAL (IconName): The :code:`switch-vertical` icon.\n SWORD (IconName): The :code:`sword` icon.\n SWORD_OFF (IconName): The :code:`sword-off` icon.\n SWORDS (IconName): The :code:`swords` icon.\n TABLE (IconName): The :code:`table` icon.\n TABLE_ALIAS (IconName): The :code:`table-alias` icon.\n TABLE_COLUMN (IconName): The :code:`table-column` icon.\n TABLE_DOWN (IconName): The :code:`table-down` icon.\n TABLE_EXPORT (IconName): The :code:`table-export` icon.\n TABLE_FILLED (IconName): The :code:`table-filled` icon.\n TABLE_HEART (IconName): The :code:`table-heart` icon.\n TABLE_IMPORT (IconName): The :code:`table-import` icon.\n TABLE_MINUS (IconName): The :code:`table-minus` icon.\n TABLE_OFF (IconName): The :code:`table-off` icon.\n TABLE_OPTIONS (IconName): The :code:`table-options` icon.\n TABLE_PLUS (IconName): The :code:`table-plus` icon.\n TABLE_ROW (IconName): The :code:`table-row` icon.\n TABLE_SHARE (IconName): The :code:`table-share` icon.\n TABLE_SHORTCUT (IconName): The :code:`table-shortcut` icon.\n TAG (IconName): The :code:`tag` icon.\n TAG_OFF (IconName): The :code:`tag-off` icon.\n TAGS (IconName): The :code:`tags` icon.\n TAGS_OFF (IconName): The :code:`tags-off` icon.\n TALLYMARK_1 (IconName): The :code:`tallymark-1` icon.\n TALLYMARK_2 (IconName): The :code:`tallymark-2` icon.\n TALLYMARK_3 (IconName): The :code:`tallymark-3` icon.\n TALLYMARK_4 (IconName): The :code:`tallymark-4` icon.\n TALLYMARKS (IconName): The :code:`tallymarks` icon.\n TANK (IconName): The :code:`tank` icon.\n TARGET (IconName): The :code:`target` icon.\n TARGET_ARROW (IconName): The :code:`target-arrow` icon.\n TARGET_OFF (IconName): The :code:`target-off` icon.\n TEAPOT (IconName): The :code:`teapot` icon.\n TELESCOPE (IconName): The :code:`telescope` icon.\n TELESCOPE_OFF (IconName): The :code:`telescope-off` icon.\n TEMPERATURE (IconName): The :code:`temperature` icon.\n TEMPERATURE_CELSIUS (IconName): The :code:`temperature-celsius` icon.\n TEMPERATURE_FAHRENHEIT (IconName): The :code:`temperature-fahrenheit` icon.\n TEMPERATURE_MINUS (IconName): The :code:`temperature-minus` icon.\n TEMPERATURE_OFF (IconName): The :code:`temperature-off` icon.\n TEMPERATURE_PLUS (IconName): The :code:`temperature-plus` icon.\n TEMPLATE (IconName): The :code:`template` icon.\n TEMPLATE_OFF (IconName): The :code:`template-off` icon.\n TENT (IconName): The :code:`tent` icon.\n TENT_OFF (IconName): The :code:`tent-off` icon.\n TERMINAL (IconName): The :code:`terminal` icon.\n TERMINAL_2 (IconName): The :code:`terminal-2` icon.\n TEST_PIPE (IconName): The :code:`test-pipe` icon.\n TEST_PIPE_2 (IconName): The :code:`test-pipe-2` icon.\n TEST_PIPE_OFF (IconName): The :code:`test-pipe-off` icon.\n TEX (IconName): The :code:`tex` icon.\n TEXT_CAPTION (IconName): The :code:`text-caption` icon.\n TEXT_COLOR (IconName): The :code:`text-color` icon.\n TEXT_DECREASE (IconName): The :code:`text-decrease` icon.\n TEXT_DIRECTION_LTR (IconName): The :code:`text-direction-ltr` icon.\n TEXT_DIRECTION_RTL (IconName): The :code:`text-direction-rtl` icon.\n TEXT_INCREASE (IconName): The :code:`text-increase` icon.\n TEXT_ORIENTATION (IconName): The :code:`text-orientation` icon.\n TEXT_PLUS (IconName): The :code:`text-plus` icon.\n TEXT_RECOGNITION (IconName): The :code:`text-recognition` icon.\n TEXT_RESIZE (IconName): The :code:`text-resize` icon.\n TEXT_SIZE (IconName): The :code:`text-size` icon.\n TEXT_SPELLCHECK (IconName): The :code:`text-spellcheck` icon.\n TEXT_WRAP (IconName): The :code:`text-wrap` icon.\n TEXT_WRAP_DISABLED (IconName): The :code:`text-wrap-disabled` icon.\n TEXTURE (IconName): The :code:`texture` icon.\n THEATER (IconName): The :code:`theater` icon.\n THERMOMETER (IconName): The :code:`thermometer` icon.\n THUMB_DOWN (IconName): The :code:`thumb-down` icon.\n THUMB_DOWN_FILLED (IconName): The :code:`thumb-down-filled` icon.\n THUMB_DOWN_OFF (IconName): The :code:`thumb-down-off` icon.\n THUMB_UP (IconName): The :code:`thumb-up` icon.\n THUMB_UP_FILLED (IconName): The :code:`thumb-up-filled` icon.\n THUMB_UP_OFF (IconName): The :code:`thumb-up-off` icon.\n TIC_TAC (IconName): The :code:`tic-tac` icon.\n TICKET (IconName): The :code:`ticket` icon.\n TICKET_OFF (IconName): The :code:`ticket-off` icon.\n TIE (IconName): The :code:`tie` icon.\n TILDE (IconName): The :code:`tilde` icon.\n TILT_SHIFT (IconName): The :code:`tilt-shift` icon.\n TILT_SHIFT_OFF (IconName): The :code:`tilt-shift-off` icon.\n TIME_DURATION_0 (IconName): The :code:`time-duration-0` icon.\n TIME_DURATION_10 (IconName): The :code:`time-duration-10` icon.\n TIME_DURATION_15 (IconName): The :code:`time-duration-15` icon.\n TIME_DURATION_30 (IconName): The :code:`time-duration-30` icon.\n TIME_DURATION_45 (IconName): The :code:`time-duration-45` icon.\n TIME_DURATION_5 (IconName): The :code:`time-duration-5` icon.\n TIME_DURATION_60 (IconName): The :code:`time-duration-60` icon.\n TIME_DURATION_90 (IconName): The :code:`time-duration-90` icon.\n TIME_DURATION_OFF (IconName): The :code:`time-duration-off` icon.\n TIMELINE (IconName): The :code:`timeline` icon.\n TIMELINE_EVENT (IconName): The :code:`timeline-event` icon.\n TIMELINE_EVENT_EXCLAMATION (IconName): The :code:`timeline-event-exclamation` icon.\n TIMELINE_EVENT_MINUS (IconName): The :code:`timeline-event-minus` icon.\n TIMELINE_EVENT_PLUS (IconName): The :code:`timeline-event-plus` icon.\n TIMELINE_EVENT_TEXT (IconName): The :code:`timeline-event-text` icon.\n TIMELINE_EVENT_X (IconName): The :code:`timeline-event-x` icon.\n TIR (IconName): The :code:`tir` icon.\n TOGGLE_LEFT (IconName): The :code:`toggle-left` icon.\n TOGGLE_RIGHT (IconName): The :code:`toggle-right` icon.\n TOILET_PAPER (IconName): The :code:`toilet-paper` icon.\n TOILET_PAPER_OFF (IconName): The :code:`toilet-paper-off` icon.\n TOML (IconName): The :code:`toml` icon.\n TOOL (IconName): The :code:`tool` icon.\n TOOLS (IconName): The :code:`tools` icon.\n TOOLS_KITCHEN (IconName): The :code:`tools-kitchen` icon.\n TOOLS_KITCHEN_2 (IconName): The :code:`tools-kitchen-2` icon.\n TOOLS_KITCHEN_2_OFF (IconName): The :code:`tools-kitchen-2-off` icon.\n TOOLS_KITCHEN_OFF (IconName): The :code:`tools-kitchen-off` icon.\n TOOLS_OFF (IconName): The :code:`tools-off` icon.\n TOOLTIP (IconName): The :code:`tooltip` icon.\n TOPOLOGY_BUS (IconName): The :code:`topology-bus` icon.\n TOPOLOGY_COMPLEX (IconName): The :code:`topology-complex` icon.\n TOPOLOGY_FULL (IconName): The :code:`topology-full` icon.\n TOPOLOGY_FULL_HIERARCHY (IconName): The :code:`topology-full-hierarchy` icon.\n TOPOLOGY_RING (IconName): The :code:`topology-ring` icon.\n TOPOLOGY_RING_2 (IconName): The :code:`topology-ring-2` icon.\n TOPOLOGY_RING_3 (IconName): The :code:`topology-ring-3` icon.\n TOPOLOGY_STAR (IconName): The :code:`topology-star` icon.\n TOPOLOGY_STAR_2 (IconName): The :code:`topology-star-2` icon.\n TOPOLOGY_STAR_3 (IconName): The :code:`topology-star-3` icon.\n TOPOLOGY_STAR_RING (IconName): The :code:`topology-star-ring` icon.\n TOPOLOGY_STAR_RING_2 (IconName): The :code:`topology-star-ring-2` icon.\n TOPOLOGY_STAR_RING_3 (IconName): The :code:`topology-star-ring-3` icon.\n TORII (IconName): The :code:`torii` icon.\n TORNADO (IconName): The :code:`tornado` icon.\n TOURNAMENT (IconName): The :code:`tournament` icon.\n TOWER (IconName): The :code:`tower` icon.\n TOWER_OFF (IconName): The :code:`tower-off` icon.\n TRACK (IconName): The :code:`track` icon.\n TRACTOR (IconName): The :code:`tractor` icon.\n TRADEMARK (IconName): The :code:`trademark` icon.\n TRAFFIC_CONE (IconName): The :code:`traffic-cone` icon.\n TRAFFIC_CONE_OFF (IconName): The :code:`traffic-cone-off` icon.\n TRAFFIC_LIGHTS (IconName): The :code:`traffic-lights` icon.\n TRAFFIC_LIGHTS_OFF (IconName): The :code:`traffic-lights-off` icon.\n TRAIN (IconName): The :code:`train` icon.\n TRANSFER_IN (IconName): The :code:`transfer-in` icon.\n TRANSFER_OUT (IconName): The :code:`transfer-out` icon.\n TRANSFORM (IconName): The :code:`transform` icon.\n TRANSFORM_FILLED (IconName): The :code:`transform-filled` icon.\n TRANSITION_BOTTOM (IconName): The :code:`transition-bottom` icon.\n TRANSITION_LEFT (IconName): The :code:`transition-left` icon.\n TRANSITION_RIGHT (IconName): The :code:`transition-right` icon.\n TRANSITION_TOP (IconName): The :code:`transition-top` icon.\n TRASH (IconName): The :code:`trash` icon.\n TRASH_FILLED (IconName): The :code:`trash-filled` icon.\n TRASH_OFF (IconName): The :code:`trash-off` icon.\n TRASH_X (IconName): The :code:`trash-x` icon.\n TRASH_X_FILLED (IconName): The :code:`trash-x-filled` icon.\n TREADMILL (IconName): The :code:`treadmill` icon.\n TREE (IconName): The :code:`tree` icon.\n TREES (IconName): The :code:`trees` icon.\n TREKKING (IconName): The :code:`trekking` icon.\n TRENDING_DOWN (IconName): The :code:`trending-down` icon.\n TRENDING_DOWN_2 (IconName): The :code:`trending-down-2` icon.\n TRENDING_DOWN_3 (IconName): The :code:`trending-down-3` icon.\n TRENDING_UP (IconName): The :code:`trending-up` icon.\n TRENDING_UP_2 (IconName): The :code:`trending-up-2` icon.\n TRENDING_UP_3 (IconName): The :code:`trending-up-3` icon.\n TRIANGLE (IconName): The :code:`triangle` icon.\n TRIANGLE_FILLED (IconName): The :code:`triangle-filled` icon.\n TRIANGLE_INVERTED (IconName): The :code:`triangle-inverted` icon.\n TRIANGLE_INVERTED_FILLED (IconName): The :code:`triangle-inverted-filled` icon.\n TRIANGLE_OFF (IconName): The :code:`triangle-off` icon.\n TRIANGLE_SQUARE_CIRCLE (IconName): The :code:`triangle-square-circle` icon.\n TRIANGLES (IconName): The :code:`triangles` icon.\n TRIDENT (IconName): The :code:`trident` icon.\n TROLLEY (IconName): The :code:`trolley` icon.\n TROPHY (IconName): The :code:`trophy` icon.\n TROPHY_FILLED (IconName): The :code:`trophy-filled` icon.\n TROPHY_OFF (IconName): The :code:`trophy-off` icon.\n TROWEL (IconName): The :code:`trowel` icon.\n TRUCK (IconName): The :code:`truck` icon.\n TRUCK_DELIVERY (IconName): The :code:`truck-delivery` icon.\n TRUCK_LOADING (IconName): The :code:`truck-loading` icon.\n TRUCK_OFF (IconName): The :code:`truck-off` icon.\n TRUCK_RETURN (IconName): The :code:`truck-return` icon.\n TXT (IconName): The :code:`txt` icon.\n TYPOGRAPHY (IconName): The :code:`typography` icon.\n TYPOGRAPHY_OFF (IconName): The :code:`typography-off` icon.\n UFO (IconName): The :code:`ufo` icon.\n UFO_OFF (IconName): The :code:`ufo-off` icon.\n UMBRELLA (IconName): The :code:`umbrella` icon.\n UMBRELLA_FILLED (IconName): The :code:`umbrella-filled` icon.\n UMBRELLA_OFF (IconName): The :code:`umbrella-off` icon.\n UNDERLINE (IconName): The :code:`underline` icon.\n UNLINK (IconName): The :code:`unlink` icon.\n UPLOAD (IconName): The :code:`upload` icon.\n URGENT (IconName): The :code:`urgent` icon.\n USB (IconName): The :code:`usb` icon.\n USER (IconName): The :code:`user` icon.\n USER_BOLT (IconName): The :code:`user-bolt` icon.\n USER_CANCEL (IconName): The :code:`user-cancel` icon.\n USER_CHECK (IconName): The :code:`user-check` icon.\n USER_CIRCLE (IconName): The :code:`user-circle` icon.\n USER_CODE (IconName): The :code:`user-code` icon.\n USER_COG (IconName): The :code:`user-cog` icon.\n USER_DOLLAR (IconName): The :code:`user-dollar` icon.\n USER_DOWN (IconName): The :code:`user-down` icon.\n USER_EDIT (IconName): The :code:`user-edit` icon.\n USER_EXCLAMATION (IconName): The :code:`user-exclamation` icon.\n USER_HEART (IconName): The :code:`user-heart` icon.\n USER_MINUS (IconName): The :code:`user-minus` icon.\n USER_OFF (IconName): The :code:`user-off` icon.\n USER_PAUSE (IconName): The :code:`user-pause` icon.\n USER_PIN (IconName): The :code:`user-pin` icon.\n USER_PLUS (IconName): The :code:`user-plus` icon.\n USER_QUESTION (IconName): The :code:`user-question` icon.\n USER_SEARCH (IconName): The :code:`user-search` icon.\n USER_SHARE (IconName): The :code:`user-share` icon.\n USER_SHIELD (IconName): The :code:`user-shield` icon.\n USER_STAR (IconName): The :code:`user-star` icon.\n USER_UP (IconName): The :code:`user-up` icon.\n USER_X (IconName): The :code:`user-x` icon.\n USERS (IconName): The :code:`users` icon.\n USERS_GROUP (IconName): The :code:`users-group` icon.\n USERS_MINUS (IconName): The :code:`users-minus` icon.\n USERS_PLUS (IconName): The :code:`users-plus` icon.\n UV_INDEX (IconName): The :code:`uv-index` icon.\n UX_CIRCLE (IconName): The :code:`ux-circle` icon.\n VACCINE (IconName): The :code:`vaccine` icon.\n VACCINE_BOTTLE (IconName): The :code:`vaccine-bottle` icon.\n VACCINE_BOTTLE_OFF (IconName): The :code:`vaccine-bottle-off` icon.\n VACCINE_OFF (IconName): The :code:`vaccine-off` icon.\n VACUUM_CLEANER (IconName): The :code:`vacuum-cleaner` icon.\n VARIABLE (IconName): The :code:`variable` icon.\n VARIABLE_MINUS (IconName): The :code:`variable-minus` icon.\n VARIABLE_OFF (IconName): The :code:`variable-off` icon.\n VARIABLE_PLUS (IconName): The :code:`variable-plus` icon.\n VECTOR (IconName): The :code:`vector` icon.\n VECTOR_BEZIER (IconName): The :code:`vector-bezier` icon.\n VECTOR_BEZIER_2 (IconName): The :code:`vector-bezier-2` icon.\n VECTOR_BEZIER_ARC (IconName): The :code:`vector-bezier-arc` icon.\n VECTOR_BEZIER_CIRCLE (IconName): The :code:`vector-bezier-circle` icon.\n VECTOR_OFF (IconName): The :code:`vector-off` icon.\n VECTOR_SPLINE (IconName): The :code:`vector-spline` icon.\n VECTOR_TRIANGLE (IconName): The :code:`vector-triangle` icon.\n VECTOR_TRIANGLE_OFF (IconName): The :code:`vector-triangle-off` icon.\n VENUS (IconName): The :code:`venus` icon.\n VERSIONS (IconName): The :code:`versions` icon.\n VERSIONS_FILLED (IconName): The :code:`versions-filled` icon.\n VERSIONS_OFF (IconName): The :code:`versions-off` icon.\n VIDEO (IconName): The :code:`video` icon.\n VIDEO_MINUS (IconName): The :code:`video-minus` icon.\n VIDEO_OFF (IconName): The :code:`video-off` icon.\n VIDEO_PLUS (IconName): The :code:`video-plus` icon.\n VIEW_360 (IconName): The :code:`view-360` icon.\n VIEW_360_OFF (IconName): The :code:`view-360-off` icon.\n VIEWFINDER (IconName): The :code:`viewfinder` icon.\n VIEWFINDER_OFF (IconName): The :code:`viewfinder-off` icon.\n VIEWPORT_NARROW (IconName): The :code:`viewport-narrow` icon.\n VIEWPORT_WIDE (IconName): The :code:`viewport-wide` icon.\n VINYL (IconName): The :code:`vinyl` icon.\n VIP (IconName): The :code:`vip` icon.\n VIP_OFF (IconName): The :code:`vip-off` icon.\n VIRUS (IconName): The :code:`virus` icon.\n VIRUS_OFF (IconName): The :code:`virus-off` icon.\n VIRUS_SEARCH (IconName): The :code:`virus-search` icon.\n VOCABULARY (IconName): The :code:`vocabulary` icon.\n VOCABULARY_OFF (IconName): The :code:`vocabulary-off` icon.\n VOLCANO (IconName): The :code:`volcano` icon.\n VOLUME (IconName): The :code:`volume` icon.\n VOLUME_2 (IconName): The :code:`volume-2` icon.\n VOLUME_3 (IconName): The :code:`volume-3` icon.\n VOLUME_OFF (IconName): The :code:`volume-off` icon.\n WALK (IconName): The :code:`walk` icon.\n WALL (IconName): The :code:`wall` icon.\n WALL_OFF (IconName): The :code:`wall-off` icon.\n WALLET (IconName): The :code:`wallet` icon.\n WALLET_OFF (IconName): The :code:`wallet-off` icon.\n WALLPAPER (IconName): The :code:`wallpaper` icon.\n WALLPAPER_OFF (IconName): The :code:`wallpaper-off` icon.\n WAND (IconName): The :code:`wand` icon.\n WAND_OFF (IconName): The :code:`wand-off` icon.\n WASH (IconName): The :code:`wash` icon.\n WASH_DRY (IconName): The :code:`wash-dry` icon.\n WASH_DRY_1 (IconName): The :code:`wash-dry-1` icon.\n WASH_DRY_2 (IconName): The :code:`wash-dry-2` icon.\n WASH_DRY_3 (IconName): The :code:`wash-dry-3` icon.\n WASH_DRY_A (IconName): The :code:`wash-dry-a` icon.\n WASH_DRY_DIP (IconName): The :code:`wash-dry-dip` icon.\n WASH_DRY_F (IconName): The :code:`wash-dry-f` icon.\n WASH_DRY_FLAT (IconName): The :code:`wash-dry-flat` icon.\n WASH_DRY_HANG (IconName): The :code:`wash-dry-hang` icon.\n WASH_DRY_OFF (IconName): The :code:`wash-dry-off` icon.\n WASH_DRY_P (IconName): The :code:`wash-dry-p` icon.\n WASH_DRY_SHADE (IconName): The :code:`wash-dry-shade` icon.\n WASH_DRY_W (IconName): The :code:`wash-dry-w` icon.\n WASH_DRYCLEAN (IconName): The :code:`wash-dryclean` icon.\n WASH_DRYCLEAN_OFF (IconName): The :code:`wash-dryclean-off` icon.\n WASH_ECO (IconName): The :code:`wash-eco` icon.\n WASH_GENTLE (IconName): The :code:`wash-gentle` icon.\n WASH_HAND (IconName): The :code:`wash-hand` icon.\n WASH_MACHINE (IconName): The :code:`wash-machine` icon.\n WASH_OFF (IconName): The :code:`wash-off` icon.\n WASH_PRESS (IconName): The :code:`wash-press` icon.\n WASH_TEMPERATURE_1 (IconName): The :code:`wash-temperature-1` icon.\n WASH_TEMPERATURE_2 (IconName): The :code:`wash-temperature-2` icon.\n WASH_TEMPERATURE_3 (IconName): The :code:`wash-temperature-3` icon.\n WASH_TEMPERATURE_4 (IconName): The :code:`wash-temperature-4` icon.\n WASH_TEMPERATURE_5 (IconName): The :code:`wash-temperature-5` icon.\n WASH_TEMPERATURE_6 (IconName): The :code:`wash-temperature-6` icon.\n WASH_TUMBLE_DRY (IconName): The :code:`wash-tumble-dry` icon.\n WASH_TUMBLE_OFF (IconName): The :code:`wash-tumble-off` icon.\n WATERPOLO (IconName): The :code:`waterpolo` icon.\n WAVE_SAW_TOOL (IconName): The :code:`wave-saw-tool` icon.\n WAVE_SINE (IconName): The :code:`wave-sine` icon.\n WAVE_SQUARE (IconName): The :code:`wave-square` icon.\n WEBHOOK (IconName): The :code:`webhook` icon.\n WEBHOOK_OFF (IconName): The :code:`webhook-off` icon.\n WEIGHT (IconName): The :code:`weight` icon.\n WHEELCHAIR (IconName): The :code:`wheelchair` icon.\n WHEELCHAIR_OFF (IconName): The :code:`wheelchair-off` icon.\n WHIRL (IconName): The :code:`whirl` icon.\n WIFI (IconName): The :code:`wifi` icon.\n WIFI_0 (IconName): The :code:`wifi-0` icon.\n WIFI_1 (IconName): The :code:`wifi-1` icon.\n WIFI_2 (IconName): The :code:`wifi-2` icon.\n WIFI_OFF (IconName): The :code:`wifi-off` icon.\n WIND (IconName): The :code:`wind` icon.\n WIND_OFF (IconName): The :code:`wind-off` icon.\n WINDMILL (IconName): The :code:`windmill` icon.\n WINDMILL_FILLED (IconName): The :code:`windmill-filled` icon.\n WINDMILL_OFF (IconName): The :code:`windmill-off` icon.\n WINDOW (IconName): The :code:`window` icon.\n WINDOW_MAXIMIZE (IconName): The :code:`window-maximize` icon.\n WINDOW_MINIMIZE (IconName): The :code:`window-minimize` icon.\n WINDOW_OFF (IconName): The :code:`window-off` icon.\n WINDSOCK (IconName): The :code:`windsock` icon.\n WIPER (IconName): The :code:`wiper` icon.\n WIPER_WASH (IconName): The :code:`wiper-wash` icon.\n WOMAN (IconName): The :code:`woman` icon.\n WOOD (IconName): The :code:`wood` icon.\n WORLD (IconName): The :code:`world` icon.\n WORLD_BOLT (IconName): The :code:`world-bolt` icon.\n WORLD_CANCEL (IconName): The :code:`world-cancel` icon.\n WORLD_CHECK (IconName): The :code:`world-check` icon.\n WORLD_CODE (IconName): The :code:`world-code` icon.\n WORLD_COG (IconName): The :code:`world-cog` icon.\n WORLD_DOLLAR (IconName): The :code:`world-dollar` icon.\n WORLD_DOWN (IconName): The :code:`world-down` icon.\n WORLD_DOWNLOAD (IconName): The :code:`world-download` icon.\n WORLD_EXCLAMATION (IconName): The :code:`world-exclamation` icon.\n WORLD_HEART (IconName): The :code:`world-heart` icon.\n WORLD_LATITUDE (IconName): The :code:`world-latitude` icon.\n WORLD_LONGITUDE (IconName): The :code:`world-longitude` icon.\n WORLD_MINUS (IconName): The :code:`world-minus` icon.\n WORLD_OFF (IconName): The :code:`world-off` icon.\n WORLD_PAUSE (IconName): The :code:`world-pause` icon.\n WORLD_PIN (IconName): The :code:`world-pin` icon.\n WORLD_PLUS (IconName): The :code:`world-plus` icon.\n WORLD_QUESTION (IconName): The :code:`world-question` icon.\n WORLD_SEARCH (IconName): The :code:`world-search` icon.\n WORLD_SHARE (IconName): The :code:`world-share` icon.\n WORLD_STAR (IconName): The :code:`world-star` icon.\n WORLD_UP (IconName): The :code:`world-up` icon.\n WORLD_UPLOAD (IconName): The :code:`world-upload` icon.\n WORLD_WWW (IconName): The :code:`world-www` icon.\n WORLD_X (IconName): The :code:`world-x` icon.\n WRECKING_BALL (IconName): The :code:`wrecking-ball` icon.\n WRITING (IconName): The :code:`writing` icon.\n WRITING_OFF (IconName): The :code:`writing-off` icon.\n WRITING_SIGN (IconName): The :code:`writing-sign` icon.\n WRITING_SIGN_OFF (IconName): The :code:`writing-sign-off` icon.\n X (IconName): The :code:`x` icon.\n XBOX_A (IconName): The :code:`xbox-a` icon.\n XBOX_B (IconName): The :code:`xbox-b` icon.\n XBOX_X (IconName): The :code:`xbox-x` icon.\n XBOX_Y (IconName): The :code:`xbox-y` icon.\n XD (IconName): The :code:`xd` icon.\n YIN_YANG (IconName): The :code:`yin-yang` icon.\n YIN_YANG_FILLED (IconName): The :code:`yin-yang-filled` icon.\n YOGA (IconName): The :code:`yoga` icon.\n ZEPPELIN (IconName): The :code:`zeppelin` icon.\n ZEPPELIN_OFF (IconName): The :code:`zeppelin-off` icon.\n ZIP (IconName): The :code:`zip` icon.\n ZODIAC_AQUARIUS (IconName): The :code:`zodiac-aquarius` icon.\n ZODIAC_ARIES (IconName): The :code:`zodiac-aries` icon.\n ZODIAC_CANCER (IconName): The :code:`zodiac-cancer` icon.\n ZODIAC_CAPRICORN (IconName): The :code:`zodiac-capricorn` icon.\n ZODIAC_GEMINI (IconName): The :code:`zodiac-gemini` icon.\n ZODIAC_LEO (IconName): The :code:`zodiac-leo` icon.\n ZODIAC_LIBRA (IconName): The :code:`zodiac-libra` icon.\n ZODIAC_PISCES (IconName): The :code:`zodiac-pisces` icon.\n ZODIAC_SAGITTARIUS (IconName): The :code:`zodiac-sagittarius` icon.\n ZODIAC_SCORPIO (IconName): The :code:`zodiac-scorpio` icon.\n ZODIAC_TAURUS (IconName): The :code:`zodiac-taurus` icon.\n ZODIAC_VIRGO (IconName): The :code:`zodiac-virgo` icon.\n ZOOM_CANCEL (IconName): The :code:`zoom-cancel` icon.\n ZOOM_CHECK (IconName): The :code:`zoom-check` icon.\n ZOOM_CHECK_FILLED (IconName): The :code:`zoom-check-filled` icon.\n ZOOM_CODE (IconName): The :code:`zoom-code` icon.\n ZOOM_EXCLAMATION (IconName): The :code:`zoom-exclamation` icon.\n ZOOM_FILLED (IconName): The :code:`zoom-filled` icon.\n ZOOM_IN (IconName): The :code:`zoom-in` icon.\n ZOOM_IN_AREA (IconName): The :code:`zoom-in-area` icon.\n ZOOM_IN_AREA_FILLED (IconName): The :code:`zoom-in-area-filled` icon.\n ZOOM_IN_FILLED (IconName): The :code:`zoom-in-filled` icon.\n ZOOM_MONEY (IconName): The :code:`zoom-money` icon.\n ZOOM_OUT (IconName): The :code:`zoom-out` icon.\n ZOOM_OUT_AREA (IconName): The :code:`zoom-out-area` icon.\n ZOOM_OUT_FILLED (IconName): The :code:`zoom-out-filled` icon.\n ZOOM_PAN (IconName): The :code:`zoom-pan` icon.\n ZOOM_QUESTION (IconName): The :code:`zoom-question` icon.\n ZOOM_REPLACE (IconName): The :code:`zoom-replace` icon.\n ZOOM_RESET (IconName): The :code:`zoom-reset` icon.\n ZZZ (IconName): The :code:`zzz` icon.\n ZZZ_OFF (IconName): The :code:`zzz-off` icon.\n \"\"\"\n" }, { "path": "viser/src/viser/_icons_enum.pyi", "content": "# Automatically generated by `_icons_generate_enum.py`\n# See https://tabler-icons.io/\nfrom typing import NewType\n\nIconName = NewType(\"IconName\", str)\n\"\"\"Name of an icon. Should be generated via `viser.Icon.*`.\"\"\"\n\nclass Icon:\n \"\"\"'Enum' class for referencing Tabler icons.\n\n We don't subclass enum.Enum for performance reasons -- importing an enum with\n thousands of names can result in import times in the hundreds of milliseconds.\n \"\"\"\n\n ICON_123: IconName = IconName(\"123\")\n ICON_24_HOURS: IconName = IconName(\"24-hours\")\n ICON_2FA: IconName = IconName(\"2fa\")\n ICON_360: IconName = IconName(\"360\")\n ICON_360_VIEW: IconName = IconName(\"360-view\")\n ICON_3D_CUBE_SPHERE: IconName = IconName(\"3d-cube-sphere\")\n ICON_3D_CUBE_SPHERE_OFF: IconName = IconName(\"3d-cube-sphere-off\")\n ICON_3D_ROTATE: IconName = IconName(\"3d-rotate\")\n A_B: IconName = IconName(\"a-b\")\n A_B_2: IconName = IconName(\"a-b-2\")\n A_B_OFF: IconName = IconName(\"a-b-off\")\n ABACUS: IconName = IconName(\"abacus\")\n ABACUS_OFF: IconName = IconName(\"abacus-off\")\n ABC: IconName = IconName(\"abc\")\n ACCESS_POINT: IconName = IconName(\"access-point\")\n ACCESS_POINT_OFF: IconName = IconName(\"access-point-off\")\n ACCESSIBLE: IconName = IconName(\"accessible\")\n ACCESSIBLE_OFF: IconName = IconName(\"accessible-off\")\n ACCESSIBLE_OFF_FILLED: IconName = IconName(\"accessible-off-filled\")\n ACTIVITY: IconName = IconName(\"activity\")\n ACTIVITY_HEARTBEAT: IconName = IconName(\"activity-heartbeat\")\n AD: IconName = IconName(\"ad\")\n AD_2: IconName = IconName(\"ad-2\")\n AD_CIRCLE: IconName = IconName(\"ad-circle\")\n AD_CIRCLE_FILLED: IconName = IconName(\"ad-circle-filled\")\n AD_CIRCLE_OFF: IconName = IconName(\"ad-circle-off\")\n AD_FILLED: IconName = IconName(\"ad-filled\")\n AD_OFF: IconName = IconName(\"ad-off\")\n ADDRESS_BOOK: IconName = IconName(\"address-book\")\n ADDRESS_BOOK_OFF: IconName = IconName(\"address-book-off\")\n ADJUSTMENTS: IconName = IconName(\"adjustments\")\n ADJUSTMENTS_ALT: IconName = IconName(\"adjustments-alt\")\n ADJUSTMENTS_BOLT: IconName = IconName(\"adjustments-bolt\")\n ADJUSTMENTS_CANCEL: IconName = IconName(\"adjustments-cancel\")\n ADJUSTMENTS_CHECK: IconName = IconName(\"adjustments-check\")\n ADJUSTMENTS_CODE: IconName = IconName(\"adjustments-code\")\n ADJUSTMENTS_COG: IconName = IconName(\"adjustments-cog\")\n ADJUSTMENTS_DOLLAR: IconName = IconName(\"adjustments-dollar\")\n ADJUSTMENTS_DOWN: IconName = IconName(\"adjustments-down\")\n ADJUSTMENTS_EXCLAMATION: IconName = IconName(\"adjustments-exclamation\")\n ADJUSTMENTS_FILLED: IconName = IconName(\"adjustments-filled\")\n ADJUSTMENTS_HEART: IconName = IconName(\"adjustments-heart\")\n ADJUSTMENTS_HORIZONTAL: IconName = IconName(\"adjustments-horizontal\")\n ADJUSTMENTS_MINUS: IconName = IconName(\"adjustments-minus\")\n ADJUSTMENTS_OFF: IconName = IconName(\"adjustments-off\")\n ADJUSTMENTS_PAUSE: IconName = IconName(\"adjustments-pause\")\n ADJUSTMENTS_PIN: IconName = IconName(\"adjustments-pin\")\n ADJUSTMENTS_PLUS: IconName = IconName(\"adjustments-plus\")\n ADJUSTMENTS_QUESTION: IconName = IconName(\"adjustments-question\")\n ADJUSTMENTS_SEARCH: IconName = IconName(\"adjustments-search\")\n ADJUSTMENTS_SHARE: IconName = IconName(\"adjustments-share\")\n ADJUSTMENTS_STAR: IconName = IconName(\"adjustments-star\")\n ADJUSTMENTS_UP: IconName = IconName(\"adjustments-up\")\n ADJUSTMENTS_X: IconName = IconName(\"adjustments-x\")\n AERIAL_LIFT: IconName = IconName(\"aerial-lift\")\n AFFILIATE: IconName = IconName(\"affiliate\")\n AFFILIATE_FILLED: IconName = IconName(\"affiliate-filled\")\n AIR_BALLOON: IconName = IconName(\"air-balloon\")\n AIR_CONDITIONING: IconName = IconName(\"air-conditioning\")\n AIR_CONDITIONING_DISABLED: IconName = IconName(\"air-conditioning-disabled\")\n ALARM: IconName = IconName(\"alarm\")\n ALARM_FILLED: IconName = IconName(\"alarm-filled\")\n ALARM_MINUS: IconName = IconName(\"alarm-minus\")\n ALARM_MINUS_FILLED: IconName = IconName(\"alarm-minus-filled\")\n ALARM_OFF: IconName = IconName(\"alarm-off\")\n ALARM_PLUS: IconName = IconName(\"alarm-plus\")\n ALARM_PLUS_FILLED: IconName = IconName(\"alarm-plus-filled\")\n ALARM_SNOOZE: IconName = IconName(\"alarm-snooze\")\n ALARM_SNOOZE_FILLED: IconName = IconName(\"alarm-snooze-filled\")\n ALBUM: IconName = IconName(\"album\")\n ALBUM_OFF: IconName = IconName(\"album-off\")\n ALERT_CIRCLE: IconName = IconName(\"alert-circle\")\n ALERT_CIRCLE_FILLED: IconName = IconName(\"alert-circle-filled\")\n ALERT_HEXAGON: IconName = IconName(\"alert-hexagon\")\n ALERT_HEXAGON_FILLED: IconName = IconName(\"alert-hexagon-filled\")\n ALERT_OCTAGON: IconName = IconName(\"alert-octagon\")\n ALERT_OCTAGON_FILLED: IconName = IconName(\"alert-octagon-filled\")\n ALERT_SMALL: IconName = IconName(\"alert-small\")\n ALERT_SQUARE: IconName = IconName(\"alert-square\")\n ALERT_SQUARE_FILLED: IconName = IconName(\"alert-square-filled\")\n ALERT_SQUARE_ROUNDED: IconName = IconName(\"alert-square-rounded\")\n ALERT_SQUARE_ROUNDED_FILLED: IconName = IconName(\"alert-square-rounded-filled\")\n ALERT_TRIANGLE: IconName = IconName(\"alert-triangle\")\n ALERT_TRIANGLE_FILLED: IconName = IconName(\"alert-triangle-filled\")\n ALIEN: IconName = IconName(\"alien\")\n ALIEN_FILLED: IconName = IconName(\"alien-filled\")\n ALIGN_BOX_BOTTOM_CENTER: IconName = IconName(\"align-box-bottom-center\")\n ALIGN_BOX_BOTTOM_CENTER_FILLED: IconName = IconName(\n \"align-box-bottom-center-filled\"\n )\n ALIGN_BOX_BOTTOM_LEFT: IconName = IconName(\"align-box-bottom-left\")\n ALIGN_BOX_BOTTOM_LEFT_FILLED: IconName = IconName(\"align-box-bottom-left-filled\")\n ALIGN_BOX_BOTTOM_RIGHT: IconName = IconName(\"align-box-bottom-right\")\n ALIGN_BOX_BOTTOM_RIGHT_FILLED: IconName = IconName(\"align-box-bottom-right-filled\")\n ALIGN_BOX_CENTER_BOTTOM: IconName = IconName(\"align-box-center-bottom\")\n ALIGN_BOX_CENTER_MIDDLE: IconName = IconName(\"align-box-center-middle\")\n ALIGN_BOX_CENTER_MIDDLE_FILLED: IconName = IconName(\n \"align-box-center-middle-filled\"\n )\n ALIGN_BOX_CENTER_STRETCH: IconName = IconName(\"align-box-center-stretch\")\n ALIGN_BOX_CENTER_TOP: IconName = IconName(\"align-box-center-top\")\n ALIGN_BOX_LEFT_BOTTOM: IconName = IconName(\"align-box-left-bottom\")\n ALIGN_BOX_LEFT_BOTTOM_FILLED: IconName = IconName(\"align-box-left-bottom-filled\")\n ALIGN_BOX_LEFT_MIDDLE: IconName = IconName(\"align-box-left-middle\")\n ALIGN_BOX_LEFT_MIDDLE_FILLED: IconName = IconName(\"align-box-left-middle-filled\")\n ALIGN_BOX_LEFT_STRETCH: IconName = IconName(\"align-box-left-stretch\")\n ALIGN_BOX_LEFT_TOP: IconName = IconName(\"align-box-left-top\")\n ALIGN_BOX_LEFT_TOP_FILLED: IconName = IconName(\"align-box-left-top-filled\")\n ALIGN_BOX_RIGHT_BOTTOM: IconName = IconName(\"align-box-right-bottom\")\n ALIGN_BOX_RIGHT_BOTTOM_FILLED: IconName = IconName(\"align-box-right-bottom-filled\")\n ALIGN_BOX_RIGHT_MIDDLE: IconName = IconName(\"align-box-right-middle\")\n ALIGN_BOX_RIGHT_MIDDLE_FILLED: IconName = IconName(\"align-box-right-middle-filled\")\n ALIGN_BOX_RIGHT_STRETCH: IconName = IconName(\"align-box-right-stretch\")\n ALIGN_BOX_RIGHT_TOP: IconName = IconName(\"align-box-right-top\")\n ALIGN_BOX_RIGHT_TOP_FILLED: IconName = IconName(\"align-box-right-top-filled\")\n ALIGN_BOX_TOP_CENTER: IconName = IconName(\"align-box-top-center\")\n ALIGN_BOX_TOP_CENTER_FILLED: IconName = IconName(\"align-box-top-center-filled\")\n ALIGN_BOX_TOP_LEFT: IconName = IconName(\"align-box-top-left\")\n ALIGN_BOX_TOP_LEFT_FILLED: IconName = IconName(\"align-box-top-left-filled\")\n ALIGN_BOX_TOP_RIGHT: IconName = IconName(\"align-box-top-right\")\n ALIGN_BOX_TOP_RIGHT_FILLED: IconName = IconName(\"align-box-top-right-filled\")\n ALIGN_CENTER: IconName = IconName(\"align-center\")\n ALIGN_JUSTIFIED: IconName = IconName(\"align-justified\")\n ALIGN_LEFT: IconName = IconName(\"align-left\")\n ALIGN_RIGHT: IconName = IconName(\"align-right\")\n ALPHA: IconName = IconName(\"alpha\")\n ALPHABET_CYRILLIC: IconName = IconName(\"alphabet-cyrillic\")\n ALPHABET_GREEK: IconName = IconName(\"alphabet-greek\")\n ALPHABET_LATIN: IconName = IconName(\"alphabet-latin\")\n AMBULANCE: IconName = IconName(\"ambulance\")\n AMPERSAND: IconName = IconName(\"ampersand\")\n ANALYZE: IconName = IconName(\"analyze\")\n ANALYZE_FILLED: IconName = IconName(\"analyze-filled\")\n ANALYZE_OFF: IconName = IconName(\"analyze-off\")\n ANCHOR: IconName = IconName(\"anchor\")\n ANCHOR_OFF: IconName = IconName(\"anchor-off\")\n ANGLE: IconName = IconName(\"angle\")\n ANKH: IconName = IconName(\"ankh\")\n ANTENNA: IconName = IconName(\"antenna\")\n ANTENNA_BARS_1: IconName = IconName(\"antenna-bars-1\")\n ANTENNA_BARS_2: IconName = IconName(\"antenna-bars-2\")\n ANTENNA_BARS_3: IconName = IconName(\"antenna-bars-3\")\n ANTENNA_BARS_4: IconName = IconName(\"antenna-bars-4\")\n ANTENNA_BARS_5: IconName = IconName(\"antenna-bars-5\")\n ANTENNA_BARS_OFF: IconName = IconName(\"antenna-bars-off\")\n ANTENNA_OFF: IconName = IconName(\"antenna-off\")\n APERTURE: IconName = IconName(\"aperture\")\n APERTURE_OFF: IconName = IconName(\"aperture-off\")\n API: IconName = IconName(\"api\")\n API_APP: IconName = IconName(\"api-app\")\n API_APP_OFF: IconName = IconName(\"api-app-off\")\n API_OFF: IconName = IconName(\"api-off\")\n APP_WINDOW: IconName = IconName(\"app-window\")\n APP_WINDOW_FILLED: IconName = IconName(\"app-window-filled\")\n APPLE: IconName = IconName(\"apple\")\n APPS: IconName = IconName(\"apps\")\n APPS_FILLED: IconName = IconName(\"apps-filled\")\n APPS_OFF: IconName = IconName(\"apps-off\")\n ARCHIVE: IconName = IconName(\"archive\")\n ARCHIVE_FILLED: IconName = IconName(\"archive-filled\")\n ARCHIVE_OFF: IconName = IconName(\"archive-off\")\n ARMCHAIR: IconName = IconName(\"armchair\")\n ARMCHAIR_2: IconName = IconName(\"armchair-2\")\n ARMCHAIR_2_OFF: IconName = IconName(\"armchair-2-off\")\n ARMCHAIR_OFF: IconName = IconName(\"armchair-off\")\n ARROW_AUTOFIT_CONTENT: IconName = IconName(\"arrow-autofit-content\")\n ARROW_AUTOFIT_CONTENT_FILLED: IconName = IconName(\"arrow-autofit-content-filled\")\n ARROW_AUTOFIT_DOWN: IconName = IconName(\"arrow-autofit-down\")\n ARROW_AUTOFIT_HEIGHT: IconName = IconName(\"arrow-autofit-height\")\n ARROW_AUTOFIT_LEFT: IconName = IconName(\"arrow-autofit-left\")\n ARROW_AUTOFIT_RIGHT: IconName = IconName(\"arrow-autofit-right\")\n ARROW_AUTOFIT_UP: IconName = IconName(\"arrow-autofit-up\")\n ARROW_AUTOFIT_WIDTH: IconName = IconName(\"arrow-autofit-width\")\n ARROW_BACK: IconName = IconName(\"arrow-back\")\n ARROW_BACK_UP: IconName = IconName(\"arrow-back-up\")\n ARROW_BACK_UP_DOUBLE: IconName = IconName(\"arrow-back-up-double\")\n ARROW_BADGE_DOWN: IconName = IconName(\"arrow-badge-down\")\n ARROW_BADGE_DOWN_FILLED: IconName = IconName(\"arrow-badge-down-filled\")\n ARROW_BADGE_LEFT: IconName = IconName(\"arrow-badge-left\")\n ARROW_BADGE_LEFT_FILLED: IconName = IconName(\"arrow-badge-left-filled\")\n ARROW_BADGE_RIGHT: IconName = IconName(\"arrow-badge-right\")\n ARROW_BADGE_RIGHT_FILLED: IconName = IconName(\"arrow-badge-right-filled\")\n ARROW_BADGE_UP: IconName = IconName(\"arrow-badge-up\")\n ARROW_BADGE_UP_FILLED: IconName = IconName(\"arrow-badge-up-filled\")\n ARROW_BAR_BOTH: IconName = IconName(\"arrow-bar-both\")\n ARROW_BAR_DOWN: IconName = IconName(\"arrow-bar-down\")\n ARROW_BAR_LEFT: IconName = IconName(\"arrow-bar-left\")\n ARROW_BAR_RIGHT: IconName = IconName(\"arrow-bar-right\")\n ARROW_BAR_TO_DOWN: IconName = IconName(\"arrow-bar-to-down\")\n ARROW_BAR_TO_LEFT: IconName = IconName(\"arrow-bar-to-left\")\n ARROW_BAR_TO_RIGHT: IconName = IconName(\"arrow-bar-to-right\")\n ARROW_BAR_TO_UP: IconName = IconName(\"arrow-bar-to-up\")\n ARROW_BAR_UP: IconName = IconName(\"arrow-bar-up\")\n ARROW_BEAR_LEFT: IconName = IconName(\"arrow-bear-left\")\n ARROW_BEAR_LEFT_2: IconName = IconName(\"arrow-bear-left-2\")\n ARROW_BEAR_RIGHT: IconName = IconName(\"arrow-bear-right\")\n ARROW_BEAR_RIGHT_2: IconName = IconName(\"arrow-bear-right-2\")\n ARROW_BIG_DOWN: IconName = IconName(\"arrow-big-down\")\n ARROW_BIG_DOWN_FILLED: IconName = IconName(\"arrow-big-down-filled\")\n ARROW_BIG_DOWN_LINE: IconName = IconName(\"arrow-big-down-line\")\n ARROW_BIG_DOWN_LINE_FILLED: IconName = IconName(\"arrow-big-down-line-filled\")\n ARROW_BIG_DOWN_LINES: IconName = IconName(\"arrow-big-down-lines\")\n ARROW_BIG_DOWN_LINES_FILLED: IconName = IconName(\"arrow-big-down-lines-filled\")\n ARROW_BIG_LEFT: IconName = IconName(\"arrow-big-left\")\n ARROW_BIG_LEFT_FILLED: IconName = IconName(\"arrow-big-left-filled\")\n ARROW_BIG_LEFT_LINE: IconName = IconName(\"arrow-big-left-line\")\n ARROW_BIG_LEFT_LINE_FILLED: IconName = IconName(\"arrow-big-left-line-filled\")\n ARROW_BIG_LEFT_LINES: IconName = IconName(\"arrow-big-left-lines\")\n ARROW_BIG_LEFT_LINES_FILLED: IconName = IconName(\"arrow-big-left-lines-filled\")\n ARROW_BIG_RIGHT: IconName = IconName(\"arrow-big-right\")\n ARROW_BIG_RIGHT_FILLED: IconName = IconName(\"arrow-big-right-filled\")\n ARROW_BIG_RIGHT_LINE: IconName = IconName(\"arrow-big-right-line\")\n ARROW_BIG_RIGHT_LINE_FILLED: IconName = IconName(\"arrow-big-right-line-filled\")\n ARROW_BIG_RIGHT_LINES: IconName = IconName(\"arrow-big-right-lines\")\n ARROW_BIG_RIGHT_LINES_FILLED: IconName = IconName(\"arrow-big-right-lines-filled\")\n ARROW_BIG_UP: IconName = IconName(\"arrow-big-up\")\n ARROW_BIG_UP_FILLED: IconName = IconName(\"arrow-big-up-filled\")\n ARROW_BIG_UP_LINE: IconName = IconName(\"arrow-big-up-line\")\n ARROW_BIG_UP_LINE_FILLED: IconName = IconName(\"arrow-big-up-line-filled\")\n ARROW_BIG_UP_LINES: IconName = IconName(\"arrow-big-up-lines\")\n ARROW_BIG_UP_LINES_FILLED: IconName = IconName(\"arrow-big-up-lines-filled\")\n ARROW_BOUNCE: IconName = IconName(\"arrow-bounce\")\n ARROW_CAPSULE: IconName = IconName(\"arrow-capsule\")\n ARROW_CURVE_LEFT: IconName = IconName(\"arrow-curve-left\")\n ARROW_CURVE_RIGHT: IconName = IconName(\"arrow-curve-right\")\n ARROW_DOWN: IconName = IconName(\"arrow-down\")\n ARROW_DOWN_BAR: IconName = IconName(\"arrow-down-bar\")\n ARROW_DOWN_CIRCLE: IconName = IconName(\"arrow-down-circle\")\n ARROW_DOWN_LEFT: IconName = IconName(\"arrow-down-left\")\n ARROW_DOWN_LEFT_CIRCLE: IconName = IconName(\"arrow-down-left-circle\")\n ARROW_DOWN_RHOMBUS: IconName = IconName(\"arrow-down-rhombus\")\n ARROW_DOWN_RIGHT: IconName = IconName(\"arrow-down-right\")\n ARROW_DOWN_RIGHT_CIRCLE: IconName = IconName(\"arrow-down-right-circle\")\n ARROW_DOWN_SQUARE: IconName = IconName(\"arrow-down-square\")\n ARROW_DOWN_TAIL: IconName = IconName(\"arrow-down-tail\")\n ARROW_ELBOW_LEFT: IconName = IconName(\"arrow-elbow-left\")\n ARROW_ELBOW_RIGHT: IconName = IconName(\"arrow-elbow-right\")\n ARROW_FORK: IconName = IconName(\"arrow-fork\")\n ARROW_FORWARD: IconName = IconName(\"arrow-forward\")\n ARROW_FORWARD_UP: IconName = IconName(\"arrow-forward-up\")\n ARROW_FORWARD_UP_DOUBLE: IconName = IconName(\"arrow-forward-up-double\")\n ARROW_GUIDE: IconName = IconName(\"arrow-guide\")\n ARROW_ITERATION: IconName = IconName(\"arrow-iteration\")\n ARROW_LEFT: IconName = IconName(\"arrow-left\")\n ARROW_LEFT_BAR: IconName = IconName(\"arrow-left-bar\")\n ARROW_LEFT_CIRCLE: IconName = IconName(\"arrow-left-circle\")\n ARROW_LEFT_RHOMBUS: IconName = IconName(\"arrow-left-rhombus\")\n ARROW_LEFT_RIGHT: IconName = IconName(\"arrow-left-right\")\n ARROW_LEFT_SQUARE: IconName = IconName(\"arrow-left-square\")\n ARROW_LEFT_TAIL: IconName = IconName(\"arrow-left-tail\")\n ARROW_LOOP_LEFT: IconName = IconName(\"arrow-loop-left\")\n ARROW_LOOP_LEFT_2: IconName = IconName(\"arrow-loop-left-2\")\n ARROW_LOOP_RIGHT: IconName = IconName(\"arrow-loop-right\")\n ARROW_LOOP_RIGHT_2: IconName = IconName(\"arrow-loop-right-2\")\n ARROW_MERGE: IconName = IconName(\"arrow-merge\")\n ARROW_MERGE_BOTH: IconName = IconName(\"arrow-merge-both\")\n ARROW_MERGE_LEFT: IconName = IconName(\"arrow-merge-left\")\n ARROW_MERGE_RIGHT: IconName = IconName(\"arrow-merge-right\")\n ARROW_MOVE_DOWN: IconName = IconName(\"arrow-move-down\")\n ARROW_MOVE_LEFT: IconName = IconName(\"arrow-move-left\")\n ARROW_MOVE_RIGHT: IconName = IconName(\"arrow-move-right\")\n ARROW_MOVE_UP: IconName = IconName(\"arrow-move-up\")\n ARROW_NARROW_DOWN: IconName = IconName(\"arrow-narrow-down\")\n ARROW_NARROW_LEFT: IconName = IconName(\"arrow-narrow-left\")\n ARROW_NARROW_RIGHT: IconName = IconName(\"arrow-narrow-right\")\n ARROW_NARROW_UP: IconName = IconName(\"arrow-narrow-up\")\n ARROW_RAMP_LEFT: IconName = IconName(\"arrow-ramp-left\")\n ARROW_RAMP_LEFT_2: IconName = IconName(\"arrow-ramp-left-2\")\n ARROW_RAMP_LEFT_3: IconName = IconName(\"arrow-ramp-left-3\")\n ARROW_RAMP_RIGHT: IconName = IconName(\"arrow-ramp-right\")\n ARROW_RAMP_RIGHT_2: IconName = IconName(\"arrow-ramp-right-2\")\n ARROW_RAMP_RIGHT_3: IconName = IconName(\"arrow-ramp-right-3\")\n ARROW_RIGHT: IconName = IconName(\"arrow-right\")\n ARROW_RIGHT_BAR: IconName = IconName(\"arrow-right-bar\")\n ARROW_RIGHT_CIRCLE: IconName = IconName(\"arrow-right-circle\")\n ARROW_RIGHT_RHOMBUS: IconName = IconName(\"arrow-right-rhombus\")\n ARROW_RIGHT_SQUARE: IconName = IconName(\"arrow-right-square\")\n ARROW_RIGHT_TAIL: IconName = IconName(\"arrow-right-tail\")\n ARROW_ROTARY_FIRST_LEFT: IconName = IconName(\"arrow-rotary-first-left\")\n ARROW_ROTARY_FIRST_RIGHT: IconName = IconName(\"arrow-rotary-first-right\")\n ARROW_ROTARY_LAST_LEFT: IconName = IconName(\"arrow-rotary-last-left\")\n ARROW_ROTARY_LAST_RIGHT: IconName = IconName(\"arrow-rotary-last-right\")\n ARROW_ROTARY_LEFT: IconName = IconName(\"arrow-rotary-left\")\n ARROW_ROTARY_RIGHT: IconName = IconName(\"arrow-rotary-right\")\n ARROW_ROTARY_STRAIGHT: IconName = IconName(\"arrow-rotary-straight\")\n ARROW_ROUNDABOUT_LEFT: IconName = IconName(\"arrow-roundabout-left\")\n ARROW_ROUNDABOUT_RIGHT: IconName = IconName(\"arrow-roundabout-right\")\n ARROW_SHARP_TURN_LEFT: IconName = IconName(\"arrow-sharp-turn-left\")\n ARROW_SHARP_TURN_RIGHT: IconName = IconName(\"arrow-sharp-turn-right\")\n ARROW_UP: IconName = IconName(\"arrow-up\")\n ARROW_UP_BAR: IconName = IconName(\"arrow-up-bar\")\n ARROW_UP_CIRCLE: IconName = IconName(\"arrow-up-circle\")\n ARROW_UP_LEFT: IconName = IconName(\"arrow-up-left\")\n ARROW_UP_LEFT_CIRCLE: IconName = IconName(\"arrow-up-left-circle\")\n ARROW_UP_RHOMBUS: IconName = IconName(\"arrow-up-rhombus\")\n ARROW_UP_RIGHT: IconName = IconName(\"arrow-up-right\")\n ARROW_UP_RIGHT_CIRCLE: IconName = IconName(\"arrow-up-right-circle\")\n ARROW_UP_SQUARE: IconName = IconName(\"arrow-up-square\")\n ARROW_UP_TAIL: IconName = IconName(\"arrow-up-tail\")\n ARROW_WAVE_LEFT_DOWN: IconName = IconName(\"arrow-wave-left-down\")\n ARROW_WAVE_LEFT_UP: IconName = IconName(\"arrow-wave-left-up\")\n ARROW_WAVE_RIGHT_DOWN: IconName = IconName(\"arrow-wave-right-down\")\n ARROW_WAVE_RIGHT_UP: IconName = IconName(\"arrow-wave-right-up\")\n ARROW_ZIG_ZAG: IconName = IconName(\"arrow-zig-zag\")\n ARROWS_CROSS: IconName = IconName(\"arrows-cross\")\n ARROWS_DIAGONAL: IconName = IconName(\"arrows-diagonal\")\n ARROWS_DIAGONAL_2: IconName = IconName(\"arrows-diagonal-2\")\n ARROWS_DIAGONAL_MINIMIZE: IconName = IconName(\"arrows-diagonal-minimize\")\n ARROWS_DIAGONAL_MINIMIZE_2: IconName = IconName(\"arrows-diagonal-minimize-2\")\n ARROWS_DIFF: IconName = IconName(\"arrows-diff\")\n ARROWS_DOUBLE_NE_SW: IconName = IconName(\"arrows-double-ne-sw\")\n ARROWS_DOUBLE_NW_SE: IconName = IconName(\"arrows-double-nw-se\")\n ARROWS_DOUBLE_SE_NW: IconName = IconName(\"arrows-double-se-nw\")\n ARROWS_DOUBLE_SW_NE: IconName = IconName(\"arrows-double-sw-ne\")\n ARROWS_DOWN: IconName = IconName(\"arrows-down\")\n ARROWS_DOWN_UP: IconName = IconName(\"arrows-down-up\")\n ARROWS_EXCHANGE: IconName = IconName(\"arrows-exchange\")\n ARROWS_EXCHANGE_2: IconName = IconName(\"arrows-exchange-2\")\n ARROWS_HORIZONTAL: IconName = IconName(\"arrows-horizontal\")\n ARROWS_JOIN: IconName = IconName(\"arrows-join\")\n ARROWS_JOIN_2: IconName = IconName(\"arrows-join-2\")\n ARROWS_LEFT: IconName = IconName(\"arrows-left\")\n ARROWS_LEFT_DOWN: IconName = IconName(\"arrows-left-down\")\n ARROWS_LEFT_RIGHT: IconName = IconName(\"arrows-left-right\")\n ARROWS_MAXIMIZE: IconName = IconName(\"arrows-maximize\")\n ARROWS_MINIMIZE: IconName = IconName(\"arrows-minimize\")\n ARROWS_MOVE: IconName = IconName(\"arrows-move\")\n ARROWS_MOVE_HORIZONTAL: IconName = IconName(\"arrows-move-horizontal\")\n ARROWS_MOVE_VERTICAL: IconName = IconName(\"arrows-move-vertical\")\n ARROWS_RANDOM: IconName = IconName(\"arrows-random\")\n ARROWS_RIGHT: IconName = IconName(\"arrows-right\")\n ARROWS_RIGHT_DOWN: IconName = IconName(\"arrows-right-down\")\n ARROWS_RIGHT_LEFT: IconName = IconName(\"arrows-right-left\")\n ARROWS_SHUFFLE: IconName = IconName(\"arrows-shuffle\")\n ARROWS_SHUFFLE_2: IconName = IconName(\"arrows-shuffle-2\")\n ARROWS_SORT: IconName = IconName(\"arrows-sort\")\n ARROWS_SPLIT: IconName = IconName(\"arrows-split\")\n ARROWS_SPLIT_2: IconName = IconName(\"arrows-split-2\")\n ARROWS_TRANSFER_DOWN: IconName = IconName(\"arrows-transfer-down\")\n ARROWS_TRANSFER_UP: IconName = IconName(\"arrows-transfer-up\")\n ARROWS_UP: IconName = IconName(\"arrows-up\")\n ARROWS_UP_DOWN: IconName = IconName(\"arrows-up-down\")\n ARROWS_UP_LEFT: IconName = IconName(\"arrows-up-left\")\n ARROWS_UP_RIGHT: IconName = IconName(\"arrows-up-right\")\n ARROWS_VERTICAL: IconName = IconName(\"arrows-vertical\")\n ARTBOARD: IconName = IconName(\"artboard\")\n ARTBOARD_FILLED: IconName = IconName(\"artboard-filled\")\n ARTBOARD_OFF: IconName = IconName(\"artboard-off\")\n ARTICLE: IconName = IconName(\"article\")\n ARTICLE_FILLED_FILLED: IconName = IconName(\"article-filled-filled\")\n ARTICLE_OFF: IconName = IconName(\"article-off\")\n ASPECT_RATIO: IconName = IconName(\"aspect-ratio\")\n ASPECT_RATIO_FILLED: IconName = IconName(\"aspect-ratio-filled\")\n ASPECT_RATIO_OFF: IconName = IconName(\"aspect-ratio-off\")\n ASSEMBLY: IconName = IconName(\"assembly\")\n ASSEMBLY_OFF: IconName = IconName(\"assembly-off\")\n ASSET: IconName = IconName(\"asset\")\n ASTERISK: IconName = IconName(\"asterisk\")\n ASTERISK_SIMPLE: IconName = IconName(\"asterisk-simple\")\n AT: IconName = IconName(\"at\")\n AT_OFF: IconName = IconName(\"at-off\")\n ATOM: IconName = IconName(\"atom\")\n ATOM_2: IconName = IconName(\"atom-2\")\n ATOM_2_FILLED: IconName = IconName(\"atom-2-filled\")\n ATOM_OFF: IconName = IconName(\"atom-off\")\n AUGMENTED_REALITY: IconName = IconName(\"augmented-reality\")\n AUGMENTED_REALITY_2: IconName = IconName(\"augmented-reality-2\")\n AUGMENTED_REALITY_OFF: IconName = IconName(\"augmented-reality-off\")\n AWARD: IconName = IconName(\"award\")\n AWARD_FILLED: IconName = IconName(\"award-filled\")\n AWARD_OFF: IconName = IconName(\"award-off\")\n AXE: IconName = IconName(\"axe\")\n AXIS_X: IconName = IconName(\"axis-x\")\n AXIS_Y: IconName = IconName(\"axis-y\")\n BABY_BOTTLE: IconName = IconName(\"baby-bottle\")\n BABY_CARRIAGE: IconName = IconName(\"baby-carriage\")\n BACKHOE: IconName = IconName(\"backhoe\")\n BACKPACK: IconName = IconName(\"backpack\")\n BACKPACK_OFF: IconName = IconName(\"backpack-off\")\n BACKSLASH: IconName = IconName(\"backslash\")\n BACKSPACE: IconName = IconName(\"backspace\")\n BACKSPACE_FILLED: IconName = IconName(\"backspace-filled\")\n BADGE: IconName = IconName(\"badge\")\n BADGE_3D: IconName = IconName(\"badge-3d\")\n BADGE_4K: IconName = IconName(\"badge-4k\")\n BADGE_8K: IconName = IconName(\"badge-8k\")\n BADGE_AD: IconName = IconName(\"badge-ad\")\n BADGE_AR: IconName = IconName(\"badge-ar\")\n BADGE_CC: IconName = IconName(\"badge-cc\")\n BADGE_FILLED: IconName = IconName(\"badge-filled\")\n BADGE_HD: IconName = IconName(\"badge-hd\")\n BADGE_OFF: IconName = IconName(\"badge-off\")\n BADGE_SD: IconName = IconName(\"badge-sd\")\n BADGE_TM: IconName = IconName(\"badge-tm\")\n BADGE_VO: IconName = IconName(\"badge-vo\")\n BADGE_VR: IconName = IconName(\"badge-vr\")\n BADGE_WC: IconName = IconName(\"badge-wc\")\n BADGES: IconName = IconName(\"badges\")\n BADGES_FILLED: IconName = IconName(\"badges-filled\")\n BADGES_OFF: IconName = IconName(\"badges-off\")\n BAGUETTE: IconName = IconName(\"baguette\")\n BALL_AMERICAN_FOOTBALL: IconName = IconName(\"ball-american-football\")\n BALL_AMERICAN_FOOTBALL_OFF: IconName = IconName(\"ball-american-football-off\")\n BALL_BASEBALL: IconName = IconName(\"ball-baseball\")\n BALL_BASKETBALL: IconName = IconName(\"ball-basketball\")\n BALL_BOWLING: IconName = IconName(\"ball-bowling\")\n BALL_FOOTBALL: IconName = IconName(\"ball-football\")\n BALL_FOOTBALL_OFF: IconName = IconName(\"ball-football-off\")\n BALL_TENNIS: IconName = IconName(\"ball-tennis\")\n BALL_VOLLEYBALL: IconName = IconName(\"ball-volleyball\")\n BALLOON: IconName = IconName(\"balloon\")\n BALLOON_FILLED: IconName = IconName(\"balloon-filled\")\n BALLOON_OFF: IconName = IconName(\"balloon-off\")\n BALLPEN: IconName = IconName(\"ballpen\")\n BALLPEN_FILLED: IconName = IconName(\"ballpen-filled\")\n BALLPEN_OFF: IconName = IconName(\"ballpen-off\")\n BAN: IconName = IconName(\"ban\")\n BANDAGE: IconName = IconName(\"bandage\")\n BANDAGE_FILLED: IconName = IconName(\"bandage-filled\")\n BANDAGE_OFF: IconName = IconName(\"bandage-off\")\n BARBELL: IconName = IconName(\"barbell\")\n BARBELL_OFF: IconName = IconName(\"barbell-off\")\n BARCODE: IconName = IconName(\"barcode\")\n BARCODE_OFF: IconName = IconName(\"barcode-off\")\n BARREL: IconName = IconName(\"barrel\")\n BARREL_OFF: IconName = IconName(\"barrel-off\")\n BARRIER_BLOCK: IconName = IconName(\"barrier-block\")\n BARRIER_BLOCK_OFF: IconName = IconName(\"barrier-block-off\")\n BASELINE: IconName = IconName(\"baseline\")\n BASELINE_DENSITY_LARGE: IconName = IconName(\"baseline-density-large\")\n BASELINE_DENSITY_MEDIUM: IconName = IconName(\"baseline-density-medium\")\n BASELINE_DENSITY_SMALL: IconName = IconName(\"baseline-density-small\")\n BASKET: IconName = IconName(\"basket\")\n BASKET_FILLED: IconName = IconName(\"basket-filled\")\n BASKET_OFF: IconName = IconName(\"basket-off\")\n BAT: IconName = IconName(\"bat\")\n BATH: IconName = IconName(\"bath\")\n BATH_FILLED: IconName = IconName(\"bath-filled\")\n BATH_OFF: IconName = IconName(\"bath-off\")\n BATTERY: IconName = IconName(\"battery\")\n BATTERY_1: IconName = IconName(\"battery-1\")\n BATTERY_1_FILLED: IconName = IconName(\"battery-1-filled\")\n BATTERY_2: IconName = IconName(\"battery-2\")\n BATTERY_2_FILLED: IconName = IconName(\"battery-2-filled\")\n BATTERY_3: IconName = IconName(\"battery-3\")\n BATTERY_3_FILLED: IconName = IconName(\"battery-3-filled\")\n BATTERY_4: IconName = IconName(\"battery-4\")\n BATTERY_4_FILLED: IconName = IconName(\"battery-4-filled\")\n BATTERY_AUTOMOTIVE: IconName = IconName(\"battery-automotive\")\n BATTERY_CHARGING: IconName = IconName(\"battery-charging\")\n BATTERY_CHARGING_2: IconName = IconName(\"battery-charging-2\")\n BATTERY_ECO: IconName = IconName(\"battery-eco\")\n BATTERY_FILLED: IconName = IconName(\"battery-filled\")\n BATTERY_OFF: IconName = IconName(\"battery-off\")\n BEACH: IconName = IconName(\"beach\")\n BEACH_OFF: IconName = IconName(\"beach-off\")\n BED: IconName = IconName(\"bed\")\n BED_FILLED: IconName = IconName(\"bed-filled\")\n BED_OFF: IconName = IconName(\"bed-off\")\n BEER: IconName = IconName(\"beer\")\n BEER_FILLED: IconName = IconName(\"beer-filled\")\n BEER_OFF: IconName = IconName(\"beer-off\")\n BELL: IconName = IconName(\"bell\")\n BELL_BOLT: IconName = IconName(\"bell-bolt\")\n BELL_CANCEL: IconName = IconName(\"bell-cancel\")\n BELL_CHECK: IconName = IconName(\"bell-check\")\n BELL_CODE: IconName = IconName(\"bell-code\")\n BELL_COG: IconName = IconName(\"bell-cog\")\n BELL_DOLLAR: IconName = IconName(\"bell-dollar\")\n BELL_DOWN: IconName = IconName(\"bell-down\")\n BELL_EXCLAMATION: IconName = IconName(\"bell-exclamation\")\n BELL_FILLED: IconName = IconName(\"bell-filled\")\n BELL_HEART: IconName = IconName(\"bell-heart\")\n BELL_MINUS: IconName = IconName(\"bell-minus\")\n BELL_MINUS_FILLED: IconName = IconName(\"bell-minus-filled\")\n BELL_OFF: IconName = IconName(\"bell-off\")\n BELL_PAUSE: IconName = IconName(\"bell-pause\")\n BELL_PIN: IconName = IconName(\"bell-pin\")\n BELL_PLUS: IconName = IconName(\"bell-plus\")\n BELL_PLUS_FILLED: IconName = IconName(\"bell-plus-filled\")\n BELL_QUESTION: IconName = IconName(\"bell-question\")\n BELL_RINGING: IconName = IconName(\"bell-ringing\")\n BELL_RINGING_2: IconName = IconName(\"bell-ringing-2\")\n BELL_RINGING_2_FILLED: IconName = IconName(\"bell-ringing-2-filled\")\n BELL_RINGING_FILLED: IconName = IconName(\"bell-ringing-filled\")\n BELL_SCHOOL: IconName = IconName(\"bell-school\")\n BELL_SEARCH: IconName = IconName(\"bell-search\")\n BELL_SHARE: IconName = IconName(\"bell-share\")\n BELL_STAR: IconName = IconName(\"bell-star\")\n BELL_UP: IconName = IconName(\"bell-up\")\n BELL_X: IconName = IconName(\"bell-x\")\n BELL_X_FILLED: IconName = IconName(\"bell-x-filled\")\n BELL_Z: IconName = IconName(\"bell-z\")\n BELL_Z_FILLED: IconName = IconName(\"bell-z-filled\")\n BETA: IconName = IconName(\"beta\")\n BIBLE: IconName = IconName(\"bible\")\n BIKE: IconName = IconName(\"bike\")\n BIKE_OFF: IconName = IconName(\"bike-off\")\n BINARY: IconName = IconName(\"binary\")\n BINARY_OFF: IconName = IconName(\"binary-off\")\n BINARY_TREE: IconName = IconName(\"binary-tree\")\n BINARY_TREE_2: IconName = IconName(\"binary-tree-2\")\n BIOHAZARD: IconName = IconName(\"biohazard\")\n BIOHAZARD_OFF: IconName = IconName(\"biohazard-off\")\n BLADE: IconName = IconName(\"blade\")\n BLADE_FILLED: IconName = IconName(\"blade-filled\")\n BLEACH: IconName = IconName(\"bleach\")\n BLEACH_CHLORINE: IconName = IconName(\"bleach-chlorine\")\n BLEACH_NO_CHLORINE: IconName = IconName(\"bleach-no-chlorine\")\n BLEACH_OFF: IconName = IconName(\"bleach-off\")\n BLOCKQUOTE: IconName = IconName(\"blockquote\")\n BLUETOOTH: IconName = IconName(\"bluetooth\")\n BLUETOOTH_CONNECTED: IconName = IconName(\"bluetooth-connected\")\n BLUETOOTH_OFF: IconName = IconName(\"bluetooth-off\")\n BLUETOOTH_X: IconName = IconName(\"bluetooth-x\")\n BLUR: IconName = IconName(\"blur\")\n BLUR_OFF: IconName = IconName(\"blur-off\")\n BMP: IconName = IconName(\"bmp\")\n BOLD: IconName = IconName(\"bold\")\n BOLD_OFF: IconName = IconName(\"bold-off\")\n BOLT: IconName = IconName(\"bolt\")\n BOLT_OFF: IconName = IconName(\"bolt-off\")\n BOMB: IconName = IconName(\"bomb\")\n BOMB_FILLED: IconName = IconName(\"bomb-filled\")\n BONE: IconName = IconName(\"bone\")\n BONE_OFF: IconName = IconName(\"bone-off\")\n BONG: IconName = IconName(\"bong\")\n BONG_OFF: IconName = IconName(\"bong-off\")\n BOOK: IconName = IconName(\"book\")\n BOOK_2: IconName = IconName(\"book-2\")\n BOOK_DOWNLOAD: IconName = IconName(\"book-download\")\n BOOK_FILLED: IconName = IconName(\"book-filled\")\n BOOK_OFF: IconName = IconName(\"book-off\")\n BOOK_UPLOAD: IconName = IconName(\"book-upload\")\n BOOKMARK: IconName = IconName(\"bookmark\")\n BOOKMARK_EDIT: IconName = IconName(\"bookmark-edit\")\n BOOKMARK_FILLED: IconName = IconName(\"bookmark-filled\")\n BOOKMARK_MINUS: IconName = IconName(\"bookmark-minus\")\n BOOKMARK_OFF: IconName = IconName(\"bookmark-off\")\n BOOKMARK_PLUS: IconName = IconName(\"bookmark-plus\")\n BOOKMARK_QUESTION: IconName = IconName(\"bookmark-question\")\n BOOKMARKS: IconName = IconName(\"bookmarks\")\n BOOKMARKS_OFF: IconName = IconName(\"bookmarks-off\")\n BOOKS: IconName = IconName(\"books\")\n BOOKS_OFF: IconName = IconName(\"books-off\")\n BORDER_ALL: IconName = IconName(\"border-all\")\n BORDER_BOTTOM: IconName = IconName(\"border-bottom\")\n BORDER_CORNERS: IconName = IconName(\"border-corners\")\n BORDER_HORIZONTAL: IconName = IconName(\"border-horizontal\")\n BORDER_INNER: IconName = IconName(\"border-inner\")\n BORDER_LEFT: IconName = IconName(\"border-left\")\n BORDER_NONE: IconName = IconName(\"border-none\")\n BORDER_OUTER: IconName = IconName(\"border-outer\")\n BORDER_RADIUS: IconName = IconName(\"border-radius\")\n BORDER_RIGHT: IconName = IconName(\"border-right\")\n BORDER_SIDES: IconName = IconName(\"border-sides\")\n BORDER_STYLE: IconName = IconName(\"border-style\")\n BORDER_STYLE_2: IconName = IconName(\"border-style-2\")\n BORDER_TOP: IconName = IconName(\"border-top\")\n BORDER_VERTICAL: IconName = IconName(\"border-vertical\")\n BOTTLE: IconName = IconName(\"bottle\")\n BOTTLE_FILLED: IconName = IconName(\"bottle-filled\")\n BOTTLE_OFF: IconName = IconName(\"bottle-off\")\n BOUNCE_LEFT: IconName = IconName(\"bounce-left\")\n BOUNCE_RIGHT: IconName = IconName(\"bounce-right\")\n BOW: IconName = IconName(\"bow\")\n BOWL: IconName = IconName(\"bowl\")\n BOX: IconName = IconName(\"box\")\n BOX_ALIGN_BOTTOM: IconName = IconName(\"box-align-bottom\")\n BOX_ALIGN_BOTTOM_FILLED: IconName = IconName(\"box-align-bottom-filled\")\n BOX_ALIGN_BOTTOM_LEFT: IconName = IconName(\"box-align-bottom-left\")\n BOX_ALIGN_BOTTOM_LEFT_FILLED: IconName = IconName(\"box-align-bottom-left-filled\")\n BOX_ALIGN_BOTTOM_RIGHT: IconName = IconName(\"box-align-bottom-right\")\n BOX_ALIGN_BOTTOM_RIGHT_FILLED: IconName = IconName(\"box-align-bottom-right-filled\")\n BOX_ALIGN_LEFT: IconName = IconName(\"box-align-left\")\n BOX_ALIGN_LEFT_FILLED: IconName = IconName(\"box-align-left-filled\")\n BOX_ALIGN_RIGHT: IconName = IconName(\"box-align-right\")\n BOX_ALIGN_RIGHT_FILLED: IconName = IconName(\"box-align-right-filled\")\n BOX_ALIGN_TOP: IconName = IconName(\"box-align-top\")\n BOX_ALIGN_TOP_FILLED: IconName = IconName(\"box-align-top-filled\")\n BOX_ALIGN_TOP_LEFT: IconName = IconName(\"box-align-top-left\")\n BOX_ALIGN_TOP_LEFT_FILLED: IconName = IconName(\"box-align-top-left-filled\")\n BOX_ALIGN_TOP_RIGHT: IconName = IconName(\"box-align-top-right\")\n BOX_ALIGN_TOP_RIGHT_FILLED: IconName = IconName(\"box-align-top-right-filled\")\n BOX_MARGIN: IconName = IconName(\"box-margin\")\n BOX_MODEL: IconName = IconName(\"box-model\")\n BOX_MODEL_2: IconName = IconName(\"box-model-2\")\n BOX_MODEL_2_OFF: IconName = IconName(\"box-model-2-off\")\n BOX_MODEL_OFF: IconName = IconName(\"box-model-off\")\n BOX_MULTIPLE: IconName = IconName(\"box-multiple\")\n BOX_MULTIPLE_0: IconName = IconName(\"box-multiple-0\")\n BOX_MULTIPLE_1: IconName = IconName(\"box-multiple-1\")\n BOX_MULTIPLE_2: IconName = IconName(\"box-multiple-2\")\n BOX_MULTIPLE_3: IconName = IconName(\"box-multiple-3\")\n BOX_MULTIPLE_4: IconName = IconName(\"box-multiple-4\")\n BOX_MULTIPLE_5: IconName = IconName(\"box-multiple-5\")\n BOX_MULTIPLE_6: IconName = IconName(\"box-multiple-6\")\n BOX_MULTIPLE_7: IconName = IconName(\"box-multiple-7\")\n BOX_MULTIPLE_8: IconName = IconName(\"box-multiple-8\")\n BOX_MULTIPLE_9: IconName = IconName(\"box-multiple-9\")\n BOX_OFF: IconName = IconName(\"box-off\")\n BOX_PADDING: IconName = IconName(\"box-padding\")\n BOX_SEAM: IconName = IconName(\"box-seam\")\n BRACES: IconName = IconName(\"braces\")\n BRACES_OFF: IconName = IconName(\"braces-off\")\n BRACKETS: IconName = IconName(\"brackets\")\n BRACKETS_CONTAIN: IconName = IconName(\"brackets-contain\")\n BRACKETS_CONTAIN_END: IconName = IconName(\"brackets-contain-end\")\n BRACKETS_CONTAIN_START: IconName = IconName(\"brackets-contain-start\")\n BRACKETS_OFF: IconName = IconName(\"brackets-off\")\n BRAILLE: IconName = IconName(\"braille\")\n BRAIN: IconName = IconName(\"brain\")\n BRAND_4CHAN: IconName = IconName(\"brand-4chan\")\n BRAND_ABSTRACT: IconName = IconName(\"brand-abstract\")\n BRAND_ADOBE: IconName = IconName(\"brand-adobe\")\n BRAND_ADONIS_JS: IconName = IconName(\"brand-adonis-js\")\n BRAND_AIRBNB: IconName = IconName(\"brand-airbnb\")\n BRAND_AIRTABLE: IconName = IconName(\"brand-airtable\")\n BRAND_ALGOLIA: IconName = IconName(\"brand-algolia\")\n BRAND_ALIPAY: IconName = IconName(\"brand-alipay\")\n BRAND_ALPINE_JS: IconName = IconName(\"brand-alpine-js\")\n BRAND_AMAZON: IconName = IconName(\"brand-amazon\")\n BRAND_AMD: IconName = IconName(\"brand-amd\")\n BRAND_AMIGO: IconName = IconName(\"brand-amigo\")\n BRAND_AMONG_US: IconName = IconName(\"brand-among-us\")\n BRAND_ANDROID: IconName = IconName(\"brand-android\")\n BRAND_ANGULAR: IconName = IconName(\"brand-angular\")\n BRAND_ANSIBLE: IconName = IconName(\"brand-ansible\")\n BRAND_AO3: IconName = IconName(\"brand-ao3\")\n BRAND_APPGALLERY: IconName = IconName(\"brand-appgallery\")\n BRAND_APPLE: IconName = IconName(\"brand-apple\")\n BRAND_APPLE_ARCADE: IconName = IconName(\"brand-apple-arcade\")\n BRAND_APPLE_PODCAST: IconName = IconName(\"brand-apple-podcast\")\n BRAND_APPSTORE: IconName = IconName(\"brand-appstore\")\n BRAND_ASANA: IconName = IconName(\"brand-asana\")\n BRAND_AWS: IconName = IconName(\"brand-aws\")\n BRAND_AZURE: IconName = IconName(\"brand-azure\")\n BRAND_BACKBONE: IconName = IconName(\"brand-backbone\")\n BRAND_BADOO: IconName = IconName(\"brand-badoo\")\n BRAND_BAIDU: IconName = IconName(\"brand-baidu\")\n BRAND_BANDCAMP: IconName = IconName(\"brand-bandcamp\")\n BRAND_BANDLAB: IconName = IconName(\"brand-bandlab\")\n BRAND_BEATS: IconName = IconName(\"brand-beats\")\n BRAND_BEHANCE: IconName = IconName(\"brand-behance\")\n BRAND_BILIBILI: IconName = IconName(\"brand-bilibili\")\n BRAND_BINANCE: IconName = IconName(\"brand-binance\")\n BRAND_BING: IconName = IconName(\"brand-bing\")\n BRAND_BITBUCKET: IconName = IconName(\"brand-bitbucket\")\n BRAND_BLACKBERRY: IconName = IconName(\"brand-blackberry\")\n BRAND_BLENDER: IconName = IconName(\"brand-blender\")\n BRAND_BLOGGER: IconName = IconName(\"brand-blogger\")\n BRAND_BOOKING: IconName = IconName(\"brand-booking\")\n BRAND_BOOTSTRAP: IconName = IconName(\"brand-bootstrap\")\n BRAND_BULMA: IconName = IconName(\"brand-bulma\")\n BRAND_BUMBLE: IconName = IconName(\"brand-bumble\")\n BRAND_BUNPO: IconName = IconName(\"brand-bunpo\")\n BRAND_C_SHARP: IconName = IconName(\"brand-c-sharp\")\n BRAND_CAKE: IconName = IconName(\"brand-cake\")\n BRAND_CAKEPHP: IconName = IconName(\"brand-cakephp\")\n BRAND_CAMPAIGNMONITOR: IconName = IconName(\"brand-campaignmonitor\")\n BRAND_CARBON: IconName = IconName(\"brand-carbon\")\n BRAND_CASHAPP: IconName = IconName(\"brand-cashapp\")\n BRAND_CHROME: IconName = IconName(\"brand-chrome\")\n BRAND_CINEMA_4D: IconName = IconName(\"brand-cinema-4d\")\n BRAND_CITYMAPPER: IconName = IconName(\"brand-citymapper\")\n BRAND_CLOUDFLARE: IconName = IconName(\"brand-cloudflare\")\n BRAND_CODECOV: IconName = IconName(\"brand-codecov\")\n BRAND_CODEPEN: IconName = IconName(\"brand-codepen\")\n BRAND_CODESANDBOX: IconName = IconName(\"brand-codesandbox\")\n BRAND_COHOST: IconName = IconName(\"brand-cohost\")\n BRAND_COINBASE: IconName = IconName(\"brand-coinbase\")\n BRAND_COMEDY_CENTRAL: IconName = IconName(\"brand-comedy-central\")\n BRAND_COREOS: IconName = IconName(\"brand-coreos\")\n BRAND_COUCHDB: IconName = IconName(\"brand-couchdb\")\n BRAND_COUCHSURFING: IconName = IconName(\"brand-couchsurfing\")\n BRAND_CPP: IconName = IconName(\"brand-cpp\")\n BRAND_CRAFT: IconName = IconName(\"brand-craft\")\n BRAND_CRUNCHBASE: IconName = IconName(\"brand-crunchbase\")\n BRAND_CSS3: IconName = IconName(\"brand-css3\")\n BRAND_CTEMPLAR: IconName = IconName(\"brand-ctemplar\")\n BRAND_CUCUMBER: IconName = IconName(\"brand-cucumber\")\n BRAND_CUPRA: IconName = IconName(\"brand-cupra\")\n BRAND_CYPRESS: IconName = IconName(\"brand-cypress\")\n BRAND_D3: IconName = IconName(\"brand-d3\")\n BRAND_DAYS_COUNTER: IconName = IconName(\"brand-days-counter\")\n BRAND_DCOS: IconName = IconName(\"brand-dcos\")\n BRAND_DEBIAN: IconName = IconName(\"brand-debian\")\n BRAND_DEEZER: IconName = IconName(\"brand-deezer\")\n BRAND_DELIVEROO: IconName = IconName(\"brand-deliveroo\")\n BRAND_DENO: IconName = IconName(\"brand-deno\")\n BRAND_DENODO: IconName = IconName(\"brand-denodo\")\n BRAND_DEVIANTART: IconName = IconName(\"brand-deviantart\")\n BRAND_DIGG: IconName = IconName(\"brand-digg\")\n BRAND_DINGTALK: IconName = IconName(\"brand-dingtalk\")\n BRAND_DISCORD: IconName = IconName(\"brand-discord\")\n BRAND_DISCORD_FILLED: IconName = IconName(\"brand-discord-filled\")\n BRAND_DISNEY: IconName = IconName(\"brand-disney\")\n BRAND_DISQUS: IconName = IconName(\"brand-disqus\")\n BRAND_DJANGO: IconName = IconName(\"brand-django\")\n BRAND_DOCKER: IconName = IconName(\"brand-docker\")\n BRAND_DOCTRINE: IconName = IconName(\"brand-doctrine\")\n BRAND_DOLBY_DIGITAL: IconName = IconName(\"brand-dolby-digital\")\n BRAND_DOUBAN: IconName = IconName(\"brand-douban\")\n BRAND_DRIBBBLE: IconName = IconName(\"brand-dribbble\")\n BRAND_DRIBBBLE_FILLED: IconName = IconName(\"brand-dribbble-filled\")\n BRAND_DROPS: IconName = IconName(\"brand-drops\")\n BRAND_DRUPAL: IconName = IconName(\"brand-drupal\")\n BRAND_EDGE: IconName = IconName(\"brand-edge\")\n BRAND_ELASTIC: IconName = IconName(\"brand-elastic\")\n BRAND_ELECTRONIC_ARTS: IconName = IconName(\"brand-electronic-arts\")\n BRAND_EMBER: IconName = IconName(\"brand-ember\")\n BRAND_ENVATO: IconName = IconName(\"brand-envato\")\n BRAND_ETSY: IconName = IconName(\"brand-etsy\")\n BRAND_EVERNOTE: IconName = IconName(\"brand-evernote\")\n BRAND_FACEBOOK: IconName = IconName(\"brand-facebook\")\n BRAND_FACEBOOK_FILLED: IconName = IconName(\"brand-facebook-filled\")\n BRAND_FEEDLY: IconName = IconName(\"brand-feedly\")\n BRAND_FIGMA: IconName = IconName(\"brand-figma\")\n BRAND_FILEZILLA: IconName = IconName(\"brand-filezilla\")\n BRAND_FINDER: IconName = IconName(\"brand-finder\")\n BRAND_FIREBASE: IconName = IconName(\"brand-firebase\")\n BRAND_FIREFOX: IconName = IconName(\"brand-firefox\")\n BRAND_FIVERR: IconName = IconName(\"brand-fiverr\")\n BRAND_FLICKR: IconName = IconName(\"brand-flickr\")\n BRAND_FLIGHTRADAR24: IconName = IconName(\"brand-flightradar24\")\n BRAND_FLIPBOARD: IconName = IconName(\"brand-flipboard\")\n BRAND_FLUTTER: IconName = IconName(\"brand-flutter\")\n BRAND_FORTNITE: IconName = IconName(\"brand-fortnite\")\n BRAND_FOURSQUARE: IconName = IconName(\"brand-foursquare\")\n BRAND_FRAMER: IconName = IconName(\"brand-framer\")\n BRAND_FRAMER_MOTION: IconName = IconName(\"brand-framer-motion\")\n BRAND_FUNIMATION: IconName = IconName(\"brand-funimation\")\n BRAND_GATSBY: IconName = IconName(\"brand-gatsby\")\n BRAND_GIT: IconName = IconName(\"brand-git\")\n BRAND_GITHUB: IconName = IconName(\"brand-github\")\n BRAND_GITHUB_COPILOT: IconName = IconName(\"brand-github-copilot\")\n BRAND_GITHUB_FILLED: IconName = IconName(\"brand-github-filled\")\n BRAND_GITLAB: IconName = IconName(\"brand-gitlab\")\n BRAND_GMAIL: IconName = IconName(\"brand-gmail\")\n BRAND_GOLANG: IconName = IconName(\"brand-golang\")\n BRAND_GOOGLE: IconName = IconName(\"brand-google\")\n BRAND_GOOGLE_ANALYTICS: IconName = IconName(\"brand-google-analytics\")\n BRAND_GOOGLE_BIG_QUERY: IconName = IconName(\"brand-google-big-query\")\n BRAND_GOOGLE_DRIVE: IconName = IconName(\"brand-google-drive\")\n BRAND_GOOGLE_FIT: IconName = IconName(\"brand-google-fit\")\n BRAND_GOOGLE_HOME: IconName = IconName(\"brand-google-home\")\n BRAND_GOOGLE_MAPS: IconName = IconName(\"brand-google-maps\")\n BRAND_GOOGLE_ONE: IconName = IconName(\"brand-google-one\")\n BRAND_GOOGLE_PHOTOS: IconName = IconName(\"brand-google-photos\")\n BRAND_GOOGLE_PLAY: IconName = IconName(\"brand-google-play\")\n BRAND_GOOGLE_PODCASTS: IconName = IconName(\"brand-google-podcasts\")\n BRAND_GRAMMARLY: IconName = IconName(\"brand-grammarly\")\n BRAND_GRAPHQL: IconName = IconName(\"brand-graphql\")\n BRAND_GRAVATAR: IconName = IconName(\"brand-gravatar\")\n BRAND_GRINDR: IconName = IconName(\"brand-grindr\")\n BRAND_GUARDIAN: IconName = IconName(\"brand-guardian\")\n BRAND_GUMROAD: IconName = IconName(\"brand-gumroad\")\n BRAND_HBO: IconName = IconName(\"brand-hbo\")\n BRAND_HEADLESSUI: IconName = IconName(\"brand-headlessui\")\n BRAND_HEXO: IconName = IconName(\"brand-hexo\")\n BRAND_HIPCHAT: IconName = IconName(\"brand-hipchat\")\n BRAND_HTML5: IconName = IconName(\"brand-html5\")\n BRAND_INERTIA: IconName = IconName(\"brand-inertia\")\n BRAND_INSTAGRAM: IconName = IconName(\"brand-instagram\")\n BRAND_INTERCOM: IconName = IconName(\"brand-intercom\")\n BRAND_ITCH: IconName = IconName(\"brand-itch\")\n BRAND_JAVASCRIPT: IconName = IconName(\"brand-javascript\")\n BRAND_JUEJIN: IconName = IconName(\"brand-juejin\")\n BRAND_KBIN: IconName = IconName(\"brand-kbin\")\n BRAND_KICK: IconName = IconName(\"brand-kick\")\n BRAND_KICKSTARTER: IconName = IconName(\"brand-kickstarter\")\n BRAND_KOTLIN: IconName = IconName(\"brand-kotlin\")\n BRAND_LARAVEL: IconName = IconName(\"brand-laravel\")\n BRAND_LASTFM: IconName = IconName(\"brand-lastfm\")\n BRAND_LEETCODE: IconName = IconName(\"brand-leetcode\")\n BRAND_LETTERBOXD: IconName = IconName(\"brand-letterboxd\")\n BRAND_LINE: IconName = IconName(\"brand-line\")\n BRAND_LINKEDIN: IconName = IconName(\"brand-linkedin\")\n BRAND_LINKTREE: IconName = IconName(\"brand-linktree\")\n BRAND_LINQPAD: IconName = IconName(\"brand-linqpad\")\n BRAND_LOOM: IconName = IconName(\"brand-loom\")\n BRAND_MAILGUN: IconName = IconName(\"brand-mailgun\")\n BRAND_MANTINE: IconName = IconName(\"brand-mantine\")\n BRAND_MASTERCARD: IconName = IconName(\"brand-mastercard\")\n BRAND_MASTODON: IconName = IconName(\"brand-mastodon\")\n BRAND_MATRIX: IconName = IconName(\"brand-matrix\")\n BRAND_MCDONALDS: IconName = IconName(\"brand-mcdonalds\")\n BRAND_MEDIUM: IconName = IconName(\"brand-medium\")\n BRAND_MERCEDES: IconName = IconName(\"brand-mercedes\")\n BRAND_MESSENGER: IconName = IconName(\"brand-messenger\")\n BRAND_META: IconName = IconName(\"brand-meta\")\n BRAND_MICROSOFT_TEAMS: IconName = IconName(\"brand-microsoft-teams\")\n BRAND_MINECRAFT: IconName = IconName(\"brand-minecraft\")\n BRAND_MINIPROGRAM: IconName = IconName(\"brand-miniprogram\")\n BRAND_MIXPANEL: IconName = IconName(\"brand-mixpanel\")\n BRAND_MONDAY: IconName = IconName(\"brand-monday\")\n BRAND_MONGODB: IconName = IconName(\"brand-mongodb\")\n BRAND_MY_OPPO: IconName = IconName(\"brand-my-oppo\")\n BRAND_MYSQL: IconName = IconName(\"brand-mysql\")\n BRAND_NATIONAL_GEOGRAPHIC: IconName = IconName(\"brand-national-geographic\")\n BRAND_NEM: IconName = IconName(\"brand-nem\")\n BRAND_NETBEANS: IconName = IconName(\"brand-netbeans\")\n BRAND_NETEASE_MUSIC: IconName = IconName(\"brand-netease-music\")\n BRAND_NETFLIX: IconName = IconName(\"brand-netflix\")\n BRAND_NEXO: IconName = IconName(\"brand-nexo\")\n BRAND_NEXTCLOUD: IconName = IconName(\"brand-nextcloud\")\n BRAND_NEXTJS: IconName = IconName(\"brand-nextjs\")\n BRAND_NODEJS: IconName = IconName(\"brand-nodejs\")\n BRAND_NORD_VPN: IconName = IconName(\"brand-nord-vpn\")\n BRAND_NOTION: IconName = IconName(\"brand-notion\")\n BRAND_NPM: IconName = IconName(\"brand-npm\")\n BRAND_NUXT: IconName = IconName(\"brand-nuxt\")\n BRAND_NYTIMES: IconName = IconName(\"brand-nytimes\")\n BRAND_OAUTH: IconName = IconName(\"brand-oauth\")\n BRAND_OFFICE: IconName = IconName(\"brand-office\")\n BRAND_OK_RU: IconName = IconName(\"brand-ok-ru\")\n BRAND_ONEDRIVE: IconName = IconName(\"brand-onedrive\")\n BRAND_ONLYFANS: IconName = IconName(\"brand-onlyfans\")\n BRAND_OPEN_SOURCE: IconName = IconName(\"brand-open-source\")\n BRAND_OPENAI: IconName = IconName(\"brand-openai\")\n BRAND_OPENVPN: IconName = IconName(\"brand-openvpn\")\n BRAND_OPERA: IconName = IconName(\"brand-opera\")\n BRAND_PAGEKIT: IconName = IconName(\"brand-pagekit\")\n BRAND_PATREON: IconName = IconName(\"brand-patreon\")\n BRAND_PAYPAL: IconName = IconName(\"brand-paypal\")\n BRAND_PAYPAL_FILLED: IconName = IconName(\"brand-paypal-filled\")\n BRAND_PAYPAY: IconName = IconName(\"brand-paypay\")\n BRAND_PEANUT: IconName = IconName(\"brand-peanut\")\n BRAND_PEPSI: IconName = IconName(\"brand-pepsi\")\n BRAND_PHP: IconName = IconName(\"brand-php\")\n BRAND_PICSART: IconName = IconName(\"brand-picsart\")\n BRAND_PINTEREST: IconName = IconName(\"brand-pinterest\")\n BRAND_PLANETSCALE: IconName = IconName(\"brand-planetscale\")\n BRAND_POCKET: IconName = IconName(\"brand-pocket\")\n BRAND_POLYMER: IconName = IconName(\"brand-polymer\")\n BRAND_POWERSHELL: IconName = IconName(\"brand-powershell\")\n BRAND_PRISMA: IconName = IconName(\"brand-prisma\")\n BRAND_PRODUCTHUNT: IconName = IconName(\"brand-producthunt\")\n BRAND_PUSHBULLET: IconName = IconName(\"brand-pushbullet\")\n BRAND_PUSHOVER: IconName = IconName(\"brand-pushover\")\n BRAND_PYTHON: IconName = IconName(\"brand-python\")\n BRAND_QQ: IconName = IconName(\"brand-qq\")\n BRAND_RADIX_UI: IconName = IconName(\"brand-radix-ui\")\n BRAND_REACT: IconName = IconName(\"brand-react\")\n BRAND_REACT_NATIVE: IconName = IconName(\"brand-react-native\")\n BRAND_REASON: IconName = IconName(\"brand-reason\")\n BRAND_REDDIT: IconName = IconName(\"brand-reddit\")\n BRAND_REDHAT: IconName = IconName(\"brand-redhat\")\n BRAND_REDUX: IconName = IconName(\"brand-redux\")\n BRAND_REVOLUT: IconName = IconName(\"brand-revolut\")\n BRAND_RUMBLE: IconName = IconName(\"brand-rumble\")\n BRAND_RUST: IconName = IconName(\"brand-rust\")\n BRAND_SAFARI: IconName = IconName(\"brand-safari\")\n BRAND_SAMSUNGPASS: IconName = IconName(\"brand-samsungpass\")\n BRAND_SASS: IconName = IconName(\"brand-sass\")\n BRAND_SENTRY: IconName = IconName(\"brand-sentry\")\n BRAND_SHARIK: IconName = IconName(\"brand-sharik\")\n BRAND_SHAZAM: IconName = IconName(\"brand-shazam\")\n BRAND_SHOPEE: IconName = IconName(\"brand-shopee\")\n BRAND_SKETCH: IconName = IconName(\"brand-sketch\")\n BRAND_SKYPE: IconName = IconName(\"brand-skype\")\n BRAND_SLACK: IconName = IconName(\"brand-slack\")\n BRAND_SNAPCHAT: IconName = IconName(\"brand-snapchat\")\n BRAND_SNAPSEED: IconName = IconName(\"brand-snapseed\")\n BRAND_SNOWFLAKE: IconName = IconName(\"brand-snowflake\")\n BRAND_SOCKET_IO: IconName = IconName(\"brand-socket-io\")\n BRAND_SOLIDJS: IconName = IconName(\"brand-solidjs\")\n BRAND_SOUNDCLOUD: IconName = IconName(\"brand-soundcloud\")\n BRAND_SPACEHEY: IconName = IconName(\"brand-spacehey\")\n BRAND_SPEEDTEST: IconName = IconName(\"brand-speedtest\")\n BRAND_SPOTIFY: IconName = IconName(\"brand-spotify\")\n BRAND_STACKOVERFLOW: IconName = IconName(\"brand-stackoverflow\")\n BRAND_STACKSHARE: IconName = IconName(\"brand-stackshare\")\n BRAND_STEAM: IconName = IconName(\"brand-steam\")\n BRAND_STORJ: IconName = IconName(\"brand-storj\")\n BRAND_STORYBOOK: IconName = IconName(\"brand-storybook\")\n BRAND_STORYTEL: IconName = IconName(\"brand-storytel\")\n BRAND_STRAVA: IconName = IconName(\"brand-strava\")\n BRAND_STRIPE: IconName = IconName(\"brand-stripe\")\n BRAND_SUBLIME_TEXT: IconName = IconName(\"brand-sublime-text\")\n BRAND_SUGARIZER: IconName = IconName(\"brand-sugarizer\")\n BRAND_SUPABASE: IconName = IconName(\"brand-supabase\")\n BRAND_SUPERHUMAN: IconName = IconName(\"brand-superhuman\")\n BRAND_SUPERNOVA: IconName = IconName(\"brand-supernova\")\n BRAND_SURFSHARK: IconName = IconName(\"brand-surfshark\")\n BRAND_SVELTE: IconName = IconName(\"brand-svelte\")\n BRAND_SWIFT: IconName = IconName(\"brand-swift\")\n BRAND_SYMFONY: IconName = IconName(\"brand-symfony\")\n BRAND_TABLER: IconName = IconName(\"brand-tabler\")\n BRAND_TAILWIND: IconName = IconName(\"brand-tailwind\")\n BRAND_TAOBAO: IconName = IconName(\"brand-taobao\")\n BRAND_TED: IconName = IconName(\"brand-ted\")\n BRAND_TELEGRAM: IconName = IconName(\"brand-telegram\")\n BRAND_TERRAFORM: IconName = IconName(\"brand-terraform\")\n BRAND_TETHER: IconName = IconName(\"brand-tether\")\n BRAND_THREEJS: IconName = IconName(\"brand-threejs\")\n BRAND_TIDAL: IconName = IconName(\"brand-tidal\")\n BRAND_TIKTO_FILLED: IconName = IconName(\"brand-tikto-filled\")\n BRAND_TIKTOK: IconName = IconName(\"brand-tiktok\")\n BRAND_TINDER: IconName = IconName(\"brand-tinder\")\n BRAND_TOPBUZZ: IconName = IconName(\"brand-topbuzz\")\n BRAND_TORCHAIN: IconName = IconName(\"brand-torchain\")\n BRAND_TOYOTA: IconName = IconName(\"brand-toyota\")\n BRAND_TRELLO: IconName = IconName(\"brand-trello\")\n BRAND_TRIPADVISOR: IconName = IconName(\"brand-tripadvisor\")\n BRAND_TUMBLR: IconName = IconName(\"brand-tumblr\")\n BRAND_TWILIO: IconName = IconName(\"brand-twilio\")\n BRAND_TWITCH: IconName = IconName(\"brand-twitch\")\n BRAND_TWITTER: IconName = IconName(\"brand-twitter\")\n BRAND_TWITTER_FILLED: IconName = IconName(\"brand-twitter-filled\")\n BRAND_TYPESCRIPT: IconName = IconName(\"brand-typescript\")\n BRAND_UBER: IconName = IconName(\"brand-uber\")\n BRAND_UBUNTU: IconName = IconName(\"brand-ubuntu\")\n BRAND_UNITY: IconName = IconName(\"brand-unity\")\n BRAND_UNSPLASH: IconName = IconName(\"brand-unsplash\")\n BRAND_UPWORK: IconName = IconName(\"brand-upwork\")\n BRAND_VALORANT: IconName = IconName(\"brand-valorant\")\n BRAND_VERCEL: IconName = IconName(\"brand-vercel\")\n BRAND_VIMEO: IconName = IconName(\"brand-vimeo\")\n BRAND_VINTED: IconName = IconName(\"brand-vinted\")\n BRAND_VISA: IconName = IconName(\"brand-visa\")\n BRAND_VISUAL_STUDIO: IconName = IconName(\"brand-visual-studio\")\n BRAND_VITE: IconName = IconName(\"brand-vite\")\n BRAND_VIVALDI: IconName = IconName(\"brand-vivaldi\")\n BRAND_VK: IconName = IconName(\"brand-vk\")\n BRAND_VLC: IconName = IconName(\"brand-vlc\")\n BRAND_VOLKSWAGEN: IconName = IconName(\"brand-volkswagen\")\n BRAND_VSCO: IconName = IconName(\"brand-vsco\")\n BRAND_VSCODE: IconName = IconName(\"brand-vscode\")\n BRAND_VUE: IconName = IconName(\"brand-vue\")\n BRAND_WALMART: IconName = IconName(\"brand-walmart\")\n BRAND_WAZE: IconName = IconName(\"brand-waze\")\n BRAND_WEBFLOW: IconName = IconName(\"brand-webflow\")\n BRAND_WECHAT: IconName = IconName(\"brand-wechat\")\n BRAND_WEIBO: IconName = IconName(\"brand-weibo\")\n BRAND_WHATSAPP: IconName = IconName(\"brand-whatsapp\")\n BRAND_WIKIPEDIA: IconName = IconName(\"brand-wikipedia\")\n BRAND_WINDOWS: IconName = IconName(\"brand-windows\")\n BRAND_WINDY: IconName = IconName(\"brand-windy\")\n BRAND_WISH: IconName = IconName(\"brand-wish\")\n BRAND_WIX: IconName = IconName(\"brand-wix\")\n BRAND_WORDPRESS: IconName = IconName(\"brand-wordpress\")\n BRAND_XAMARIN: IconName = IconName(\"brand-xamarin\")\n BRAND_XBOX: IconName = IconName(\"brand-xbox\")\n BRAND_XING: IconName = IconName(\"brand-xing\")\n BRAND_YAHOO: IconName = IconName(\"brand-yahoo\")\n BRAND_YANDEX: IconName = IconName(\"brand-yandex\")\n BRAND_YATSE: IconName = IconName(\"brand-yatse\")\n BRAND_YCOMBINATOR: IconName = IconName(\"brand-ycombinator\")\n BRAND_YOUTUBE: IconName = IconName(\"brand-youtube\")\n BRAND_YOUTUBE_KIDS: IconName = IconName(\"brand-youtube-kids\")\n BRAND_ZALANDO: IconName = IconName(\"brand-zalando\")\n BRAND_ZAPIER: IconName = IconName(\"brand-zapier\")\n BRAND_ZEIT: IconName = IconName(\"brand-zeit\")\n BRAND_ZHIHU: IconName = IconName(\"brand-zhihu\")\n BRAND_ZOOM: IconName = IconName(\"brand-zoom\")\n BRAND_ZULIP: IconName = IconName(\"brand-zulip\")\n BRAND_ZWIFT: IconName = IconName(\"brand-zwift\")\n BREAD: IconName = IconName(\"bread\")\n BREAD_OFF: IconName = IconName(\"bread-off\")\n BRIEFCASE: IconName = IconName(\"briefcase\")\n BRIEFCASE_OFF: IconName = IconName(\"briefcase-off\")\n BRIGHTNESS: IconName = IconName(\"brightness\")\n BRIGHTNESS_2: IconName = IconName(\"brightness-2\")\n BRIGHTNESS_DOWN: IconName = IconName(\"brightness-down\")\n BRIGHTNESS_HALF: IconName = IconName(\"brightness-half\")\n BRIGHTNESS_OFF: IconName = IconName(\"brightness-off\")\n BRIGHTNESS_UP: IconName = IconName(\"brightness-up\")\n BROADCAST: IconName = IconName(\"broadcast\")\n BROADCAST_OFF: IconName = IconName(\"broadcast-off\")\n BROWSER: IconName = IconName(\"browser\")\n BROWSER_CHECK: IconName = IconName(\"browser-check\")\n BROWSER_OFF: IconName = IconName(\"browser-off\")\n BROWSER_PLUS: IconName = IconName(\"browser-plus\")\n BROWSER_X: IconName = IconName(\"browser-x\")\n BRUSH: IconName = IconName(\"brush\")\n BRUSH_OFF: IconName = IconName(\"brush-off\")\n BUCKET: IconName = IconName(\"bucket\")\n BUCKET_DROPLET: IconName = IconName(\"bucket-droplet\")\n BUCKET_OFF: IconName = IconName(\"bucket-off\")\n BUG: IconName = IconName(\"bug\")\n BUG_OFF: IconName = IconName(\"bug-off\")\n BUILDING: IconName = IconName(\"building\")\n BUILDING_ARCH: IconName = IconName(\"building-arch\")\n BUILDING_BANK: IconName = IconName(\"building-bank\")\n BUILDING_BRIDGE: IconName = IconName(\"building-bridge\")\n BUILDING_BRIDGE_2: IconName = IconName(\"building-bridge-2\")\n BUILDING_BROADCAST_TOWER: IconName = IconName(\"building-broadcast-tower\")\n BUILDING_CAROUSEL: IconName = IconName(\"building-carousel\")\n BUILDING_CASTLE: IconName = IconName(\"building-castle\")\n BUILDING_CHURCH: IconName = IconName(\"building-church\")\n BUILDING_CIRCUS: IconName = IconName(\"building-circus\")\n BUILDING_COMMUNITY: IconName = IconName(\"building-community\")\n BUILDING_COTTAGE: IconName = IconName(\"building-cottage\")\n BUILDING_ESTATE: IconName = IconName(\"building-estate\")\n BUILDING_FACTORY: IconName = IconName(\"building-factory\")\n BUILDING_FACTORY_2: IconName = IconName(\"building-factory-2\")\n BUILDING_FORTRESS: IconName = IconName(\"building-fortress\")\n BUILDING_HOSPITAL: IconName = IconName(\"building-hospital\")\n BUILDING_LIGHTHOUSE: IconName = IconName(\"building-lighthouse\")\n BUILDING_MONUMENT: IconName = IconName(\"building-monument\")\n BUILDING_MOSQUE: IconName = IconName(\"building-mosque\")\n BUILDING_PAVILION: IconName = IconName(\"building-pavilion\")\n BUILDING_SKYSCRAPER: IconName = IconName(\"building-skyscraper\")\n BUILDING_STADIUM: IconName = IconName(\"building-stadium\")\n BUILDING_STORE: IconName = IconName(\"building-store\")\n BUILDING_TUNNEL: IconName = IconName(\"building-tunnel\")\n BUILDING_WAREHOUSE: IconName = IconName(\"building-warehouse\")\n BUILDING_WIND_TURBINE: IconName = IconName(\"building-wind-turbine\")\n BULB: IconName = IconName(\"bulb\")\n BULB_FILLED: IconName = IconName(\"bulb-filled\")\n BULB_OFF: IconName = IconName(\"bulb-off\")\n BULLDOZER: IconName = IconName(\"bulldozer\")\n BUS: IconName = IconName(\"bus\")\n BUS_OFF: IconName = IconName(\"bus-off\")\n BUS_STOP: IconName = IconName(\"bus-stop\")\n BUSINESSPLAN: IconName = IconName(\"businessplan\")\n BUTTERFLY: IconName = IconName(\"butterfly\")\n CACTUS: IconName = IconName(\"cactus\")\n CACTUS_OFF: IconName = IconName(\"cactus-off\")\n CAKE: IconName = IconName(\"cake\")\n CAKE_OFF: IconName = IconName(\"cake-off\")\n CALCULATOR: IconName = IconName(\"calculator\")\n CALCULATOR_OFF: IconName = IconName(\"calculator-off\")\n CALENDAR: IconName = IconName(\"calendar\")\n CALENDAR_BOLT: IconName = IconName(\"calendar-bolt\")\n CALENDAR_CANCEL: IconName = IconName(\"calendar-cancel\")\n CALENDAR_CHECK: IconName = IconName(\"calendar-check\")\n CALENDAR_CODE: IconName = IconName(\"calendar-code\")\n CALENDAR_COG: IconName = IconName(\"calendar-cog\")\n CALENDAR_DOLLAR: IconName = IconName(\"calendar-dollar\")\n CALENDAR_DOWN: IconName = IconName(\"calendar-down\")\n CALENDAR_DUE: IconName = IconName(\"calendar-due\")\n CALENDAR_EVENT: IconName = IconName(\"calendar-event\")\n CALENDAR_EXCLAMATION: IconName = IconName(\"calendar-exclamation\")\n CALENDAR_HEART: IconName = IconName(\"calendar-heart\")\n CALENDAR_MINUS: IconName = IconName(\"calendar-minus\")\n CALENDAR_OFF: IconName = IconName(\"calendar-off\")\n CALENDAR_PAUSE: IconName = IconName(\"calendar-pause\")\n CALENDAR_PIN: IconName = IconName(\"calendar-pin\")\n CALENDAR_PLUS: IconName = IconName(\"calendar-plus\")\n CALENDAR_QUESTION: IconName = IconName(\"calendar-question\")\n CALENDAR_REPEAT: IconName = IconName(\"calendar-repeat\")\n CALENDAR_SEARCH: IconName = IconName(\"calendar-search\")\n CALENDAR_SHARE: IconName = IconName(\"calendar-share\")\n CALENDAR_STAR: IconName = IconName(\"calendar-star\")\n CALENDAR_STATS: IconName = IconName(\"calendar-stats\")\n CALENDAR_TIME: IconName = IconName(\"calendar-time\")\n CALENDAR_UP: IconName = IconName(\"calendar-up\")\n CALENDAR_X: IconName = IconName(\"calendar-x\")\n CAMERA: IconName = IconName(\"camera\")\n CAMERA_BOLT: IconName = IconName(\"camera-bolt\")\n CAMERA_CANCEL: IconName = IconName(\"camera-cancel\")\n CAMERA_CHECK: IconName = IconName(\"camera-check\")\n CAMERA_CODE: IconName = IconName(\"camera-code\")\n CAMERA_COG: IconName = IconName(\"camera-cog\")\n CAMERA_DOLLAR: IconName = IconName(\"camera-dollar\")\n CAMERA_DOWN: IconName = IconName(\"camera-down\")\n CAMERA_EXCLAMATION: IconName = IconName(\"camera-exclamation\")\n CAMERA_FILLED: IconName = IconName(\"camera-filled\")\n CAMERA_HEART: IconName = IconName(\"camera-heart\")\n CAMERA_MINUS: IconName = IconName(\"camera-minus\")\n CAMERA_OFF: IconName = IconName(\"camera-off\")\n CAMERA_PAUSE: IconName = IconName(\"camera-pause\")\n CAMERA_PIN: IconName = IconName(\"camera-pin\")\n CAMERA_PLUS: IconName = IconName(\"camera-plus\")\n CAMERA_QUESTION: IconName = IconName(\"camera-question\")\n CAMERA_ROTATE: IconName = IconName(\"camera-rotate\")\n CAMERA_SEARCH: IconName = IconName(\"camera-search\")\n CAMERA_SELFIE: IconName = IconName(\"camera-selfie\")\n CAMERA_SHARE: IconName = IconName(\"camera-share\")\n CAMERA_STAR: IconName = IconName(\"camera-star\")\n CAMERA_UP: IconName = IconName(\"camera-up\")\n CAMERA_X: IconName = IconName(\"camera-x\")\n CAMPER: IconName = IconName(\"camper\")\n CAMPFIRE: IconName = IconName(\"campfire\")\n CANDLE: IconName = IconName(\"candle\")\n CANDY: IconName = IconName(\"candy\")\n CANDY_OFF: IconName = IconName(\"candy-off\")\n CANE: IconName = IconName(\"cane\")\n CANNABIS: IconName = IconName(\"cannabis\")\n CAPSULE: IconName = IconName(\"capsule\")\n CAPSULE_HORIZONTAL: IconName = IconName(\"capsule-horizontal\")\n CAPTURE: IconName = IconName(\"capture\")\n CAPTURE_OFF: IconName = IconName(\"capture-off\")\n CAR: IconName = IconName(\"car\")\n CAR_CRANE: IconName = IconName(\"car-crane\")\n CAR_CRASH: IconName = IconName(\"car-crash\")\n CAR_OFF: IconName = IconName(\"car-off\")\n CAR_TURBINE: IconName = IconName(\"car-turbine\")\n CARAVAN: IconName = IconName(\"caravan\")\n CARDBOARDS: IconName = IconName(\"cardboards\")\n CARDBOARDS_OFF: IconName = IconName(\"cardboards-off\")\n CARDS: IconName = IconName(\"cards\")\n CARET_DOWN: IconName = IconName(\"caret-down\")\n CARET_LEFT: IconName = IconName(\"caret-left\")\n CARET_RIGHT: IconName = IconName(\"caret-right\")\n CARET_UP: IconName = IconName(\"caret-up\")\n CAROUSEL_HORIZONTAL: IconName = IconName(\"carousel-horizontal\")\n CAROUSEL_HORIZONTAL_FILLED: IconName = IconName(\"carousel-horizontal-filled\")\n CAROUSEL_VERTICAL: IconName = IconName(\"carousel-vertical\")\n CAROUSEL_VERTICAL_FILLED: IconName = IconName(\"carousel-vertical-filled\")\n CARROT: IconName = IconName(\"carrot\")\n CARROT_OFF: IconName = IconName(\"carrot-off\")\n CASH: IconName = IconName(\"cash\")\n CASH_BANKNOTE: IconName = IconName(\"cash-banknote\")\n CASH_BANKNOTE_OFF: IconName = IconName(\"cash-banknote-off\")\n CASH_OFF: IconName = IconName(\"cash-off\")\n CAST: IconName = IconName(\"cast\")\n CAST_OFF: IconName = IconName(\"cast-off\")\n CAT: IconName = IconName(\"cat\")\n CATEGORY: IconName = IconName(\"category\")\n CATEGORY_2: IconName = IconName(\"category-2\")\n CE: IconName = IconName(\"ce\")\n CE_OFF: IconName = IconName(\"ce-off\")\n CELL: IconName = IconName(\"cell\")\n CELL_SIGNAL_1: IconName = IconName(\"cell-signal-1\")\n CELL_SIGNAL_2: IconName = IconName(\"cell-signal-2\")\n CELL_SIGNAL_3: IconName = IconName(\"cell-signal-3\")\n CELL_SIGNAL_4: IconName = IconName(\"cell-signal-4\")\n CELL_SIGNAL_5: IconName = IconName(\"cell-signal-5\")\n CELL_SIGNAL_OFF: IconName = IconName(\"cell-signal-off\")\n CERTIFICATE: IconName = IconName(\"certificate\")\n CERTIFICATE_2: IconName = IconName(\"certificate-2\")\n CERTIFICATE_2_OFF: IconName = IconName(\"certificate-2-off\")\n CERTIFICATE_OFF: IconName = IconName(\"certificate-off\")\n CHAIR_DIRECTOR: IconName = IconName(\"chair-director\")\n CHALKBOARD: IconName = IconName(\"chalkboard\")\n CHALKBOARD_OFF: IconName = IconName(\"chalkboard-off\")\n CHARGING_PILE: IconName = IconName(\"charging-pile\")\n CHART_ARCS: IconName = IconName(\"chart-arcs\")\n CHART_ARCS_3: IconName = IconName(\"chart-arcs-3\")\n CHART_AREA: IconName = IconName(\"chart-area\")\n CHART_AREA_FILLED: IconName = IconName(\"chart-area-filled\")\n CHART_AREA_LINE: IconName = IconName(\"chart-area-line\")\n CHART_AREA_LINE_FILLED: IconName = IconName(\"chart-area-line-filled\")\n CHART_ARROWS: IconName = IconName(\"chart-arrows\")\n CHART_ARROWS_VERTICAL: IconName = IconName(\"chart-arrows-vertical\")\n CHART_BAR: IconName = IconName(\"chart-bar\")\n CHART_BAR_OFF: IconName = IconName(\"chart-bar-off\")\n CHART_BUBBLE: IconName = IconName(\"chart-bubble\")\n CHART_BUBBLE_FILLED: IconName = IconName(\"chart-bubble-filled\")\n CHART_CANDLE: IconName = IconName(\"chart-candle\")\n CHART_CANDLE_FILLED: IconName = IconName(\"chart-candle-filled\")\n CHART_CIRCLES: IconName = IconName(\"chart-circles\")\n CHART_DONUT: IconName = IconName(\"chart-donut\")\n CHART_DONUT_2: IconName = IconName(\"chart-donut-2\")\n CHART_DONUT_3: IconName = IconName(\"chart-donut-3\")\n CHART_DONUT_4: IconName = IconName(\"chart-donut-4\")\n CHART_DONUT_FILLED: IconName = IconName(\"chart-donut-filled\")\n CHART_DOTS: IconName = IconName(\"chart-dots\")\n CHART_DOTS_2: IconName = IconName(\"chart-dots-2\")\n CHART_DOTS_3: IconName = IconName(\"chart-dots-3\")\n CHART_GRID_DOTS: IconName = IconName(\"chart-grid-dots\")\n CHART_HISTOGRAM: IconName = IconName(\"chart-histogram\")\n CHART_INFOGRAPHIC: IconName = IconName(\"chart-infographic\")\n CHART_LINE: IconName = IconName(\"chart-line\")\n CHART_PIE: IconName = IconName(\"chart-pie\")\n CHART_PIE_2: IconName = IconName(\"chart-pie-2\")\n CHART_PIE_3: IconName = IconName(\"chart-pie-3\")\n CHART_PIE_4: IconName = IconName(\"chart-pie-4\")\n CHART_PIE_FILLED: IconName = IconName(\"chart-pie-filled\")\n CHART_PIE_OFF: IconName = IconName(\"chart-pie-off\")\n CHART_PPF: IconName = IconName(\"chart-ppf\")\n CHART_RADAR: IconName = IconName(\"chart-radar\")\n CHART_SANKEY: IconName = IconName(\"chart-sankey\")\n CHART_TREEMAP: IconName = IconName(\"chart-treemap\")\n CHECK: IconName = IconName(\"check\")\n CHECKBOX: IconName = IconName(\"checkbox\")\n CHECKLIST: IconName = IconName(\"checklist\")\n CHECKS: IconName = IconName(\"checks\")\n CHECKUP_LIST: IconName = IconName(\"checkup-list\")\n CHEESE: IconName = IconName(\"cheese\")\n CHEF_HAT: IconName = IconName(\"chef-hat\")\n CHEF_HAT_OFF: IconName = IconName(\"chef-hat-off\")\n CHERRY: IconName = IconName(\"cherry\")\n CHERRY_FILLED: IconName = IconName(\"cherry-filled\")\n CHESS: IconName = IconName(\"chess\")\n CHESS_BISHOP: IconName = IconName(\"chess-bishop\")\n CHESS_BISHOP_FILLED: IconName = IconName(\"chess-bishop-filled\")\n CHESS_FILLED: IconName = IconName(\"chess-filled\")\n CHESS_KING: IconName = IconName(\"chess-king\")\n CHESS_KING_FILLED: IconName = IconName(\"chess-king-filled\")\n CHESS_KNIGHT: IconName = IconName(\"chess-knight\")\n CHESS_KNIGHT_FILLED: IconName = IconName(\"chess-knight-filled\")\n CHESS_QUEEN: IconName = IconName(\"chess-queen\")\n CHESS_QUEEN_FILLED: IconName = IconName(\"chess-queen-filled\")\n CHESS_ROOK: IconName = IconName(\"chess-rook\")\n CHESS_ROOK_FILLED: IconName = IconName(\"chess-rook-filled\")\n CHEVRON_COMPACT_DOWN: IconName = IconName(\"chevron-compact-down\")\n CHEVRON_COMPACT_LEFT: IconName = IconName(\"chevron-compact-left\")\n CHEVRON_COMPACT_RIGHT: IconName = IconName(\"chevron-compact-right\")\n CHEVRON_COMPACT_UP: IconName = IconName(\"chevron-compact-up\")\n CHEVRON_DOWN: IconName = IconName(\"chevron-down\")\n CHEVRON_DOWN_LEFT: IconName = IconName(\"chevron-down-left\")\n CHEVRON_DOWN_RIGHT: IconName = IconName(\"chevron-down-right\")\n CHEVRON_LEFT: IconName = IconName(\"chevron-left\")\n CHEVRON_LEFT_PIPE: IconName = IconName(\"chevron-left-pipe\")\n CHEVRON_RIGHT: IconName = IconName(\"chevron-right\")\n CHEVRON_RIGHT_PIPE: IconName = IconName(\"chevron-right-pipe\")\n CHEVRON_UP: IconName = IconName(\"chevron-up\")\n CHEVRON_UP_LEFT: IconName = IconName(\"chevron-up-left\")\n CHEVRON_UP_RIGHT: IconName = IconName(\"chevron-up-right\")\n CHEVRONS_DOWN: IconName = IconName(\"chevrons-down\")\n CHEVRONS_DOWN_LEFT: IconName = IconName(\"chevrons-down-left\")\n CHEVRONS_DOWN_RIGHT: IconName = IconName(\"chevrons-down-right\")\n CHEVRONS_LEFT: IconName = IconName(\"chevrons-left\")\n CHEVRONS_RIGHT: IconName = IconName(\"chevrons-right\")\n CHEVRONS_UP: IconName = IconName(\"chevrons-up\")\n CHEVRONS_UP_LEFT: IconName = IconName(\"chevrons-up-left\")\n CHEVRONS_UP_RIGHT: IconName = IconName(\"chevrons-up-right\")\n CHISEL: IconName = IconName(\"chisel\")\n CHRISTMAS_TREE: IconName = IconName(\"christmas-tree\")\n CHRISTMAS_TREE_OFF: IconName = IconName(\"christmas-tree-off\")\n CIRCLE: IconName = IconName(\"circle\")\n CIRCLE_0_FILLED: IconName = IconName(\"circle-0-filled\")\n CIRCLE_1_FILLED: IconName = IconName(\"circle-1-filled\")\n CIRCLE_2_FILLED: IconName = IconName(\"circle-2-filled\")\n CIRCLE_3_FILLED: IconName = IconName(\"circle-3-filled\")\n CIRCLE_4_FILLED: IconName = IconName(\"circle-4-filled\")\n CIRCLE_5_FILLED: IconName = IconName(\"circle-5-filled\")\n CIRCLE_6_FILLED: IconName = IconName(\"circle-6-filled\")\n CIRCLE_7_FILLED: IconName = IconName(\"circle-7-filled\")\n CIRCLE_8_FILLED: IconName = IconName(\"circle-8-filled\")\n CIRCLE_9_FILLED: IconName = IconName(\"circle-9-filled\")\n CIRCLE_ARROW_DOWN: IconName = IconName(\"circle-arrow-down\")\n CIRCLE_ARROW_DOWN_FILLED: IconName = IconName(\"circle-arrow-down-filled\")\n CIRCLE_ARROW_DOWN_LEFT: IconName = IconName(\"circle-arrow-down-left\")\n CIRCLE_ARROW_DOWN_LEFT_FILLED: IconName = IconName(\"circle-arrow-down-left-filled\")\n CIRCLE_ARROW_DOWN_RIGHT: IconName = IconName(\"circle-arrow-down-right\")\n CIRCLE_ARROW_DOWN_RIGHT_FILLED: IconName = IconName(\n \"circle-arrow-down-right-filled\"\n )\n CIRCLE_ARROW_LEFT: IconName = IconName(\"circle-arrow-left\")\n CIRCLE_ARROW_LEFT_FILLED: IconName = IconName(\"circle-arrow-left-filled\")\n CIRCLE_ARROW_RIGHT: IconName = IconName(\"circle-arrow-right\")\n CIRCLE_ARROW_RIGHT_FILLED: IconName = IconName(\"circle-arrow-right-filled\")\n CIRCLE_ARROW_UP: IconName = IconName(\"circle-arrow-up\")\n CIRCLE_ARROW_UP_FILLED: IconName = IconName(\"circle-arrow-up-filled\")\n CIRCLE_ARROW_UP_LEFT: IconName = IconName(\"circle-arrow-up-left\")\n CIRCLE_ARROW_UP_LEFT_FILLED: IconName = IconName(\"circle-arrow-up-left-filled\")\n CIRCLE_ARROW_UP_RIGHT: IconName = IconName(\"circle-arrow-up-right\")\n CIRCLE_ARROW_UP_RIGHT_FILLED: IconName = IconName(\"circle-arrow-up-right-filled\")\n CIRCLE_CARET_DOWN: IconName = IconName(\"circle-caret-down\")\n CIRCLE_CARET_LEFT: IconName = IconName(\"circle-caret-left\")\n CIRCLE_CARET_RIGHT: IconName = IconName(\"circle-caret-right\")\n CIRCLE_CARET_UP: IconName = IconName(\"circle-caret-up\")\n CIRCLE_CHECK: IconName = IconName(\"circle-check\")\n CIRCLE_CHECK_FILLED: IconName = IconName(\"circle-check-filled\")\n CIRCLE_CHEVRON_DOWN: IconName = IconName(\"circle-chevron-down\")\n CIRCLE_CHEVRON_LEFT: IconName = IconName(\"circle-chevron-left\")\n CIRCLE_CHEVRON_RIGHT: IconName = IconName(\"circle-chevron-right\")\n CIRCLE_CHEVRON_UP: IconName = IconName(\"circle-chevron-up\")\n CIRCLE_CHEVRONS_DOWN: IconName = IconName(\"circle-chevrons-down\")\n CIRCLE_CHEVRONS_LEFT: IconName = IconName(\"circle-chevrons-left\")\n CIRCLE_CHEVRONS_RIGHT: IconName = IconName(\"circle-chevrons-right\")\n CIRCLE_CHEVRONS_UP: IconName = IconName(\"circle-chevrons-up\")\n CIRCLE_DASHED: IconName = IconName(\"circle-dashed\")\n CIRCLE_DOT: IconName = IconName(\"circle-dot\")\n CIRCLE_DOT_FILLED: IconName = IconName(\"circle-dot-filled\")\n CIRCLE_DOTTED: IconName = IconName(\"circle-dotted\")\n CIRCLE_FILLED: IconName = IconName(\"circle-filled\")\n CIRCLE_HALF: IconName = IconName(\"circle-half\")\n CIRCLE_HALF_2: IconName = IconName(\"circle-half-2\")\n CIRCLE_HALF_VERTICAL: IconName = IconName(\"circle-half-vertical\")\n CIRCLE_KEY: IconName = IconName(\"circle-key\")\n CIRCLE_KEY_FILLED: IconName = IconName(\"circle-key-filled\")\n CIRCLE_LETTER_A: IconName = IconName(\"circle-letter-a\")\n CIRCLE_LETTER_B: IconName = IconName(\"circle-letter-b\")\n CIRCLE_LETTER_C: IconName = IconName(\"circle-letter-c\")\n CIRCLE_LETTER_D: IconName = IconName(\"circle-letter-d\")\n CIRCLE_LETTER_E: IconName = IconName(\"circle-letter-e\")\n CIRCLE_LETTER_F: IconName = IconName(\"circle-letter-f\")\n CIRCLE_LETTER_G: IconName = IconName(\"circle-letter-g\")\n CIRCLE_LETTER_H: IconName = IconName(\"circle-letter-h\")\n CIRCLE_LETTER_I: IconName = IconName(\"circle-letter-i\")\n CIRCLE_LETTER_J: IconName = IconName(\"circle-letter-j\")\n CIRCLE_LETTER_K: IconName = IconName(\"circle-letter-k\")\n CIRCLE_LETTER_L: IconName = IconName(\"circle-letter-l\")\n CIRCLE_LETTER_M: IconName = IconName(\"circle-letter-m\")\n CIRCLE_LETTER_N: IconName = IconName(\"circle-letter-n\")\n CIRCLE_LETTER_O: IconName = IconName(\"circle-letter-o\")\n CIRCLE_LETTER_P: IconName = IconName(\"circle-letter-p\")\n CIRCLE_LETTER_Q: IconName = IconName(\"circle-letter-q\")\n CIRCLE_LETTER_R: IconName = IconName(\"circle-letter-r\")\n CIRCLE_LETTER_S: IconName = IconName(\"circle-letter-s\")\n CIRCLE_LETTER_T: IconName = IconName(\"circle-letter-t\")\n CIRCLE_LETTER_U: IconName = IconName(\"circle-letter-u\")\n CIRCLE_LETTER_V: IconName = IconName(\"circle-letter-v\")\n CIRCLE_LETTER_W: IconName = IconName(\"circle-letter-w\")\n CIRCLE_LETTER_X: IconName = IconName(\"circle-letter-x\")\n CIRCLE_LETTER_Y: IconName = IconName(\"circle-letter-y\")\n CIRCLE_LETTER_Z: IconName = IconName(\"circle-letter-z\")\n CIRCLE_MINUS: IconName = IconName(\"circle-minus\")\n CIRCLE_NUMBER_0: IconName = IconName(\"circle-number-0\")\n CIRCLE_NUMBER_1: IconName = IconName(\"circle-number-1\")\n CIRCLE_NUMBER_2: IconName = IconName(\"circle-number-2\")\n CIRCLE_NUMBER_3: IconName = IconName(\"circle-number-3\")\n CIRCLE_NUMBER_4: IconName = IconName(\"circle-number-4\")\n CIRCLE_NUMBER_5: IconName = IconName(\"circle-number-5\")\n CIRCLE_NUMBER_6: IconName = IconName(\"circle-number-6\")\n CIRCLE_NUMBER_7: IconName = IconName(\"circle-number-7\")\n CIRCLE_NUMBER_8: IconName = IconName(\"circle-number-8\")\n CIRCLE_NUMBER_9: IconName = IconName(\"circle-number-9\")\n CIRCLE_OFF: IconName = IconName(\"circle-off\")\n CIRCLE_PLUS: IconName = IconName(\"circle-plus\")\n CIRCLE_RECTANGLE: IconName = IconName(\"circle-rectangle\")\n CIRCLE_RECTANGLE_OFF: IconName = IconName(\"circle-rectangle-off\")\n CIRCLE_SQUARE: IconName = IconName(\"circle-square\")\n CIRCLE_TRIANGLE: IconName = IconName(\"circle-triangle\")\n CIRCLE_X: IconName = IconName(\"circle-x\")\n CIRCLE_X_FILLED: IconName = IconName(\"circle-x-filled\")\n CIRCLES: IconName = IconName(\"circles\")\n CIRCLES_FILLED: IconName = IconName(\"circles-filled\")\n CIRCLES_RELATION: IconName = IconName(\"circles-relation\")\n CIRCUIT_AMMETER: IconName = IconName(\"circuit-ammeter\")\n CIRCUIT_BATTERY: IconName = IconName(\"circuit-battery\")\n CIRCUIT_BULB: IconName = IconName(\"circuit-bulb\")\n CIRCUIT_CAPACITOR: IconName = IconName(\"circuit-capacitor\")\n CIRCUIT_CAPACITOR_POLARIZED: IconName = IconName(\"circuit-capacitor-polarized\")\n CIRCUIT_CELL: IconName = IconName(\"circuit-cell\")\n CIRCUIT_CELL_PLUS: IconName = IconName(\"circuit-cell-plus\")\n CIRCUIT_CHANGEOVER: IconName = IconName(\"circuit-changeover\")\n CIRCUIT_DIODE: IconName = IconName(\"circuit-diode\")\n CIRCUIT_DIODE_ZENER: IconName = IconName(\"circuit-diode-zener\")\n CIRCUIT_GROUND: IconName = IconName(\"circuit-ground\")\n CIRCUIT_GROUND_DIGITAL: IconName = IconName(\"circuit-ground-digital\")\n CIRCUIT_INDUCTOR: IconName = IconName(\"circuit-inductor\")\n CIRCUIT_MOTOR: IconName = IconName(\"circuit-motor\")\n CIRCUIT_PUSHBUTTON: IconName = IconName(\"circuit-pushbutton\")\n CIRCUIT_RESISTOR: IconName = IconName(\"circuit-resistor\")\n CIRCUIT_SWITCH_CLOSED: IconName = IconName(\"circuit-switch-closed\")\n CIRCUIT_SWITCH_OPEN: IconName = IconName(\"circuit-switch-open\")\n CIRCUIT_VOLTMETER: IconName = IconName(\"circuit-voltmeter\")\n CLEAR_ALL: IconName = IconName(\"clear-all\")\n CLEAR_FORMATTING: IconName = IconName(\"clear-formatting\")\n CLICK: IconName = IconName(\"click\")\n CLIPBOARD: IconName = IconName(\"clipboard\")\n CLIPBOARD_CHECK: IconName = IconName(\"clipboard-check\")\n CLIPBOARD_COPY: IconName = IconName(\"clipboard-copy\")\n CLIPBOARD_DATA: IconName = IconName(\"clipboard-data\")\n CLIPBOARD_HEART: IconName = IconName(\"clipboard-heart\")\n CLIPBOARD_LIST: IconName = IconName(\"clipboard-list\")\n CLIPBOARD_OFF: IconName = IconName(\"clipboard-off\")\n CLIPBOARD_PLUS: IconName = IconName(\"clipboard-plus\")\n CLIPBOARD_TEXT: IconName = IconName(\"clipboard-text\")\n CLIPBOARD_TYPOGRAPHY: IconName = IconName(\"clipboard-typography\")\n CLIPBOARD_X: IconName = IconName(\"clipboard-x\")\n CLOCK: IconName = IconName(\"clock\")\n CLOCK_2: IconName = IconName(\"clock-2\")\n CLOCK_BOLT: IconName = IconName(\"clock-bolt\")\n CLOCK_CANCEL: IconName = IconName(\"clock-cancel\")\n CLOCK_CHECK: IconName = IconName(\"clock-check\")\n CLOCK_CODE: IconName = IconName(\"clock-code\")\n CLOCK_COG: IconName = IconName(\"clock-cog\")\n CLOCK_DOLLAR: IconName = IconName(\"clock-dollar\")\n CLOCK_DOWN: IconName = IconName(\"clock-down\")\n CLOCK_EDIT: IconName = IconName(\"clock-edit\")\n CLOCK_EXCLAMATION: IconName = IconName(\"clock-exclamation\")\n CLOCK_FILLED: IconName = IconName(\"clock-filled\")\n CLOCK_HEART: IconName = IconName(\"clock-heart\")\n CLOCK_HOUR_1: IconName = IconName(\"clock-hour-1\")\n CLOCK_HOUR_10: IconName = IconName(\"clock-hour-10\")\n CLOCK_HOUR_11: IconName = IconName(\"clock-hour-11\")\n CLOCK_HOUR_12: IconName = IconName(\"clock-hour-12\")\n CLOCK_HOUR_2: IconName = IconName(\"clock-hour-2\")\n CLOCK_HOUR_3: IconName = IconName(\"clock-hour-3\")\n CLOCK_HOUR_4: IconName = IconName(\"clock-hour-4\")\n CLOCK_HOUR_5: IconName = IconName(\"clock-hour-5\")\n CLOCK_HOUR_6: IconName = IconName(\"clock-hour-6\")\n CLOCK_HOUR_7: IconName = IconName(\"clock-hour-7\")\n CLOCK_HOUR_8: IconName = IconName(\"clock-hour-8\")\n CLOCK_HOUR_9: IconName = IconName(\"clock-hour-9\")\n CLOCK_MINUS: IconName = IconName(\"clock-minus\")\n CLOCK_OFF: IconName = IconName(\"clock-off\")\n CLOCK_PAUSE: IconName = IconName(\"clock-pause\")\n CLOCK_PIN: IconName = IconName(\"clock-pin\")\n CLOCK_PLAY: IconName = IconName(\"clock-play\")\n CLOCK_PLUS: IconName = IconName(\"clock-plus\")\n CLOCK_QUESTION: IconName = IconName(\"clock-question\")\n CLOCK_RECORD: IconName = IconName(\"clock-record\")\n CLOCK_SEARCH: IconName = IconName(\"clock-search\")\n CLOCK_SHARE: IconName = IconName(\"clock-share\")\n CLOCK_SHIELD: IconName = IconName(\"clock-shield\")\n CLOCK_STAR: IconName = IconName(\"clock-star\")\n CLOCK_STOP: IconName = IconName(\"clock-stop\")\n CLOCK_UP: IconName = IconName(\"clock-up\")\n CLOCK_X: IconName = IconName(\"clock-x\")\n CLOTHES_RACK: IconName = IconName(\"clothes-rack\")\n CLOTHES_RACK_OFF: IconName = IconName(\"clothes-rack-off\")\n CLOUD: IconName = IconName(\"cloud\")\n CLOUD_BOLT: IconName = IconName(\"cloud-bolt\")\n CLOUD_CANCEL: IconName = IconName(\"cloud-cancel\")\n CLOUD_CHECK: IconName = IconName(\"cloud-check\")\n CLOUD_CODE: IconName = IconName(\"cloud-code\")\n CLOUD_COG: IconName = IconName(\"cloud-cog\")\n CLOUD_COMPUTING: IconName = IconName(\"cloud-computing\")\n CLOUD_DATA_CONNECTION: IconName = IconName(\"cloud-data-connection\")\n CLOUD_DOLLAR: IconName = IconName(\"cloud-dollar\")\n CLOUD_DOWN: IconName = IconName(\"cloud-down\")\n CLOUD_DOWNLOAD: IconName = IconName(\"cloud-download\")\n CLOUD_EXCLAMATION: IconName = IconName(\"cloud-exclamation\")\n CLOUD_FILLED: IconName = IconName(\"cloud-filled\")\n CLOUD_FOG: IconName = IconName(\"cloud-fog\")\n CLOUD_HEART: IconName = IconName(\"cloud-heart\")\n CLOUD_LOCK: IconName = IconName(\"cloud-lock\")\n CLOUD_LOCK_OPEN: IconName = IconName(\"cloud-lock-open\")\n CLOUD_MINUS: IconName = IconName(\"cloud-minus\")\n CLOUD_OFF: IconName = IconName(\"cloud-off\")\n CLOUD_PAUSE: IconName = IconName(\"cloud-pause\")\n CLOUD_PIN: IconName = IconName(\"cloud-pin\")\n CLOUD_PLUS: IconName = IconName(\"cloud-plus\")\n CLOUD_QUESTION: IconName = IconName(\"cloud-question\")\n CLOUD_RAIN: IconName = IconName(\"cloud-rain\")\n CLOUD_SEARCH: IconName = IconName(\"cloud-search\")\n CLOUD_SHARE: IconName = IconName(\"cloud-share\")\n CLOUD_SNOW: IconName = IconName(\"cloud-snow\")\n CLOUD_STAR: IconName = IconName(\"cloud-star\")\n CLOUD_STORM: IconName = IconName(\"cloud-storm\")\n CLOUD_UP: IconName = IconName(\"cloud-up\")\n CLOUD_UPLOAD: IconName = IconName(\"cloud-upload\")\n CLOUD_X: IconName = IconName(\"cloud-x\")\n CLOVER: IconName = IconName(\"clover\")\n CLOVER_2: IconName = IconName(\"clover-2\")\n CLUBS: IconName = IconName(\"clubs\")\n CLUBS_FILLED: IconName = IconName(\"clubs-filled\")\n CODE: IconName = IconName(\"code\")\n CODE_ASTERIX: IconName = IconName(\"code-asterix\")\n CODE_CIRCLE: IconName = IconName(\"code-circle\")\n CODE_CIRCLE_2: IconName = IconName(\"code-circle-2\")\n CODE_DOTS: IconName = IconName(\"code-dots\")\n CODE_MINUS: IconName = IconName(\"code-minus\")\n CODE_OFF: IconName = IconName(\"code-off\")\n CODE_PLUS: IconName = IconName(\"code-plus\")\n COFFEE: IconName = IconName(\"coffee\")\n COFFEE_OFF: IconName = IconName(\"coffee-off\")\n COFFIN: IconName = IconName(\"coffin\")\n COIN: IconName = IconName(\"coin\")\n COIN_BITCOIN: IconName = IconName(\"coin-bitcoin\")\n COIN_EURO: IconName = IconName(\"coin-euro\")\n COIN_MONERO: IconName = IconName(\"coin-monero\")\n COIN_OFF: IconName = IconName(\"coin-off\")\n COIN_POUND: IconName = IconName(\"coin-pound\")\n COIN_RUPEE: IconName = IconName(\"coin-rupee\")\n COIN_YEN: IconName = IconName(\"coin-yen\")\n COIN_YUAN: IconName = IconName(\"coin-yuan\")\n COINS: IconName = IconName(\"coins\")\n COLOR_FILTER: IconName = IconName(\"color-filter\")\n COLOR_PICKER: IconName = IconName(\"color-picker\")\n COLOR_PICKER_OFF: IconName = IconName(\"color-picker-off\")\n COLOR_SWATCH: IconName = IconName(\"color-swatch\")\n COLOR_SWATCH_OFF: IconName = IconName(\"color-swatch-off\")\n COLUMN_INSERT_LEFT: IconName = IconName(\"column-insert-left\")\n COLUMN_INSERT_RIGHT: IconName = IconName(\"column-insert-right\")\n COLUMN_REMOVE: IconName = IconName(\"column-remove\")\n COLUMNS: IconName = IconName(\"columns\")\n COLUMNS_1: IconName = IconName(\"columns-1\")\n COLUMNS_2: IconName = IconName(\"columns-2\")\n COLUMNS_3: IconName = IconName(\"columns-3\")\n COLUMNS_OFF: IconName = IconName(\"columns-off\")\n COMET: IconName = IconName(\"comet\")\n COMMAND: IconName = IconName(\"command\")\n COMMAND_OFF: IconName = IconName(\"command-off\")\n COMPASS: IconName = IconName(\"compass\")\n COMPASS_OFF: IconName = IconName(\"compass-off\")\n COMPONENTS: IconName = IconName(\"components\")\n COMPONENTS_OFF: IconName = IconName(\"components-off\")\n CONE: IconName = IconName(\"cone\")\n CONE_2: IconName = IconName(\"cone-2\")\n CONE_OFF: IconName = IconName(\"cone-off\")\n CONE_PLUS: IconName = IconName(\"cone-plus\")\n CONFETTI: IconName = IconName(\"confetti\")\n CONFETTI_OFF: IconName = IconName(\"confetti-off\")\n CONFUCIUS: IconName = IconName(\"confucius\")\n CONTAINER: IconName = IconName(\"container\")\n CONTAINER_OFF: IconName = IconName(\"container-off\")\n CONTRAST: IconName = IconName(\"contrast\")\n CONTRAST_2: IconName = IconName(\"contrast-2\")\n CONTRAST_2_OFF: IconName = IconName(\"contrast-2-off\")\n CONTRAST_OFF: IconName = IconName(\"contrast-off\")\n COOKER: IconName = IconName(\"cooker\")\n COOKIE: IconName = IconName(\"cookie\")\n COOKIE_MAN: IconName = IconName(\"cookie-man\")\n COOKIE_OFF: IconName = IconName(\"cookie-off\")\n COPY: IconName = IconName(\"copy\")\n COPY_OFF: IconName = IconName(\"copy-off\")\n COPYLEFT: IconName = IconName(\"copyleft\")\n COPYLEFT_FILLED: IconName = IconName(\"copyleft-filled\")\n COPYLEFT_OFF: IconName = IconName(\"copyleft-off\")\n COPYRIGHT: IconName = IconName(\"copyright\")\n COPYRIGHT_FILLED: IconName = IconName(\"copyright-filled\")\n COPYRIGHT_OFF: IconName = IconName(\"copyright-off\")\n CORNER_DOWN_LEFT: IconName = IconName(\"corner-down-left\")\n CORNER_DOWN_LEFT_DOUBLE: IconName = IconName(\"corner-down-left-double\")\n CORNER_DOWN_RIGHT: IconName = IconName(\"corner-down-right\")\n CORNER_DOWN_RIGHT_DOUBLE: IconName = IconName(\"corner-down-right-double\")\n CORNER_LEFT_DOWN: IconName = IconName(\"corner-left-down\")\n CORNER_LEFT_DOWN_DOUBLE: IconName = IconName(\"corner-left-down-double\")\n CORNER_LEFT_UP: IconName = IconName(\"corner-left-up\")\n CORNER_LEFT_UP_DOUBLE: IconName = IconName(\"corner-left-up-double\")\n CORNER_RIGHT_DOWN: IconName = IconName(\"corner-right-down\")\n CORNER_RIGHT_DOWN_DOUBLE: IconName = IconName(\"corner-right-down-double\")\n CORNER_RIGHT_UP: IconName = IconName(\"corner-right-up\")\n CORNER_RIGHT_UP_DOUBLE: IconName = IconName(\"corner-right-up-double\")\n CORNER_UP_LEFT: IconName = IconName(\"corner-up-left\")\n CORNER_UP_LEFT_DOUBLE: IconName = IconName(\"corner-up-left-double\")\n CORNER_UP_RIGHT: IconName = IconName(\"corner-up-right\")\n CORNER_UP_RIGHT_DOUBLE: IconName = IconName(\"corner-up-right-double\")\n CPU: IconName = IconName(\"cpu\")\n CPU_2: IconName = IconName(\"cpu-2\")\n CPU_OFF: IconName = IconName(\"cpu-off\")\n CRANE: IconName = IconName(\"crane\")\n CRANE_OFF: IconName = IconName(\"crane-off\")\n CREATIVE_COMMONS: IconName = IconName(\"creative-commons\")\n CREATIVE_COMMONS_BY: IconName = IconName(\"creative-commons-by\")\n CREATIVE_COMMONS_NC: IconName = IconName(\"creative-commons-nc\")\n CREATIVE_COMMONS_ND: IconName = IconName(\"creative-commons-nd\")\n CREATIVE_COMMONS_OFF: IconName = IconName(\"creative-commons-off\")\n CREATIVE_COMMONS_SA: IconName = IconName(\"creative-commons-sa\")\n CREATIVE_COMMONS_ZERO: IconName = IconName(\"creative-commons-zero\")\n CREDIT_CARD: IconName = IconName(\"credit-card\")\n CREDIT_CARD_OFF: IconName = IconName(\"credit-card-off\")\n CRICKET: IconName = IconName(\"cricket\")\n CROP: IconName = IconName(\"crop\")\n CROSS: IconName = IconName(\"cross\")\n CROSS_FILLED: IconName = IconName(\"cross-filled\")\n CROSS_OFF: IconName = IconName(\"cross-off\")\n CROSSHAIR: IconName = IconName(\"crosshair\")\n CROWN: IconName = IconName(\"crown\")\n CROWN_OFF: IconName = IconName(\"crown-off\")\n CRUTCHES: IconName = IconName(\"crutches\")\n CRUTCHES_OFF: IconName = IconName(\"crutches-off\")\n CRYSTAL_BALL: IconName = IconName(\"crystal-ball\")\n CSV: IconName = IconName(\"csv\")\n CUBE: IconName = IconName(\"cube\")\n CUBE_OFF: IconName = IconName(\"cube-off\")\n CUBE_PLUS: IconName = IconName(\"cube-plus\")\n CUBE_SEND: IconName = IconName(\"cube-send\")\n CUBE_UNFOLDED: IconName = IconName(\"cube-unfolded\")\n CUP: IconName = IconName(\"cup\")\n CUP_OFF: IconName = IconName(\"cup-off\")\n CURLING: IconName = IconName(\"curling\")\n CURLY_LOOP: IconName = IconName(\"curly-loop\")\n CURRENCY: IconName = IconName(\"currency\")\n CURRENCY_AFGHANI: IconName = IconName(\"currency-afghani\")\n CURRENCY_BAHRAINI: IconName = IconName(\"currency-bahraini\")\n CURRENCY_BAHT: IconName = IconName(\"currency-baht\")\n CURRENCY_BITCOIN: IconName = IconName(\"currency-bitcoin\")\n CURRENCY_CENT: IconName = IconName(\"currency-cent\")\n CURRENCY_DINAR: IconName = IconName(\"currency-dinar\")\n CURRENCY_DIRHAM: IconName = IconName(\"currency-dirham\")\n CURRENCY_DOGECOIN: IconName = IconName(\"currency-dogecoin\")\n CURRENCY_DOLLAR: IconName = IconName(\"currency-dollar\")\n CURRENCY_DOLLAR_AUSTRALIAN: IconName = IconName(\"currency-dollar-australian\")\n CURRENCY_DOLLAR_BRUNEI: IconName = IconName(\"currency-dollar-brunei\")\n CURRENCY_DOLLAR_CANADIAN: IconName = IconName(\"currency-dollar-canadian\")\n CURRENCY_DOLLAR_GUYANESE: IconName = IconName(\"currency-dollar-guyanese\")\n CURRENCY_DOLLAR_OFF: IconName = IconName(\"currency-dollar-off\")\n CURRENCY_DOLLAR_SINGAPORE: IconName = IconName(\"currency-dollar-singapore\")\n CURRENCY_DOLLAR_ZIMBABWEAN: IconName = IconName(\"currency-dollar-zimbabwean\")\n CURRENCY_DONG: IconName = IconName(\"currency-dong\")\n CURRENCY_DRAM: IconName = IconName(\"currency-dram\")\n CURRENCY_ETHEREUM: IconName = IconName(\"currency-ethereum\")\n CURRENCY_EURO: IconName = IconName(\"currency-euro\")\n CURRENCY_EURO_OFF: IconName = IconName(\"currency-euro-off\")\n CURRENCY_FLORIN: IconName = IconName(\"currency-florin\")\n CURRENCY_FORINT: IconName = IconName(\"currency-forint\")\n CURRENCY_FRANK: IconName = IconName(\"currency-frank\")\n CURRENCY_GUARANI: IconName = IconName(\"currency-guarani\")\n CURRENCY_HRYVNIA: IconName = IconName(\"currency-hryvnia\")\n CURRENCY_IRANIAN_RIAL: IconName = IconName(\"currency-iranian-rial\")\n CURRENCY_KIP: IconName = IconName(\"currency-kip\")\n CURRENCY_KRONE_CZECH: IconName = IconName(\"currency-krone-czech\")\n CURRENCY_KRONE_DANISH: IconName = IconName(\"currency-krone-danish\")\n CURRENCY_KRONE_SWEDISH: IconName = IconName(\"currency-krone-swedish\")\n CURRENCY_LARI: IconName = IconName(\"currency-lari\")\n CURRENCY_LEU: IconName = IconName(\"currency-leu\")\n CURRENCY_LIRA: IconName = IconName(\"currency-lira\")\n CURRENCY_LITECOIN: IconName = IconName(\"currency-litecoin\")\n CURRENCY_LYD: IconName = IconName(\"currency-lyd\")\n CURRENCY_MANAT: IconName = IconName(\"currency-manat\")\n CURRENCY_MONERO: IconName = IconName(\"currency-monero\")\n CURRENCY_NAIRA: IconName = IconName(\"currency-naira\")\n CURRENCY_NANO: IconName = IconName(\"currency-nano\")\n CURRENCY_OFF: IconName = IconName(\"currency-off\")\n CURRENCY_PAANGA: IconName = IconName(\"currency-paanga\")\n CURRENCY_PESO: IconName = IconName(\"currency-peso\")\n CURRENCY_POUND: IconName = IconName(\"currency-pound\")\n CURRENCY_POUND_OFF: IconName = IconName(\"currency-pound-off\")\n CURRENCY_QUETZAL: IconName = IconName(\"currency-quetzal\")\n CURRENCY_REAL: IconName = IconName(\"currency-real\")\n CURRENCY_RENMINBI: IconName = IconName(\"currency-renminbi\")\n CURRENCY_RIPPLE: IconName = IconName(\"currency-ripple\")\n CURRENCY_RIYAL: IconName = IconName(\"currency-riyal\")\n CURRENCY_RUBEL: IconName = IconName(\"currency-rubel\")\n CURRENCY_RUFIYAA: IconName = IconName(\"currency-rufiyaa\")\n CURRENCY_RUPEE: IconName = IconName(\"currency-rupee\")\n CURRENCY_RUPEE_NEPALESE: IconName = IconName(\"currency-rupee-nepalese\")\n CURRENCY_SHEKEL: IconName = IconName(\"currency-shekel\")\n CURRENCY_SOLANA: IconName = IconName(\"currency-solana\")\n CURRENCY_SOM: IconName = IconName(\"currency-som\")\n CURRENCY_TAKA: IconName = IconName(\"currency-taka\")\n CURRENCY_TENGE: IconName = IconName(\"currency-tenge\")\n CURRENCY_TUGRIK: IconName = IconName(\"currency-tugrik\")\n CURRENCY_WON: IconName = IconName(\"currency-won\")\n CURRENCY_YEN: IconName = IconName(\"currency-yen\")\n CURRENCY_YEN_OFF: IconName = IconName(\"currency-yen-off\")\n CURRENCY_YUAN: IconName = IconName(\"currency-yuan\")\n CURRENCY_ZLOTY: IconName = IconName(\"currency-zloty\")\n CURRENT_LOCATION: IconName = IconName(\"current-location\")\n CURRENT_LOCATION_OFF: IconName = IconName(\"current-location-off\")\n CURSOR_OFF: IconName = IconName(\"cursor-off\")\n CURSOR_TEXT: IconName = IconName(\"cursor-text\")\n CUT: IconName = IconName(\"cut\")\n CYLINDER: IconName = IconName(\"cylinder\")\n CYLINDER_OFF: IconName = IconName(\"cylinder-off\")\n CYLINDER_PLUS: IconName = IconName(\"cylinder-plus\")\n DASHBOARD: IconName = IconName(\"dashboard\")\n DASHBOARD_OFF: IconName = IconName(\"dashboard-off\")\n DATABASE: IconName = IconName(\"database\")\n DATABASE_COG: IconName = IconName(\"database-cog\")\n DATABASE_DOLLAR: IconName = IconName(\"database-dollar\")\n DATABASE_EDIT: IconName = IconName(\"database-edit\")\n DATABASE_EXCLAMATION: IconName = IconName(\"database-exclamation\")\n DATABASE_EXPORT: IconName = IconName(\"database-export\")\n DATABASE_HEART: IconName = IconName(\"database-heart\")\n DATABASE_IMPORT: IconName = IconName(\"database-import\")\n DATABASE_LEAK: IconName = IconName(\"database-leak\")\n DATABASE_MINUS: IconName = IconName(\"database-minus\")\n DATABASE_OFF: IconName = IconName(\"database-off\")\n DATABASE_PLUS: IconName = IconName(\"database-plus\")\n DATABASE_SEARCH: IconName = IconName(\"database-search\")\n DATABASE_SHARE: IconName = IconName(\"database-share\")\n DATABASE_STAR: IconName = IconName(\"database-star\")\n DATABASE_X: IconName = IconName(\"database-x\")\n DECIMAL: IconName = IconName(\"decimal\")\n DEER: IconName = IconName(\"deer\")\n DELTA: IconName = IconName(\"delta\")\n DENTAL: IconName = IconName(\"dental\")\n DENTAL_BROKEN: IconName = IconName(\"dental-broken\")\n DENTAL_OFF: IconName = IconName(\"dental-off\")\n DESELECT: IconName = IconName(\"deselect\")\n DETAILS: IconName = IconName(\"details\")\n DETAILS_OFF: IconName = IconName(\"details-off\")\n DEVICE_AIRPODS: IconName = IconName(\"device-airpods\")\n DEVICE_AIRPODS_CASE: IconName = IconName(\"device-airpods-case\")\n DEVICE_AIRTAG: IconName = IconName(\"device-airtag\")\n DEVICE_ANALYTICS: IconName = IconName(\"device-analytics\")\n DEVICE_AUDIO_TAPE: IconName = IconName(\"device-audio-tape\")\n DEVICE_CAMERA_PHONE: IconName = IconName(\"device-camera-phone\")\n DEVICE_CCTV: IconName = IconName(\"device-cctv\")\n DEVICE_CCTV_OFF: IconName = IconName(\"device-cctv-off\")\n DEVICE_COMPUTER_CAMERA: IconName = IconName(\"device-computer-camera\")\n DEVICE_COMPUTER_CAMERA_OFF: IconName = IconName(\"device-computer-camera-off\")\n DEVICE_DESKTOP: IconName = IconName(\"device-desktop\")\n DEVICE_DESKTOP_ANALYTICS: IconName = IconName(\"device-desktop-analytics\")\n DEVICE_DESKTOP_BOLT: IconName = IconName(\"device-desktop-bolt\")\n DEVICE_DESKTOP_CANCEL: IconName = IconName(\"device-desktop-cancel\")\n DEVICE_DESKTOP_CHECK: IconName = IconName(\"device-desktop-check\")\n DEVICE_DESKTOP_CODE: IconName = IconName(\"device-desktop-code\")\n DEVICE_DESKTOP_COG: IconName = IconName(\"device-desktop-cog\")\n DEVICE_DESKTOP_DOLLAR: IconName = IconName(\"device-desktop-dollar\")\n DEVICE_DESKTOP_DOWN: IconName = IconName(\"device-desktop-down\")\n DEVICE_DESKTOP_EXCLAMATION: IconName = IconName(\"device-desktop-exclamation\")\n DEVICE_DESKTOP_HEART: IconName = IconName(\"device-desktop-heart\")\n DEVICE_DESKTOP_MINUS: IconName = IconName(\"device-desktop-minus\")\n DEVICE_DESKTOP_OFF: IconName = IconName(\"device-desktop-off\")\n DEVICE_DESKTOP_PAUSE: IconName = IconName(\"device-desktop-pause\")\n DEVICE_DESKTOP_PIN: IconName = IconName(\"device-desktop-pin\")\n DEVICE_DESKTOP_PLUS: IconName = IconName(\"device-desktop-plus\")\n DEVICE_DESKTOP_QUESTION: IconName = IconName(\"device-desktop-question\")\n DEVICE_DESKTOP_SEARCH: IconName = IconName(\"device-desktop-search\")\n DEVICE_DESKTOP_SHARE: IconName = IconName(\"device-desktop-share\")\n DEVICE_DESKTOP_STAR: IconName = IconName(\"device-desktop-star\")\n DEVICE_DESKTOP_UP: IconName = IconName(\"device-desktop-up\")\n DEVICE_DESKTOP_X: IconName = IconName(\"device-desktop-x\")\n DEVICE_FLOPPY: IconName = IconName(\"device-floppy\")\n DEVICE_GAMEPAD: IconName = IconName(\"device-gamepad\")\n DEVICE_GAMEPAD_2: IconName = IconName(\"device-gamepad-2\")\n DEVICE_HEART_MONITOR: IconName = IconName(\"device-heart-monitor\")\n DEVICE_HEART_MONITOR_FILLED: IconName = IconName(\"device-heart-monitor-filled\")\n DEVICE_IMAC: IconName = IconName(\"device-imac\")\n DEVICE_IMAC_BOLT: IconName = IconName(\"device-imac-bolt\")\n DEVICE_IMAC_CANCEL: IconName = IconName(\"device-imac-cancel\")\n DEVICE_IMAC_CHECK: IconName = IconName(\"device-imac-check\")\n DEVICE_IMAC_CODE: IconName = IconName(\"device-imac-code\")\n DEVICE_IMAC_COG: IconName = IconName(\"device-imac-cog\")\n DEVICE_IMAC_DOLLAR: IconName = IconName(\"device-imac-dollar\")\n DEVICE_IMAC_DOWN: IconName = IconName(\"device-imac-down\")\n DEVICE_IMAC_EXCLAMATION: IconName = IconName(\"device-imac-exclamation\")\n DEVICE_IMAC_HEART: IconName = IconName(\"device-imac-heart\")\n DEVICE_IMAC_MINUS: IconName = IconName(\"device-imac-minus\")\n DEVICE_IMAC_OFF: IconName = IconName(\"device-imac-off\")\n DEVICE_IMAC_PAUSE: IconName = IconName(\"device-imac-pause\")\n DEVICE_IMAC_PIN: IconName = IconName(\"device-imac-pin\")\n DEVICE_IMAC_PLUS: IconName = IconName(\"device-imac-plus\")\n DEVICE_IMAC_QUESTION: IconName = IconName(\"device-imac-question\")\n DEVICE_IMAC_SEARCH: IconName = IconName(\"device-imac-search\")\n DEVICE_IMAC_SHARE: IconName = IconName(\"device-imac-share\")\n DEVICE_IMAC_STAR: IconName = IconName(\"device-imac-star\")\n DEVICE_IMAC_UP: IconName = IconName(\"device-imac-up\")\n DEVICE_IMAC_X: IconName = IconName(\"device-imac-x\")\n DEVICE_IPAD: IconName = IconName(\"device-ipad\")\n DEVICE_IPAD_BOLT: IconName = IconName(\"device-ipad-bolt\")\n DEVICE_IPAD_CANCEL: IconName = IconName(\"device-ipad-cancel\")\n DEVICE_IPAD_CHECK: IconName = IconName(\"device-ipad-check\")\n DEVICE_IPAD_CODE: IconName = IconName(\"device-ipad-code\")\n DEVICE_IPAD_COG: IconName = IconName(\"device-ipad-cog\")\n DEVICE_IPAD_DOLLAR: IconName = IconName(\"device-ipad-dollar\")\n DEVICE_IPAD_DOWN: IconName = IconName(\"device-ipad-down\")\n DEVICE_IPAD_EXCLAMATION: IconName = IconName(\"device-ipad-exclamation\")\n DEVICE_IPAD_HEART: IconName = IconName(\"device-ipad-heart\")\n DEVICE_IPAD_HORIZONTAL: IconName = IconName(\"device-ipad-horizontal\")\n DEVICE_IPAD_HORIZONTAL_BOLT: IconName = IconName(\"device-ipad-horizontal-bolt\")\n DEVICE_IPAD_HORIZONTAL_CANCEL: IconName = IconName(\"device-ipad-horizontal-cancel\")\n DEVICE_IPAD_HORIZONTAL_CHECK: IconName = IconName(\"device-ipad-horizontal-check\")\n DEVICE_IPAD_HORIZONTAL_CODE: IconName = IconName(\"device-ipad-horizontal-code\")\n DEVICE_IPAD_HORIZONTAL_COG: IconName = IconName(\"device-ipad-horizontal-cog\")\n DEVICE_IPAD_HORIZONTAL_DOLLAR: IconName = IconName(\"device-ipad-horizontal-dollar\")\n DEVICE_IPAD_HORIZONTAL_DOWN: IconName = IconName(\"device-ipad-horizontal-down\")\n DEVICE_IPAD_HORIZONTAL_EXCLAMATION: IconName = IconName(\n \"device-ipad-horizontal-exclamation\"\n )\n DEVICE_IPAD_HORIZONTAL_HEART: IconName = IconName(\"device-ipad-horizontal-heart\")\n DEVICE_IPAD_HORIZONTAL_MINUS: IconName = IconName(\"device-ipad-horizontal-minus\")\n DEVICE_IPAD_HORIZONTAL_OFF: IconName = IconName(\"device-ipad-horizontal-off\")\n DEVICE_IPAD_HORIZONTAL_PAUSE: IconName = IconName(\"device-ipad-horizontal-pause\")\n DEVICE_IPAD_HORIZONTAL_PIN: IconName = IconName(\"device-ipad-horizontal-pin\")\n DEVICE_IPAD_HORIZONTAL_PLUS: IconName = IconName(\"device-ipad-horizontal-plus\")\n DEVICE_IPAD_HORIZONTAL_QUESTION: IconName = IconName(\n \"device-ipad-horizontal-question\"\n )\n DEVICE_IPAD_HORIZONTAL_SEARCH: IconName = IconName(\"device-ipad-horizontal-search\")\n DEVICE_IPAD_HORIZONTAL_SHARE: IconName = IconName(\"device-ipad-horizontal-share\")\n DEVICE_IPAD_HORIZONTAL_STAR: IconName = IconName(\"device-ipad-horizontal-star\")\n DEVICE_IPAD_HORIZONTAL_UP: IconName = IconName(\"device-ipad-horizontal-up\")\n DEVICE_IPAD_HORIZONTAL_X: IconName = IconName(\"device-ipad-horizontal-x\")\n DEVICE_IPAD_MINUS: IconName = IconName(\"device-ipad-minus\")\n DEVICE_IPAD_OFF: IconName = IconName(\"device-ipad-off\")\n DEVICE_IPAD_PAUSE: IconName = IconName(\"device-ipad-pause\")\n DEVICE_IPAD_PIN: IconName = IconName(\"device-ipad-pin\")\n DEVICE_IPAD_PLUS: IconName = IconName(\"device-ipad-plus\")\n DEVICE_IPAD_QUESTION: IconName = IconName(\"device-ipad-question\")\n DEVICE_IPAD_SEARCH: IconName = IconName(\"device-ipad-search\")\n DEVICE_IPAD_SHARE: IconName = IconName(\"device-ipad-share\")\n DEVICE_IPAD_STAR: IconName = IconName(\"device-ipad-star\")\n DEVICE_IPAD_UP: IconName = IconName(\"device-ipad-up\")\n DEVICE_IPAD_X: IconName = IconName(\"device-ipad-x\")\n DEVICE_LANDLINE_PHONE: IconName = IconName(\"device-landline-phone\")\n DEVICE_LAPTOP: IconName = IconName(\"device-laptop\")\n DEVICE_LAPTOP_OFF: IconName = IconName(\"device-laptop-off\")\n DEVICE_MOBILE: IconName = IconName(\"device-mobile\")\n DEVICE_MOBILE_BOLT: IconName = IconName(\"device-mobile-bolt\")\n DEVICE_MOBILE_CANCEL: IconName = IconName(\"device-mobile-cancel\")\n DEVICE_MOBILE_CHARGING: IconName = IconName(\"device-mobile-charging\")\n DEVICE_MOBILE_CHECK: IconName = IconName(\"device-mobile-check\")\n DEVICE_MOBILE_CODE: IconName = IconName(\"device-mobile-code\")\n DEVICE_MOBILE_COG: IconName = IconName(\"device-mobile-cog\")\n DEVICE_MOBILE_DOLLAR: IconName = IconName(\"device-mobile-dollar\")\n DEVICE_MOBILE_DOWN: IconName = IconName(\"device-mobile-down\")\n DEVICE_MOBILE_EXCLAMATION: IconName = IconName(\"device-mobile-exclamation\")\n DEVICE_MOBILE_FILLED: IconName = IconName(\"device-mobile-filled\")\n DEVICE_MOBILE_HEART: IconName = IconName(\"device-mobile-heart\")\n DEVICE_MOBILE_MESSAGE: IconName = IconName(\"device-mobile-message\")\n DEVICE_MOBILE_MINUS: IconName = IconName(\"device-mobile-minus\")\n DEVICE_MOBILE_OFF: IconName = IconName(\"device-mobile-off\")\n DEVICE_MOBILE_PAUSE: IconName = IconName(\"device-mobile-pause\")\n DEVICE_MOBILE_PIN: IconName = IconName(\"device-mobile-pin\")\n DEVICE_MOBILE_PLUS: IconName = IconName(\"device-mobile-plus\")\n DEVICE_MOBILE_QUESTION: IconName = IconName(\"device-mobile-question\")\n DEVICE_MOBILE_ROTATED: IconName = IconName(\"device-mobile-rotated\")\n DEVICE_MOBILE_SEARCH: IconName = IconName(\"device-mobile-search\")\n DEVICE_MOBILE_SHARE: IconName = IconName(\"device-mobile-share\")\n DEVICE_MOBILE_STAR: IconName = IconName(\"device-mobile-star\")\n DEVICE_MOBILE_UP: IconName = IconName(\"device-mobile-up\")\n DEVICE_MOBILE_VIBRATION: IconName = IconName(\"device-mobile-vibration\")\n DEVICE_MOBILE_X: IconName = IconName(\"device-mobile-x\")\n DEVICE_NINTENDO: IconName = IconName(\"device-nintendo\")\n DEVICE_NINTENDO_OFF: IconName = IconName(\"device-nintendo-off\")\n DEVICE_REMOTE: IconName = IconName(\"device-remote\")\n DEVICE_SD_CARD: IconName = IconName(\"device-sd-card\")\n DEVICE_SIM: IconName = IconName(\"device-sim\")\n DEVICE_SIM_1: IconName = IconName(\"device-sim-1\")\n DEVICE_SIM_2: IconName = IconName(\"device-sim-2\")\n DEVICE_SIM_3: IconName = IconName(\"device-sim-3\")\n DEVICE_SPEAKER: IconName = IconName(\"device-speaker\")\n DEVICE_SPEAKER_OFF: IconName = IconName(\"device-speaker-off\")\n DEVICE_TABLET: IconName = IconName(\"device-tablet\")\n DEVICE_TABLET_BOLT: IconName = IconName(\"device-tablet-bolt\")\n DEVICE_TABLET_CANCEL: IconName = IconName(\"device-tablet-cancel\")\n DEVICE_TABLET_CHECK: IconName = IconName(\"device-tablet-check\")\n DEVICE_TABLET_CODE: IconName = IconName(\"device-tablet-code\")\n DEVICE_TABLET_COG: IconName = IconName(\"device-tablet-cog\")\n DEVICE_TABLET_DOLLAR: IconName = IconName(\"device-tablet-dollar\")\n DEVICE_TABLET_DOWN: IconName = IconName(\"device-tablet-down\")\n DEVICE_TABLET_EXCLAMATION: IconName = IconName(\"device-tablet-exclamation\")\n DEVICE_TABLET_FILLED: IconName = IconName(\"device-tablet-filled\")\n DEVICE_TABLET_HEART: IconName = IconName(\"device-tablet-heart\")\n DEVICE_TABLET_MINUS: IconName = IconName(\"device-tablet-minus\")\n DEVICE_TABLET_OFF: IconName = IconName(\"device-tablet-off\")\n DEVICE_TABLET_PAUSE: IconName = IconName(\"device-tablet-pause\")\n DEVICE_TABLET_PIN: IconName = IconName(\"device-tablet-pin\")\n DEVICE_TABLET_PLUS: IconName = IconName(\"device-tablet-plus\")\n DEVICE_TABLET_QUESTION: IconName = IconName(\"device-tablet-question\")\n DEVICE_TABLET_SEARCH: IconName = IconName(\"device-tablet-search\")\n DEVICE_TABLET_SHARE: IconName = IconName(\"device-tablet-share\")\n DEVICE_TABLET_STAR: IconName = IconName(\"device-tablet-star\")\n DEVICE_TABLET_UP: IconName = IconName(\"device-tablet-up\")\n DEVICE_TABLET_X: IconName = IconName(\"device-tablet-x\")\n DEVICE_TV: IconName = IconName(\"device-tv\")\n DEVICE_TV_OFF: IconName = IconName(\"device-tv-off\")\n DEVICE_TV_OLD: IconName = IconName(\"device-tv-old\")\n DEVICE_VISION_PRO: IconName = IconName(\"device-vision-pro\")\n DEVICE_WATCH: IconName = IconName(\"device-watch\")\n DEVICE_WATCH_BOLT: IconName = IconName(\"device-watch-bolt\")\n DEVICE_WATCH_CANCEL: IconName = IconName(\"device-watch-cancel\")\n DEVICE_WATCH_CHECK: IconName = IconName(\"device-watch-check\")\n DEVICE_WATCH_CODE: IconName = IconName(\"device-watch-code\")\n DEVICE_WATCH_COG: IconName = IconName(\"device-watch-cog\")\n DEVICE_WATCH_DOLLAR: IconName = IconName(\"device-watch-dollar\")\n DEVICE_WATCH_DOWN: IconName = IconName(\"device-watch-down\")\n DEVICE_WATCH_EXCLAMATION: IconName = IconName(\"device-watch-exclamation\")\n DEVICE_WATCH_HEART: IconName = IconName(\"device-watch-heart\")\n DEVICE_WATCH_MINUS: IconName = IconName(\"device-watch-minus\")\n DEVICE_WATCH_OFF: IconName = IconName(\"device-watch-off\")\n DEVICE_WATCH_PAUSE: IconName = IconName(\"device-watch-pause\")\n DEVICE_WATCH_PIN: IconName = IconName(\"device-watch-pin\")\n DEVICE_WATCH_PLUS: IconName = IconName(\"device-watch-plus\")\n DEVICE_WATCH_QUESTION: IconName = IconName(\"device-watch-question\")\n DEVICE_WATCH_SEARCH: IconName = IconName(\"device-watch-search\")\n DEVICE_WATCH_SHARE: IconName = IconName(\"device-watch-share\")\n DEVICE_WATCH_STAR: IconName = IconName(\"device-watch-star\")\n DEVICE_WATCH_STATS: IconName = IconName(\"device-watch-stats\")\n DEVICE_WATCH_STATS_2: IconName = IconName(\"device-watch-stats-2\")\n DEVICE_WATCH_UP: IconName = IconName(\"device-watch-up\")\n DEVICE_WATCH_X: IconName = IconName(\"device-watch-x\")\n DEVICES: IconName = IconName(\"devices\")\n DEVICES_2: IconName = IconName(\"devices-2\")\n DEVICES_BOLT: IconName = IconName(\"devices-bolt\")\n DEVICES_CANCEL: IconName = IconName(\"devices-cancel\")\n DEVICES_CHECK: IconName = IconName(\"devices-check\")\n DEVICES_CODE: IconName = IconName(\"devices-code\")\n DEVICES_COG: IconName = IconName(\"devices-cog\")\n DEVICES_DOLLAR: IconName = IconName(\"devices-dollar\")\n DEVICES_DOWN: IconName = IconName(\"devices-down\")\n DEVICES_EXCLAMATION: IconName = IconName(\"devices-exclamation\")\n DEVICES_HEART: IconName = IconName(\"devices-heart\")\n DEVICES_MINUS: IconName = IconName(\"devices-minus\")\n DEVICES_OFF: IconName = IconName(\"devices-off\")\n DEVICES_PAUSE: IconName = IconName(\"devices-pause\")\n DEVICES_PC: IconName = IconName(\"devices-pc\")\n DEVICES_PC_OFF: IconName = IconName(\"devices-pc-off\")\n DEVICES_PIN: IconName = IconName(\"devices-pin\")\n DEVICES_PLUS: IconName = IconName(\"devices-plus\")\n DEVICES_QUESTION: IconName = IconName(\"devices-question\")\n DEVICES_SEARCH: IconName = IconName(\"devices-search\")\n DEVICES_SHARE: IconName = IconName(\"devices-share\")\n DEVICES_STAR: IconName = IconName(\"devices-star\")\n DEVICES_UP: IconName = IconName(\"devices-up\")\n DEVICES_X: IconName = IconName(\"devices-x\")\n DIABOLO: IconName = IconName(\"diabolo\")\n DIABOLO_OFF: IconName = IconName(\"diabolo-off\")\n DIABOLO_PLUS: IconName = IconName(\"diabolo-plus\")\n DIALPAD: IconName = IconName(\"dialpad\")\n DIALPAD_FILLED: IconName = IconName(\"dialpad-filled\")\n DIALPAD_OFF: IconName = IconName(\"dialpad-off\")\n DIAMOND: IconName = IconName(\"diamond\")\n DIAMOND_FILLED: IconName = IconName(\"diamond-filled\")\n DIAMOND_OFF: IconName = IconName(\"diamond-off\")\n DIAMONDS: IconName = IconName(\"diamonds\")\n DIAMONDS_FILLED: IconName = IconName(\"diamonds-filled\")\n DICE: IconName = IconName(\"dice\")\n DICE_1: IconName = IconName(\"dice-1\")\n DICE_1_FILLED: IconName = IconName(\"dice-1-filled\")\n DICE_2: IconName = IconName(\"dice-2\")\n DICE_2_FILLED: IconName = IconName(\"dice-2-filled\")\n DICE_3: IconName = IconName(\"dice-3\")\n DICE_3_FILLED: IconName = IconName(\"dice-3-filled\")\n DICE_4: IconName = IconName(\"dice-4\")\n DICE_4_FILLED: IconName = IconName(\"dice-4-filled\")\n DICE_5: IconName = IconName(\"dice-5\")\n DICE_5_FILLED: IconName = IconName(\"dice-5-filled\")\n DICE_6: IconName = IconName(\"dice-6\")\n DICE_6_FILLED: IconName = IconName(\"dice-6-filled\")\n DICE_FILLED: IconName = IconName(\"dice-filled\")\n DIMENSIONS: IconName = IconName(\"dimensions\")\n DIRECTION: IconName = IconName(\"direction\")\n DIRECTION_HORIZONTAL: IconName = IconName(\"direction-horizontal\")\n DIRECTION_SIGN: IconName = IconName(\"direction-sign\")\n DIRECTION_SIGN_FILLED: IconName = IconName(\"direction-sign-filled\")\n DIRECTION_SIGN_OFF: IconName = IconName(\"direction-sign-off\")\n DIRECTIONS: IconName = IconName(\"directions\")\n DIRECTIONS_OFF: IconName = IconName(\"directions-off\")\n DISABLED: IconName = IconName(\"disabled\")\n DISABLED_2: IconName = IconName(\"disabled-2\")\n DISABLED_OFF: IconName = IconName(\"disabled-off\")\n DISC: IconName = IconName(\"disc\")\n DISC_GOLF: IconName = IconName(\"disc-golf\")\n DISC_OFF: IconName = IconName(\"disc-off\")\n DISCOUNT: IconName = IconName(\"discount\")\n DISCOUNT_2: IconName = IconName(\"discount-2\")\n DISCOUNT_2_OFF: IconName = IconName(\"discount-2-off\")\n DISCOUNT_CHECK: IconName = IconName(\"discount-check\")\n DISCOUNT_CHECK_FILLED: IconName = IconName(\"discount-check-filled\")\n DISCOUNT_OFF: IconName = IconName(\"discount-off\")\n DIVIDE: IconName = IconName(\"divide\")\n DNA: IconName = IconName(\"dna\")\n DNA_2: IconName = IconName(\"dna-2\")\n DNA_2_OFF: IconName = IconName(\"dna-2-off\")\n DNA_OFF: IconName = IconName(\"dna-off\")\n DOG: IconName = IconName(\"dog\")\n DOG_BOWL: IconName = IconName(\"dog-bowl\")\n DOOR: IconName = IconName(\"door\")\n DOOR_ENTER: IconName = IconName(\"door-enter\")\n DOOR_EXIT: IconName = IconName(\"door-exit\")\n DOOR_OFF: IconName = IconName(\"door-off\")\n DOTS: IconName = IconName(\"dots\")\n DOTS_CIRCLE_HORIZONTAL: IconName = IconName(\"dots-circle-horizontal\")\n DOTS_DIAGONAL: IconName = IconName(\"dots-diagonal\")\n DOTS_DIAGONAL_2: IconName = IconName(\"dots-diagonal-2\")\n DOTS_VERTICAL: IconName = IconName(\"dots-vertical\")\n DOWNLOAD: IconName = IconName(\"download\")\n DOWNLOAD_OFF: IconName = IconName(\"download-off\")\n DRAG_DROP: IconName = IconName(\"drag-drop\")\n DRAG_DROP_2: IconName = IconName(\"drag-drop-2\")\n DRONE: IconName = IconName(\"drone\")\n DRONE_OFF: IconName = IconName(\"drone-off\")\n DROP_CIRCLE: IconName = IconName(\"drop-circle\")\n DROPLET: IconName = IconName(\"droplet\")\n DROPLET_BOLT: IconName = IconName(\"droplet-bolt\")\n DROPLET_CANCEL: IconName = IconName(\"droplet-cancel\")\n DROPLET_CHECK: IconName = IconName(\"droplet-check\")\n DROPLET_CODE: IconName = IconName(\"droplet-code\")\n DROPLET_COG: IconName = IconName(\"droplet-cog\")\n DROPLET_DOLLAR: IconName = IconName(\"droplet-dollar\")\n DROPLET_DOWN: IconName = IconName(\"droplet-down\")\n DROPLET_EXCLAMATION: IconName = IconName(\"droplet-exclamation\")\n DROPLET_FILLED: IconName = IconName(\"droplet-filled\")\n DROPLET_FILLED_2: IconName = IconName(\"droplet-filled-2\")\n DROPLET_HALF: IconName = IconName(\"droplet-half\")\n DROPLET_HALF_2: IconName = IconName(\"droplet-half-2\")\n DROPLET_HALF_FILLED: IconName = IconName(\"droplet-half-filled\")\n DROPLET_HEART: IconName = IconName(\"droplet-heart\")\n DROPLET_MINUS: IconName = IconName(\"droplet-minus\")\n DROPLET_OFF: IconName = IconName(\"droplet-off\")\n DROPLET_PAUSE: IconName = IconName(\"droplet-pause\")\n DROPLET_PIN: IconName = IconName(\"droplet-pin\")\n DROPLET_PLUS: IconName = IconName(\"droplet-plus\")\n DROPLET_QUESTION: IconName = IconName(\"droplet-question\")\n DROPLET_SEARCH: IconName = IconName(\"droplet-search\")\n DROPLET_SHARE: IconName = IconName(\"droplet-share\")\n DROPLET_STAR: IconName = IconName(\"droplet-star\")\n DROPLET_UP: IconName = IconName(\"droplet-up\")\n DROPLET_X: IconName = IconName(\"droplet-x\")\n DUAL_SCREEN: IconName = IconName(\"dual-screen\")\n E_PASSPORT: IconName = IconName(\"e-passport\")\n EAR: IconName = IconName(\"ear\")\n EAR_OFF: IconName = IconName(\"ear-off\")\n EASE_IN: IconName = IconName(\"ease-in\")\n EASE_IN_CONTROL_POINT: IconName = IconName(\"ease-in-control-point\")\n EASE_IN_OUT: IconName = IconName(\"ease-in-out\")\n EASE_IN_OUT_CONTROL_POINTS: IconName = IconName(\"ease-in-out-control-points\")\n EASE_OUT: IconName = IconName(\"ease-out\")\n EASE_OUT_CONTROL_POINT: IconName = IconName(\"ease-out-control-point\")\n EDIT: IconName = IconName(\"edit\")\n EDIT_CIRCLE: IconName = IconName(\"edit-circle\")\n EDIT_CIRCLE_OFF: IconName = IconName(\"edit-circle-off\")\n EDIT_OFF: IconName = IconName(\"edit-off\")\n EGG: IconName = IconName(\"egg\")\n EGG_CRACKED: IconName = IconName(\"egg-cracked\")\n EGG_FILLED: IconName = IconName(\"egg-filled\")\n EGG_FRIED: IconName = IconName(\"egg-fried\")\n EGG_OFF: IconName = IconName(\"egg-off\")\n EGGS: IconName = IconName(\"eggs\")\n ELEVATOR: IconName = IconName(\"elevator\")\n ELEVATOR_OFF: IconName = IconName(\"elevator-off\")\n EMERGENCY_BED: IconName = IconName(\"emergency-bed\")\n EMPATHIZE: IconName = IconName(\"empathize\")\n EMPATHIZE_OFF: IconName = IconName(\"empathize-off\")\n EMPHASIS: IconName = IconName(\"emphasis\")\n ENGINE: IconName = IconName(\"engine\")\n ENGINE_OFF: IconName = IconName(\"engine-off\")\n EQUAL: IconName = IconName(\"equal\")\n EQUAL_DOUBLE: IconName = IconName(\"equal-double\")\n EQUAL_NOT: IconName = IconName(\"equal-not\")\n ERASER: IconName = IconName(\"eraser\")\n ERASER_OFF: IconName = IconName(\"eraser-off\")\n ERROR_404: IconName = IconName(\"error-404\")\n ERROR_404_OFF: IconName = IconName(\"error-404-off\")\n EXCHANGE: IconName = IconName(\"exchange\")\n EXCHANGE_OFF: IconName = IconName(\"exchange-off\")\n EXCLAMATION_CIRCLE: IconName = IconName(\"exclamation-circle\")\n EXCLAMATION_MARK: IconName = IconName(\"exclamation-mark\")\n EXCLAMATION_MARK_OFF: IconName = IconName(\"exclamation-mark-off\")\n EXPLICIT: IconName = IconName(\"explicit\")\n EXPLICIT_OFF: IconName = IconName(\"explicit-off\")\n EXPOSURE: IconName = IconName(\"exposure\")\n EXPOSURE_0: IconName = IconName(\"exposure-0\")\n EXPOSURE_MINUS_1: IconName = IconName(\"exposure-minus-1\")\n EXPOSURE_MINUS_2: IconName = IconName(\"exposure-minus-2\")\n EXPOSURE_OFF: IconName = IconName(\"exposure-off\")\n EXPOSURE_PLUS_1: IconName = IconName(\"exposure-plus-1\")\n EXPOSURE_PLUS_2: IconName = IconName(\"exposure-plus-2\")\n EXTERNAL_LINK: IconName = IconName(\"external-link\")\n EXTERNAL_LINK_OFF: IconName = IconName(\"external-link-off\")\n EYE: IconName = IconName(\"eye\")\n EYE_CHECK: IconName = IconName(\"eye-check\")\n EYE_CLOSED: IconName = IconName(\"eye-closed\")\n EYE_COG: IconName = IconName(\"eye-cog\")\n EYE_EDIT: IconName = IconName(\"eye-edit\")\n EYE_EXCLAMATION: IconName = IconName(\"eye-exclamation\")\n EYE_FILLED: IconName = IconName(\"eye-filled\")\n EYE_HEART: IconName = IconName(\"eye-heart\")\n EYE_OFF: IconName = IconName(\"eye-off\")\n EYE_TABLE: IconName = IconName(\"eye-table\")\n EYE_X: IconName = IconName(\"eye-x\")\n EYEGLASS: IconName = IconName(\"eyeglass\")\n EYEGLASS_2: IconName = IconName(\"eyeglass-2\")\n EYEGLASS_OFF: IconName = IconName(\"eyeglass-off\")\n FACE_ID: IconName = IconName(\"face-id\")\n FACE_ID_ERROR: IconName = IconName(\"face-id-error\")\n FACE_MASK: IconName = IconName(\"face-mask\")\n FACE_MASK_OFF: IconName = IconName(\"face-mask-off\")\n FALL: IconName = IconName(\"fall\")\n FEATHER: IconName = IconName(\"feather\")\n FEATHER_OFF: IconName = IconName(\"feather-off\")\n FENCE: IconName = IconName(\"fence\")\n FENCE_OFF: IconName = IconName(\"fence-off\")\n FIDGET_SPINNER: IconName = IconName(\"fidget-spinner\")\n FILE: IconName = IconName(\"file\")\n FILE_3D: IconName = IconName(\"file-3d\")\n FILE_ALERT: IconName = IconName(\"file-alert\")\n FILE_ANALYTICS: IconName = IconName(\"file-analytics\")\n FILE_ARROW_LEFT: IconName = IconName(\"file-arrow-left\")\n FILE_ARROW_RIGHT: IconName = IconName(\"file-arrow-right\")\n FILE_BARCODE: IconName = IconName(\"file-barcode\")\n FILE_BROKEN: IconName = IconName(\"file-broken\")\n FILE_CERTIFICATE: IconName = IconName(\"file-certificate\")\n FILE_CHART: IconName = IconName(\"file-chart\")\n FILE_CHECK: IconName = IconName(\"file-check\")\n FILE_CODE: IconName = IconName(\"file-code\")\n FILE_CODE_2: IconName = IconName(\"file-code-2\")\n FILE_CV: IconName = IconName(\"file-cv\")\n FILE_DATABASE: IconName = IconName(\"file-database\")\n FILE_DELTA: IconName = IconName(\"file-delta\")\n FILE_DESCRIPTION: IconName = IconName(\"file-description\")\n FILE_DIFF: IconName = IconName(\"file-diff\")\n FILE_DIGIT: IconName = IconName(\"file-digit\")\n FILE_DISLIKE: IconName = IconName(\"file-dislike\")\n FILE_DOLLAR: IconName = IconName(\"file-dollar\")\n FILE_DOTS: IconName = IconName(\"file-dots\")\n FILE_DOWNLOAD: IconName = IconName(\"file-download\")\n FILE_EURO: IconName = IconName(\"file-euro\")\n FILE_EXPORT: IconName = IconName(\"file-export\")\n FILE_FILLED: IconName = IconName(\"file-filled\")\n FILE_FUNCTION: IconName = IconName(\"file-function\")\n FILE_HORIZONTAL: IconName = IconName(\"file-horizontal\")\n FILE_IMPORT: IconName = IconName(\"file-import\")\n FILE_INFINITY: IconName = IconName(\"file-infinity\")\n FILE_INFO: IconName = IconName(\"file-info\")\n FILE_INVOICE: IconName = IconName(\"file-invoice\")\n FILE_LAMBDA: IconName = IconName(\"file-lambda\")\n FILE_LIKE: IconName = IconName(\"file-like\")\n FILE_MINUS: IconName = IconName(\"file-minus\")\n FILE_MUSIC: IconName = IconName(\"file-music\")\n FILE_OFF: IconName = IconName(\"file-off\")\n FILE_ORIENTATION: IconName = IconName(\"file-orientation\")\n FILE_PENCIL: IconName = IconName(\"file-pencil\")\n FILE_PERCENT: IconName = IconName(\"file-percent\")\n FILE_PHONE: IconName = IconName(\"file-phone\")\n FILE_PLUS: IconName = IconName(\"file-plus\")\n FILE_POWER: IconName = IconName(\"file-power\")\n FILE_REPORT: IconName = IconName(\"file-report\")\n FILE_RSS: IconName = IconName(\"file-rss\")\n FILE_SCISSORS: IconName = IconName(\"file-scissors\")\n FILE_SEARCH: IconName = IconName(\"file-search\")\n FILE_SETTINGS: IconName = IconName(\"file-settings\")\n FILE_SHREDDER: IconName = IconName(\"file-shredder\")\n FILE_SIGNAL: IconName = IconName(\"file-signal\")\n FILE_SPREADSHEET: IconName = IconName(\"file-spreadsheet\")\n FILE_STACK: IconName = IconName(\"file-stack\")\n FILE_STAR: IconName = IconName(\"file-star\")\n FILE_SYMLINK: IconName = IconName(\"file-symlink\")\n FILE_TEXT: IconName = IconName(\"file-text\")\n FILE_TEXT_AI: IconName = IconName(\"file-text-ai\")\n FILE_TIME: IconName = IconName(\"file-time\")\n FILE_TYPOGRAPHY: IconName = IconName(\"file-typography\")\n FILE_UNKNOWN: IconName = IconName(\"file-unknown\")\n FILE_UPLOAD: IconName = IconName(\"file-upload\")\n FILE_VECTOR: IconName = IconName(\"file-vector\")\n FILE_X: IconName = IconName(\"file-x\")\n FILE_X_FILLED: IconName = IconName(\"file-x-filled\")\n FILE_ZIP: IconName = IconName(\"file-zip\")\n FILES: IconName = IconName(\"files\")\n FILES_OFF: IconName = IconName(\"files-off\")\n FILTER: IconName = IconName(\"filter\")\n FILTER_COG: IconName = IconName(\"filter-cog\")\n FILTER_DOLLAR: IconName = IconName(\"filter-dollar\")\n FILTER_EDIT: IconName = IconName(\"filter-edit\")\n FILTER_MINUS: IconName = IconName(\"filter-minus\")\n FILTER_OFF: IconName = IconName(\"filter-off\")\n FILTER_PLUS: IconName = IconName(\"filter-plus\")\n FILTER_STAR: IconName = IconName(\"filter-star\")\n FILTER_X: IconName = IconName(\"filter-x\")\n FILTERS: IconName = IconName(\"filters\")\n FINGERPRINT: IconName = IconName(\"fingerprint\")\n FINGERPRINT_OFF: IconName = IconName(\"fingerprint-off\")\n FIRE_EXTINGUISHER: IconName = IconName(\"fire-extinguisher\")\n FIRE_HYDRANT: IconName = IconName(\"fire-hydrant\")\n FIRE_HYDRANT_OFF: IconName = IconName(\"fire-hydrant-off\")\n FIRETRUCK: IconName = IconName(\"firetruck\")\n FIRST_AID_KIT: IconName = IconName(\"first-aid-kit\")\n FIRST_AID_KIT_OFF: IconName = IconName(\"first-aid-kit-off\")\n FISH: IconName = IconName(\"fish\")\n FISH_BONE: IconName = IconName(\"fish-bone\")\n FISH_CHRISTIANITY: IconName = IconName(\"fish-christianity\")\n FISH_HOOK: IconName = IconName(\"fish-hook\")\n FISH_HOOK_OFF: IconName = IconName(\"fish-hook-off\")\n FISH_OFF: IconName = IconName(\"fish-off\")\n FLAG: IconName = IconName(\"flag\")\n FLAG_2: IconName = IconName(\"flag-2\")\n FLAG_2_FILLED: IconName = IconName(\"flag-2-filled\")\n FLAG_2_OFF: IconName = IconName(\"flag-2-off\")\n FLAG_3: IconName = IconName(\"flag-3\")\n FLAG_3_FILLED: IconName = IconName(\"flag-3-filled\")\n FLAG_FILLED: IconName = IconName(\"flag-filled\")\n FLAG_OFF: IconName = IconName(\"flag-off\")\n FLAME: IconName = IconName(\"flame\")\n FLAME_OFF: IconName = IconName(\"flame-off\")\n FLARE: IconName = IconName(\"flare\")\n FLASK: IconName = IconName(\"flask\")\n FLASK_2: IconName = IconName(\"flask-2\")\n FLASK_2_OFF: IconName = IconName(\"flask-2-off\")\n FLASK_OFF: IconName = IconName(\"flask-off\")\n FLIP_FLOPS: IconName = IconName(\"flip-flops\")\n FLIP_HORIZONTAL: IconName = IconName(\"flip-horizontal\")\n FLIP_VERTICAL: IconName = IconName(\"flip-vertical\")\n FLOAT_CENTER: IconName = IconName(\"float-center\")\n FLOAT_LEFT: IconName = IconName(\"float-left\")\n FLOAT_NONE: IconName = IconName(\"float-none\")\n FLOAT_RIGHT: IconName = IconName(\"float-right\")\n FLOWER: IconName = IconName(\"flower\")\n FLOWER_OFF: IconName = IconName(\"flower-off\")\n FOCUS: IconName = IconName(\"focus\")\n FOCUS_2: IconName = IconName(\"focus-2\")\n FOCUS_AUTO: IconName = IconName(\"focus-auto\")\n FOCUS_CENTERED: IconName = IconName(\"focus-centered\")\n FOLD: IconName = IconName(\"fold\")\n FOLD_DOWN: IconName = IconName(\"fold-down\")\n FOLD_UP: IconName = IconName(\"fold-up\")\n FOLDER: IconName = IconName(\"folder\")\n FOLDER_BOLT: IconName = IconName(\"folder-bolt\")\n FOLDER_CANCEL: IconName = IconName(\"folder-cancel\")\n FOLDER_CHECK: IconName = IconName(\"folder-check\")\n FOLDER_CODE: IconName = IconName(\"folder-code\")\n FOLDER_COG: IconName = IconName(\"folder-cog\")\n FOLDER_DOLLAR: IconName = IconName(\"folder-dollar\")\n FOLDER_DOWN: IconName = IconName(\"folder-down\")\n FOLDER_EXCLAMATION: IconName = IconName(\"folder-exclamation\")\n FOLDER_FILLED: IconName = IconName(\"folder-filled\")\n FOLDER_HEART: IconName = IconName(\"folder-heart\")\n FOLDER_MINUS: IconName = IconName(\"folder-minus\")\n FOLDER_OFF: IconName = IconName(\"folder-off\")\n FOLDER_OPEN: IconName = IconName(\"folder-open\")\n FOLDER_PAUSE: IconName = IconName(\"folder-pause\")\n FOLDER_PIN: IconName = IconName(\"folder-pin\")\n FOLDER_PLUS: IconName = IconName(\"folder-plus\")\n FOLDER_QUESTION: IconName = IconName(\"folder-question\")\n FOLDER_SEARCH: IconName = IconName(\"folder-search\")\n FOLDER_SHARE: IconName = IconName(\"folder-share\")\n FOLDER_STAR: IconName = IconName(\"folder-star\")\n FOLDER_SYMLINK: IconName = IconName(\"folder-symlink\")\n FOLDER_UP: IconName = IconName(\"folder-up\")\n FOLDER_X: IconName = IconName(\"folder-x\")\n FOLDERS: IconName = IconName(\"folders\")\n FOLDERS_OFF: IconName = IconName(\"folders-off\")\n FORBID: IconName = IconName(\"forbid\")\n FORBID_2: IconName = IconName(\"forbid-2\")\n FORKLIFT: IconName = IconName(\"forklift\")\n FORMS: IconName = IconName(\"forms\")\n FOUNTAIN: IconName = IconName(\"fountain\")\n FOUNTAIN_OFF: IconName = IconName(\"fountain-off\")\n FRAME: IconName = IconName(\"frame\")\n FRAME_OFF: IconName = IconName(\"frame-off\")\n FREE_RIGHTS: IconName = IconName(\"free-rights\")\n FREEZE_COLUMN: IconName = IconName(\"freeze-column\")\n FREEZE_ROW: IconName = IconName(\"freeze-row\")\n FREEZE_ROW_COLUMN: IconName = IconName(\"freeze-row-column\")\n FRIDGE: IconName = IconName(\"fridge\")\n FRIDGE_OFF: IconName = IconName(\"fridge-off\")\n FRIENDS: IconName = IconName(\"friends\")\n FRIENDS_OFF: IconName = IconName(\"friends-off\")\n FRUSTUM: IconName = IconName(\"frustum\")\n FRUSTUM_OFF: IconName = IconName(\"frustum-off\")\n FRUSTUM_PLUS: IconName = IconName(\"frustum-plus\")\n FUNCTION: IconName = IconName(\"function\")\n FUNCTION_OFF: IconName = IconName(\"function-off\")\n GARDEN_CART: IconName = IconName(\"garden-cart\")\n GARDEN_CART_OFF: IconName = IconName(\"garden-cart-off\")\n GAS_STATION: IconName = IconName(\"gas-station\")\n GAS_STATION_OFF: IconName = IconName(\"gas-station-off\")\n GAUGE: IconName = IconName(\"gauge\")\n GAUGE_OFF: IconName = IconName(\"gauge-off\")\n GAVEL: IconName = IconName(\"gavel\")\n GENDER_AGENDER: IconName = IconName(\"gender-agender\")\n GENDER_ANDROGYNE: IconName = IconName(\"gender-androgyne\")\n GENDER_BIGENDER: IconName = IconName(\"gender-bigender\")\n GENDER_DEMIBOY: IconName = IconName(\"gender-demiboy\")\n GENDER_DEMIGIRL: IconName = IconName(\"gender-demigirl\")\n GENDER_EPICENE: IconName = IconName(\"gender-epicene\")\n GENDER_FEMALE: IconName = IconName(\"gender-female\")\n GENDER_FEMME: IconName = IconName(\"gender-femme\")\n GENDER_GENDERFLUID: IconName = IconName(\"gender-genderfluid\")\n GENDER_GENDERLESS: IconName = IconName(\"gender-genderless\")\n GENDER_GENDERQUEER: IconName = IconName(\"gender-genderqueer\")\n GENDER_HERMAPHRODITE: IconName = IconName(\"gender-hermaphrodite\")\n GENDER_INTERGENDER: IconName = IconName(\"gender-intergender\")\n GENDER_MALE: IconName = IconName(\"gender-male\")\n GENDER_NEUTROIS: IconName = IconName(\"gender-neutrois\")\n GENDER_THIRD: IconName = IconName(\"gender-third\")\n GENDER_TRANSGENDER: IconName = IconName(\"gender-transgender\")\n GENDER_TRASVESTI: IconName = IconName(\"gender-trasvesti\")\n GEOMETRY: IconName = IconName(\"geometry\")\n GHOST: IconName = IconName(\"ghost\")\n GHOST_2: IconName = IconName(\"ghost-2\")\n GHOST_2_FILLED: IconName = IconName(\"ghost-2-filled\")\n GHOST_FILLED: IconName = IconName(\"ghost-filled\")\n GHOST_OFF: IconName = IconName(\"ghost-off\")\n GIF: IconName = IconName(\"gif\")\n GIFT: IconName = IconName(\"gift\")\n GIFT_CARD: IconName = IconName(\"gift-card\")\n GIFT_OFF: IconName = IconName(\"gift-off\")\n GIT_BRANCH: IconName = IconName(\"git-branch\")\n GIT_BRANCH_DELETED: IconName = IconName(\"git-branch-deleted\")\n GIT_CHERRY_PICK: IconName = IconName(\"git-cherry-pick\")\n GIT_COMMIT: IconName = IconName(\"git-commit\")\n GIT_COMPARE: IconName = IconName(\"git-compare\")\n GIT_FORK: IconName = IconName(\"git-fork\")\n GIT_MERGE: IconName = IconName(\"git-merge\")\n GIT_PULL_REQUEST: IconName = IconName(\"git-pull-request\")\n GIT_PULL_REQUEST_CLOSED: IconName = IconName(\"git-pull-request-closed\")\n GIT_PULL_REQUEST_DRAFT: IconName = IconName(\"git-pull-request-draft\")\n GIZMO: IconName = IconName(\"gizmo\")\n GLASS: IconName = IconName(\"glass\")\n GLASS_FULL: IconName = IconName(\"glass-full\")\n GLASS_OFF: IconName = IconName(\"glass-off\")\n GLOBE: IconName = IconName(\"globe\")\n GLOBE_OFF: IconName = IconName(\"globe-off\")\n GO_GAME: IconName = IconName(\"go-game\")\n GOLF: IconName = IconName(\"golf\")\n GOLF_OFF: IconName = IconName(\"golf-off\")\n GPS: IconName = IconName(\"gps\")\n GRADIENTER: IconName = IconName(\"gradienter\")\n GRAIN: IconName = IconName(\"grain\")\n GRAPH: IconName = IconName(\"graph\")\n GRAPH_OFF: IconName = IconName(\"graph-off\")\n GRAVE: IconName = IconName(\"grave\")\n GRAVE_2: IconName = IconName(\"grave-2\")\n GRID_DOTS: IconName = IconName(\"grid-dots\")\n GRID_PATTERN: IconName = IconName(\"grid-pattern\")\n GRILL: IconName = IconName(\"grill\")\n GRILL_FORK: IconName = IconName(\"grill-fork\")\n GRILL_OFF: IconName = IconName(\"grill-off\")\n GRILL_SPATULA: IconName = IconName(\"grill-spatula\")\n GRIP_HORIZONTAL: IconName = IconName(\"grip-horizontal\")\n GRIP_VERTICAL: IconName = IconName(\"grip-vertical\")\n GROWTH: IconName = IconName(\"growth\")\n GUITAR_PICK: IconName = IconName(\"guitar-pick\")\n GUITAR_PICK_FILLED: IconName = IconName(\"guitar-pick-filled\")\n H_1: IconName = IconName(\"h-1\")\n H_2: IconName = IconName(\"h-2\")\n H_3: IconName = IconName(\"h-3\")\n H_4: IconName = IconName(\"h-4\")\n H_5: IconName = IconName(\"h-5\")\n H_6: IconName = IconName(\"h-6\")\n HAMMER: IconName = IconName(\"hammer\")\n HAMMER_OFF: IconName = IconName(\"hammer-off\")\n HAND_CLICK: IconName = IconName(\"hand-click\")\n HAND_FINGER: IconName = IconName(\"hand-finger\")\n HAND_FINGER_OFF: IconName = IconName(\"hand-finger-off\")\n HAND_GRAB: IconName = IconName(\"hand-grab\")\n HAND_LITTLE_FINGER: IconName = IconName(\"hand-little-finger\")\n HAND_MIDDLE_FINGER: IconName = IconName(\"hand-middle-finger\")\n HAND_MOVE: IconName = IconName(\"hand-move\")\n HAND_OFF: IconName = IconName(\"hand-off\")\n HAND_RING_FINGER: IconName = IconName(\"hand-ring-finger\")\n HAND_ROCK: IconName = IconName(\"hand-rock\")\n HAND_SANITIZER: IconName = IconName(\"hand-sanitizer\")\n HAND_STOP: IconName = IconName(\"hand-stop\")\n HAND_THREE_FINGERS: IconName = IconName(\"hand-three-fingers\")\n HAND_TWO_FINGERS: IconName = IconName(\"hand-two-fingers\")\n HANGER: IconName = IconName(\"hanger\")\n HANGER_2: IconName = IconName(\"hanger-2\")\n HANGER_OFF: IconName = IconName(\"hanger-off\")\n HASH: IconName = IconName(\"hash\")\n HAZE: IconName = IconName(\"haze\")\n HAZE_MOON: IconName = IconName(\"haze-moon\")\n HDR: IconName = IconName(\"hdr\")\n HEADING: IconName = IconName(\"heading\")\n HEADING_OFF: IconName = IconName(\"heading-off\")\n HEADPHONES: IconName = IconName(\"headphones\")\n HEADPHONES_FILLED: IconName = IconName(\"headphones-filled\")\n HEADPHONES_OFF: IconName = IconName(\"headphones-off\")\n HEADSET: IconName = IconName(\"headset\")\n HEADSET_OFF: IconName = IconName(\"headset-off\")\n HEALTH_RECOGNITION: IconName = IconName(\"health-recognition\")\n HEART: IconName = IconName(\"heart\")\n HEART_BROKEN: IconName = IconName(\"heart-broken\")\n HEART_FILLED: IconName = IconName(\"heart-filled\")\n HEART_HANDSHAKE: IconName = IconName(\"heart-handshake\")\n HEART_MINUS: IconName = IconName(\"heart-minus\")\n HEART_OFF: IconName = IconName(\"heart-off\")\n HEART_PLUS: IconName = IconName(\"heart-plus\")\n HEART_RATE_MONITOR: IconName = IconName(\"heart-rate-monitor\")\n HEARTBEAT: IconName = IconName(\"heartbeat\")\n HEARTS: IconName = IconName(\"hearts\")\n HEARTS_OFF: IconName = IconName(\"hearts-off\")\n HELICOPTER: IconName = IconName(\"helicopter\")\n HELICOPTER_LANDING: IconName = IconName(\"helicopter-landing\")\n HELMET: IconName = IconName(\"helmet\")\n HELMET_OFF: IconName = IconName(\"helmet-off\")\n HELP: IconName = IconName(\"help\")\n HELP_CIRCLE: IconName = IconName(\"help-circle\")\n HELP_CIRCLE_FILLED: IconName = IconName(\"help-circle-filled\")\n HELP_HEXAGON: IconName = IconName(\"help-hexagon\")\n HELP_HEXAGON_FILLED: IconName = IconName(\"help-hexagon-filled\")\n HELP_OCTAGON: IconName = IconName(\"help-octagon\")\n HELP_OCTAGON_FILLED: IconName = IconName(\"help-octagon-filled\")\n HELP_OFF: IconName = IconName(\"help-off\")\n HELP_SMALL: IconName = IconName(\"help-small\")\n HELP_SQUARE: IconName = IconName(\"help-square\")\n HELP_SQUARE_FILLED: IconName = IconName(\"help-square-filled\")\n HELP_SQUARE_ROUNDED: IconName = IconName(\"help-square-rounded\")\n HELP_SQUARE_ROUNDED_FILLED: IconName = IconName(\"help-square-rounded-filled\")\n HELP_TRIANGLE: IconName = IconName(\"help-triangle\")\n HELP_TRIANGLE_FILLED: IconName = IconName(\"help-triangle-filled\")\n HEMISPHERE: IconName = IconName(\"hemisphere\")\n HEMISPHERE_OFF: IconName = IconName(\"hemisphere-off\")\n HEMISPHERE_PLUS: IconName = IconName(\"hemisphere-plus\")\n HEXAGON: IconName = IconName(\"hexagon\")\n HEXAGON_0_FILLED: IconName = IconName(\"hexagon-0-filled\")\n HEXAGON_1_FILLED: IconName = IconName(\"hexagon-1-filled\")\n HEXAGON_2_FILLED: IconName = IconName(\"hexagon-2-filled\")\n HEXAGON_3_FILLED: IconName = IconName(\"hexagon-3-filled\")\n HEXAGON_3D: IconName = IconName(\"hexagon-3d\")\n HEXAGON_4_FILLED: IconName = IconName(\"hexagon-4-filled\")\n HEXAGON_5_FILLED: IconName = IconName(\"hexagon-5-filled\")\n HEXAGON_6_FILLED: IconName = IconName(\"hexagon-6-filled\")\n HEXAGON_7_FILLED: IconName = IconName(\"hexagon-7-filled\")\n HEXAGON_8_FILLED: IconName = IconName(\"hexagon-8-filled\")\n HEXAGON_9_FILLED: IconName = IconName(\"hexagon-9-filled\")\n HEXAGON_FILLED: IconName = IconName(\"hexagon-filled\")\n HEXAGON_LETTER_A: IconName = IconName(\"hexagon-letter-a\")\n HEXAGON_LETTER_B: IconName = IconName(\"hexagon-letter-b\")\n HEXAGON_LETTER_C: IconName = IconName(\"hexagon-letter-c\")\n HEXAGON_LETTER_D: IconName = IconName(\"hexagon-letter-d\")\n HEXAGON_LETTER_E: IconName = IconName(\"hexagon-letter-e\")\n HEXAGON_LETTER_F: IconName = IconName(\"hexagon-letter-f\")\n HEXAGON_LETTER_G: IconName = IconName(\"hexagon-letter-g\")\n HEXAGON_LETTER_H: IconName = IconName(\"hexagon-letter-h\")\n HEXAGON_LETTER_I: IconName = IconName(\"hexagon-letter-i\")\n HEXAGON_LETTER_J: IconName = IconName(\"hexagon-letter-j\")\n HEXAGON_LETTER_K: IconName = IconName(\"hexagon-letter-k\")\n HEXAGON_LETTER_L: IconName = IconName(\"hexagon-letter-l\")\n HEXAGON_LETTER_M: IconName = IconName(\"hexagon-letter-m\")\n HEXAGON_LETTER_N: IconName = IconName(\"hexagon-letter-n\")\n HEXAGON_LETTER_O: IconName = IconName(\"hexagon-letter-o\")\n HEXAGON_LETTER_P: IconName = IconName(\"hexagon-letter-p\")\n HEXAGON_LETTER_Q: IconName = IconName(\"hexagon-letter-q\")\n HEXAGON_LETTER_R: IconName = IconName(\"hexagon-letter-r\")\n HEXAGON_LETTER_S: IconName = IconName(\"hexagon-letter-s\")\n HEXAGON_LETTER_T: IconName = IconName(\"hexagon-letter-t\")\n HEXAGON_LETTER_U: IconName = IconName(\"hexagon-letter-u\")\n HEXAGON_LETTER_V: IconName = IconName(\"hexagon-letter-v\")\n HEXAGON_LETTER_W: IconName = IconName(\"hexagon-letter-w\")\n HEXAGON_LETTER_X: IconName = IconName(\"hexagon-letter-x\")\n HEXAGON_LETTER_Y: IconName = IconName(\"hexagon-letter-y\")\n HEXAGON_LETTER_Z: IconName = IconName(\"hexagon-letter-z\")\n HEXAGON_NUMBER_0: IconName = IconName(\"hexagon-number-0\")\n HEXAGON_NUMBER_1: IconName = IconName(\"hexagon-number-1\")\n HEXAGON_NUMBER_2: IconName = IconName(\"hexagon-number-2\")\n HEXAGON_NUMBER_3: IconName = IconName(\"hexagon-number-3\")\n HEXAGON_NUMBER_4: IconName = IconName(\"hexagon-number-4\")\n HEXAGON_NUMBER_5: IconName = IconName(\"hexagon-number-5\")\n HEXAGON_NUMBER_6: IconName = IconName(\"hexagon-number-6\")\n HEXAGON_NUMBER_7: IconName = IconName(\"hexagon-number-7\")\n HEXAGON_NUMBER_8: IconName = IconName(\"hexagon-number-8\")\n HEXAGON_NUMBER_9: IconName = IconName(\"hexagon-number-9\")\n HEXAGON_OFF: IconName = IconName(\"hexagon-off\")\n HEXAGONAL_PRISM: IconName = IconName(\"hexagonal-prism\")\n HEXAGONAL_PRISM_OFF: IconName = IconName(\"hexagonal-prism-off\")\n HEXAGONAL_PRISM_PLUS: IconName = IconName(\"hexagonal-prism-plus\")\n HEXAGONAL_PYRAMID: IconName = IconName(\"hexagonal-pyramid\")\n HEXAGONAL_PYRAMID_OFF: IconName = IconName(\"hexagonal-pyramid-off\")\n HEXAGONAL_PYRAMID_PLUS: IconName = IconName(\"hexagonal-pyramid-plus\")\n HEXAGONS: IconName = IconName(\"hexagons\")\n HEXAGONS_OFF: IconName = IconName(\"hexagons-off\")\n HIERARCHY: IconName = IconName(\"hierarchy\")\n HIERARCHY_2: IconName = IconName(\"hierarchy-2\")\n HIERARCHY_3: IconName = IconName(\"hierarchy-3\")\n HIERARCHY_OFF: IconName = IconName(\"hierarchy-off\")\n HIGHLIGHT: IconName = IconName(\"highlight\")\n HIGHLIGHT_OFF: IconName = IconName(\"highlight-off\")\n HISTORY: IconName = IconName(\"history\")\n HISTORY_OFF: IconName = IconName(\"history-off\")\n HISTORY_TOGGLE: IconName = IconName(\"history-toggle\")\n HOME: IconName = IconName(\"home\")\n HOME_2: IconName = IconName(\"home-2\")\n HOME_BOLT: IconName = IconName(\"home-bolt\")\n HOME_CANCEL: IconName = IconName(\"home-cancel\")\n HOME_CHECK: IconName = IconName(\"home-check\")\n HOME_COG: IconName = IconName(\"home-cog\")\n HOME_DOLLAR: IconName = IconName(\"home-dollar\")\n HOME_DOT: IconName = IconName(\"home-dot\")\n HOME_DOWN: IconName = IconName(\"home-down\")\n HOME_ECO: IconName = IconName(\"home-eco\")\n HOME_EDIT: IconName = IconName(\"home-edit\")\n HOME_EXCLAMATION: IconName = IconName(\"home-exclamation\")\n HOME_HAND: IconName = IconName(\"home-hand\")\n HOME_HEART: IconName = IconName(\"home-heart\")\n HOME_INFINITY: IconName = IconName(\"home-infinity\")\n HOME_LINK: IconName = IconName(\"home-link\")\n HOME_MINUS: IconName = IconName(\"home-minus\")\n HOME_MOVE: IconName = IconName(\"home-move\")\n HOME_OFF: IconName = IconName(\"home-off\")\n HOME_PLUS: IconName = IconName(\"home-plus\")\n HOME_QUESTION: IconName = IconName(\"home-question\")\n HOME_RIBBON: IconName = IconName(\"home-ribbon\")\n HOME_SEARCH: IconName = IconName(\"home-search\")\n HOME_SHARE: IconName = IconName(\"home-share\")\n HOME_SHIELD: IconName = IconName(\"home-shield\")\n HOME_SIGNAL: IconName = IconName(\"home-signal\")\n HOME_STAR: IconName = IconName(\"home-star\")\n HOME_STATS: IconName = IconName(\"home-stats\")\n HOME_UP: IconName = IconName(\"home-up\")\n HOME_X: IconName = IconName(\"home-x\")\n HORSE_TOY: IconName = IconName(\"horse-toy\")\n HOTEL_SERVICE: IconName = IconName(\"hotel-service\")\n HOURGLASS: IconName = IconName(\"hourglass\")\n HOURGLASS_EMPTY: IconName = IconName(\"hourglass-empty\")\n HOURGLASS_FILLED: IconName = IconName(\"hourglass-filled\")\n HOURGLASS_HIGH: IconName = IconName(\"hourglass-high\")\n HOURGLASS_LOW: IconName = IconName(\"hourglass-low\")\n HOURGLASS_OFF: IconName = IconName(\"hourglass-off\")\n HTML: IconName = IconName(\"html\")\n HTTP_CONNECT: IconName = IconName(\"http-connect\")\n HTTP_DELETE: IconName = IconName(\"http-delete\")\n HTTP_GET: IconName = IconName(\"http-get\")\n HTTP_HEAD: IconName = IconName(\"http-head\")\n HTTP_OPTIONS: IconName = IconName(\"http-options\")\n HTTP_PATCH: IconName = IconName(\"http-patch\")\n HTTP_POST: IconName = IconName(\"http-post\")\n HTTP_PUT: IconName = IconName(\"http-put\")\n HTTP_QUE: IconName = IconName(\"http-que\")\n HTTP_TRACE: IconName = IconName(\"http-trace\")\n ICE_CREAM: IconName = IconName(\"ice-cream\")\n ICE_CREAM_2: IconName = IconName(\"ice-cream-2\")\n ICE_CREAM_OFF: IconName = IconName(\"ice-cream-off\")\n ICE_SKATING: IconName = IconName(\"ice-skating\")\n ICONS: IconName = IconName(\"icons\")\n ICONS_OFF: IconName = IconName(\"icons-off\")\n ID: IconName = IconName(\"id\")\n ID_BADGE: IconName = IconName(\"id-badge\")\n ID_BADGE_2: IconName = IconName(\"id-badge-2\")\n ID_BADGE_OFF: IconName = IconName(\"id-badge-off\")\n ID_OFF: IconName = IconName(\"id-off\")\n INBOX: IconName = IconName(\"inbox\")\n INBOX_OFF: IconName = IconName(\"inbox-off\")\n INDENT_DECREASE: IconName = IconName(\"indent-decrease\")\n INDENT_INCREASE: IconName = IconName(\"indent-increase\")\n INFINITY: IconName = IconName(\"infinity\")\n INFINITY_OFF: IconName = IconName(\"infinity-off\")\n INFO_CIRCLE: IconName = IconName(\"info-circle\")\n INFO_CIRCLE_FILLED: IconName = IconName(\"info-circle-filled\")\n INFO_HEXAGON: IconName = IconName(\"info-hexagon\")\n INFO_HEXAGON_FILLED: IconName = IconName(\"info-hexagon-filled\")\n INFO_OCTAGON: IconName = IconName(\"info-octagon\")\n INFO_OCTAGON_FILLED: IconName = IconName(\"info-octagon-filled\")\n INFO_SMALL: IconName = IconName(\"info-small\")\n INFO_SQUARE: IconName = IconName(\"info-square\")\n INFO_SQUARE_FILLED: IconName = IconName(\"info-square-filled\")\n INFO_SQUARE_ROUNDED: IconName = IconName(\"info-square-rounded\")\n INFO_SQUARE_ROUNDED_FILLED: IconName = IconName(\"info-square-rounded-filled\")\n INFO_TRIANGLE: IconName = IconName(\"info-triangle\")\n INFO_TRIANGLE_FILLED: IconName = IconName(\"info-triangle-filled\")\n INNER_SHADOW_BOTTOM: IconName = IconName(\"inner-shadow-bottom\")\n INNER_SHADOW_BOTTOM_FILLED: IconName = IconName(\"inner-shadow-bottom-filled\")\n INNER_SHADOW_BOTTOM_LEFT: IconName = IconName(\"inner-shadow-bottom-left\")\n INNER_SHADOW_BOTTOM_LEFT_FILLED: IconName = IconName(\n \"inner-shadow-bottom-left-filled\"\n )\n INNER_SHADOW_BOTTOM_RIGHT: IconName = IconName(\"inner-shadow-bottom-right\")\n INNER_SHADOW_BOTTOM_RIGHT_FILLED: IconName = IconName(\n \"inner-shadow-bottom-right-filled\"\n )\n INNER_SHADOW_LEFT: IconName = IconName(\"inner-shadow-left\")\n INNER_SHADOW_LEFT_FILLED: IconName = IconName(\"inner-shadow-left-filled\")\n INNER_SHADOW_RIGHT: IconName = IconName(\"inner-shadow-right\")\n INNER_SHADOW_RIGHT_FILLED: IconName = IconName(\"inner-shadow-right-filled\")\n INNER_SHADOW_TOP: IconName = IconName(\"inner-shadow-top\")\n INNER_SHADOW_TOP_FILLED: IconName = IconName(\"inner-shadow-top-filled\")\n INNER_SHADOW_TOP_LEFT: IconName = IconName(\"inner-shadow-top-left\")\n INNER_SHADOW_TOP_LEFT_FILLED: IconName = IconName(\"inner-shadow-top-left-filled\")\n INNER_SHADOW_TOP_RIGHT: IconName = IconName(\"inner-shadow-top-right\")\n INNER_SHADOW_TOP_RIGHT_FILLED: IconName = IconName(\"inner-shadow-top-right-filled\")\n INPUT_SEARCH: IconName = IconName(\"input-search\")\n IRONING: IconName = IconName(\"ironing\")\n IRONING_1: IconName = IconName(\"ironing-1\")\n IRONING_2: IconName = IconName(\"ironing-2\")\n IRONING_3: IconName = IconName(\"ironing-3\")\n IRONING_OFF: IconName = IconName(\"ironing-off\")\n IRONING_STEAM: IconName = IconName(\"ironing-steam\")\n IRONING_STEAM_OFF: IconName = IconName(\"ironing-steam-off\")\n IRREGULAR_POLYHEDRON: IconName = IconName(\"irregular-polyhedron\")\n IRREGULAR_POLYHEDRON_OFF: IconName = IconName(\"irregular-polyhedron-off\")\n IRREGULAR_POLYHEDRON_PLUS: IconName = IconName(\"irregular-polyhedron-plus\")\n ITALIC: IconName = IconName(\"italic\")\n JACKET: IconName = IconName(\"jacket\")\n JETPACK: IconName = IconName(\"jetpack\")\n JEWISH_STAR: IconName = IconName(\"jewish-star\")\n JEWISH_STAR_FILLED: IconName = IconName(\"jewish-star-filled\")\n JPG: IconName = IconName(\"jpg\")\n JSON: IconName = IconName(\"json\")\n JUMP_ROPE: IconName = IconName(\"jump-rope\")\n KARATE: IconName = IconName(\"karate\")\n KAYAK: IconName = IconName(\"kayak\")\n KERING: IconName = IconName(\"kering\")\n KEY: IconName = IconName(\"key\")\n KEY_OFF: IconName = IconName(\"key-off\")\n KEYBOARD: IconName = IconName(\"keyboard\")\n KEYBOARD_HIDE: IconName = IconName(\"keyboard-hide\")\n KEYBOARD_OFF: IconName = IconName(\"keyboard-off\")\n KEYBOARD_SHOW: IconName = IconName(\"keyboard-show\")\n KEYFRAME: IconName = IconName(\"keyframe\")\n KEYFRAME_ALIGN_CENTER: IconName = IconName(\"keyframe-align-center\")\n KEYFRAME_ALIGN_HORIZONTAL: IconName = IconName(\"keyframe-align-horizontal\")\n KEYFRAME_ALIGN_VERTICAL: IconName = IconName(\"keyframe-align-vertical\")\n KEYFRAMES: IconName = IconName(\"keyframes\")\n LADDER: IconName = IconName(\"ladder\")\n LADDER_OFF: IconName = IconName(\"ladder-off\")\n LAMBDA: IconName = IconName(\"lambda\")\n LAMP: IconName = IconName(\"lamp\")\n LAMP_2: IconName = IconName(\"lamp-2\")\n LAMP_OFF: IconName = IconName(\"lamp-off\")\n LANE: IconName = IconName(\"lane\")\n LANGUAGE: IconName = IconName(\"language\")\n LANGUAGE_HIRAGANA: IconName = IconName(\"language-hiragana\")\n LANGUAGE_KATAKANA: IconName = IconName(\"language-katakana\")\n LANGUAGE_OFF: IconName = IconName(\"language-off\")\n LASSO: IconName = IconName(\"lasso\")\n LASSO_OFF: IconName = IconName(\"lasso-off\")\n LASSO_POLYGON: IconName = IconName(\"lasso-polygon\")\n LAYERS_DIFFERENCE: IconName = IconName(\"layers-difference\")\n LAYERS_INTERSECT: IconName = IconName(\"layers-intersect\")\n LAYERS_INTERSECT_2: IconName = IconName(\"layers-intersect-2\")\n LAYERS_LINKED: IconName = IconName(\"layers-linked\")\n LAYERS_OFF: IconName = IconName(\"layers-off\")\n LAYERS_SUBTRACT: IconName = IconName(\"layers-subtract\")\n LAYERS_UNION: IconName = IconName(\"layers-union\")\n LAYOUT: IconName = IconName(\"layout\")\n LAYOUT_2: IconName = IconName(\"layout-2\")\n LAYOUT_ALIGN_BOTTOM: IconName = IconName(\"layout-align-bottom\")\n LAYOUT_ALIGN_CENTER: IconName = IconName(\"layout-align-center\")\n LAYOUT_ALIGN_LEFT: IconName = IconName(\"layout-align-left\")\n LAYOUT_ALIGN_MIDDLE: IconName = IconName(\"layout-align-middle\")\n LAYOUT_ALIGN_RIGHT: IconName = IconName(\"layout-align-right\")\n LAYOUT_ALIGN_TOP: IconName = IconName(\"layout-align-top\")\n LAYOUT_BOARD: IconName = IconName(\"layout-board\")\n LAYOUT_BOARD_SPLIT: IconName = IconName(\"layout-board-split\")\n LAYOUT_BOTTOMBAR: IconName = IconName(\"layout-bottombar\")\n LAYOUT_BOTTOMBAR_COLLAPSE: IconName = IconName(\"layout-bottombar-collapse\")\n LAYOUT_BOTTOMBAR_EXPAND: IconName = IconName(\"layout-bottombar-expand\")\n LAYOUT_CARDS: IconName = IconName(\"layout-cards\")\n LAYOUT_COLLAGE: IconName = IconName(\"layout-collage\")\n LAYOUT_COLUMNS: IconName = IconName(\"layout-columns\")\n LAYOUT_DASHBOARD: IconName = IconName(\"layout-dashboard\")\n LAYOUT_DISTRIBUTE_HORIZONTAL: IconName = IconName(\"layout-distribute-horizontal\")\n LAYOUT_DISTRIBUTE_VERTICAL: IconName = IconName(\"layout-distribute-vertical\")\n LAYOUT_GRID: IconName = IconName(\"layout-grid\")\n LAYOUT_GRID_ADD: IconName = IconName(\"layout-grid-add\")\n LAYOUT_GRID_REMOVE: IconName = IconName(\"layout-grid-remove\")\n LAYOUT_KANBAN: IconName = IconName(\"layout-kanban\")\n LAYOUT_LIST: IconName = IconName(\"layout-list\")\n LAYOUT_NAVBAR: IconName = IconName(\"layout-navbar\")\n LAYOUT_NAVBAR_COLLAPSE: IconName = IconName(\"layout-navbar-collapse\")\n LAYOUT_NAVBAR_EXPAND: IconName = IconName(\"layout-navbar-expand\")\n LAYOUT_OFF: IconName = IconName(\"layout-off\")\n LAYOUT_ROWS: IconName = IconName(\"layout-rows\")\n LAYOUT_SIDEBAR: IconName = IconName(\"layout-sidebar\")\n LAYOUT_SIDEBAR_LEFT_COLLAPSE: IconName = IconName(\"layout-sidebar-left-collapse\")\n LAYOUT_SIDEBAR_LEFT_EXPAND: IconName = IconName(\"layout-sidebar-left-expand\")\n LAYOUT_SIDEBAR_RIGHT: IconName = IconName(\"layout-sidebar-right\")\n LAYOUT_SIDEBAR_RIGHT_COLLAPSE: IconName = IconName(\"layout-sidebar-right-collapse\")\n LAYOUT_SIDEBAR_RIGHT_EXPAND: IconName = IconName(\"layout-sidebar-right-expand\")\n LEAF: IconName = IconName(\"leaf\")\n LEAF_OFF: IconName = IconName(\"leaf-off\")\n LEGO: IconName = IconName(\"lego\")\n LEGO_OFF: IconName = IconName(\"lego-off\")\n LEMON: IconName = IconName(\"lemon\")\n LEMON_2: IconName = IconName(\"lemon-2\")\n LETTER_A: IconName = IconName(\"letter-a\")\n LETTER_B: IconName = IconName(\"letter-b\")\n LETTER_C: IconName = IconName(\"letter-c\")\n LETTER_CASE: IconName = IconName(\"letter-case\")\n LETTER_CASE_LOWER: IconName = IconName(\"letter-case-lower\")\n LETTER_CASE_TOGGLE: IconName = IconName(\"letter-case-toggle\")\n LETTER_CASE_UPPER: IconName = IconName(\"letter-case-upper\")\n LETTER_D: IconName = IconName(\"letter-d\")\n LETTER_E: IconName = IconName(\"letter-e\")\n LETTER_F: IconName = IconName(\"letter-f\")\n LETTER_G: IconName = IconName(\"letter-g\")\n LETTER_H: IconName = IconName(\"letter-h\")\n LETTER_I: IconName = IconName(\"letter-i\")\n LETTER_J: IconName = IconName(\"letter-j\")\n LETTER_K: IconName = IconName(\"letter-k\")\n LETTER_L: IconName = IconName(\"letter-l\")\n LETTER_M: IconName = IconName(\"letter-m\")\n LETTER_N: IconName = IconName(\"letter-n\")\n LETTER_O: IconName = IconName(\"letter-o\")\n LETTER_P: IconName = IconName(\"letter-p\")\n LETTER_Q: IconName = IconName(\"letter-q\")\n LETTER_R: IconName = IconName(\"letter-r\")\n LETTER_S: IconName = IconName(\"letter-s\")\n LETTER_SPACING: IconName = IconName(\"letter-spacing\")\n LETTER_T: IconName = IconName(\"letter-t\")\n LETTER_U: IconName = IconName(\"letter-u\")\n LETTER_V: IconName = IconName(\"letter-v\")\n LETTER_W: IconName = IconName(\"letter-w\")\n LETTER_X: IconName = IconName(\"letter-x\")\n LETTER_Y: IconName = IconName(\"letter-y\")\n LETTER_Z: IconName = IconName(\"letter-z\")\n LICENSE: IconName = IconName(\"license\")\n LICENSE_OFF: IconName = IconName(\"license-off\")\n LIFEBUOY: IconName = IconName(\"lifebuoy\")\n LIFEBUOY_OFF: IconName = IconName(\"lifebuoy-off\")\n LIGHTER: IconName = IconName(\"lighter\")\n LINE: IconName = IconName(\"line\")\n LINE_DASHED: IconName = IconName(\"line-dashed\")\n LINE_DOTTED: IconName = IconName(\"line-dotted\")\n LINE_HEIGHT: IconName = IconName(\"line-height\")\n LINK: IconName = IconName(\"link\")\n LINK_OFF: IconName = IconName(\"link-off\")\n LIST: IconName = IconName(\"list\")\n LIST_CHECK: IconName = IconName(\"list-check\")\n LIST_DETAILS: IconName = IconName(\"list-details\")\n LIST_NUMBERS: IconName = IconName(\"list-numbers\")\n LIST_SEARCH: IconName = IconName(\"list-search\")\n LIST_TREE: IconName = IconName(\"list-tree\")\n LIVE_PHOTO: IconName = IconName(\"live-photo\")\n LIVE_PHOTO_OFF: IconName = IconName(\"live-photo-off\")\n LIVE_VIEW: IconName = IconName(\"live-view\")\n LOAD_BALANCER: IconName = IconName(\"load-balancer\")\n LOADER: IconName = IconName(\"loader\")\n LOADER_2: IconName = IconName(\"loader-2\")\n LOADER_3: IconName = IconName(\"loader-3\")\n LOADER_QUARTER: IconName = IconName(\"loader-quarter\")\n LOCATION: IconName = IconName(\"location\")\n LOCATION_BROKEN: IconName = IconName(\"location-broken\")\n LOCATION_FILLED: IconName = IconName(\"location-filled\")\n LOCATION_OFF: IconName = IconName(\"location-off\")\n LOCK: IconName = IconName(\"lock\")\n LOCK_ACCESS: IconName = IconName(\"lock-access\")\n LOCK_ACCESS_OFF: IconName = IconName(\"lock-access-off\")\n LOCK_BOLT: IconName = IconName(\"lock-bolt\")\n LOCK_CANCEL: IconName = IconName(\"lock-cancel\")\n LOCK_CHECK: IconName = IconName(\"lock-check\")\n LOCK_CODE: IconName = IconName(\"lock-code\")\n LOCK_COG: IconName = IconName(\"lock-cog\")\n LOCK_DOLLAR: IconName = IconName(\"lock-dollar\")\n LOCK_DOWN: IconName = IconName(\"lock-down\")\n LOCK_EXCLAMATION: IconName = IconName(\"lock-exclamation\")\n LOCK_HEART: IconName = IconName(\"lock-heart\")\n LOCK_MINUS: IconName = IconName(\"lock-minus\")\n LOCK_OFF: IconName = IconName(\"lock-off\")\n LOCK_OPEN: IconName = IconName(\"lock-open\")\n LOCK_OPEN_OFF: IconName = IconName(\"lock-open-off\")\n LOCK_PAUSE: IconName = IconName(\"lock-pause\")\n LOCK_PIN: IconName = IconName(\"lock-pin\")\n LOCK_PLUS: IconName = IconName(\"lock-plus\")\n LOCK_QUESTION: IconName = IconName(\"lock-question\")\n LOCK_SEARCH: IconName = IconName(\"lock-search\")\n LOCK_SHARE: IconName = IconName(\"lock-share\")\n LOCK_SQUARE: IconName = IconName(\"lock-square\")\n LOCK_SQUARE_ROUNDED: IconName = IconName(\"lock-square-rounded\")\n LOCK_SQUARE_ROUNDED_FILLED: IconName = IconName(\"lock-square-rounded-filled\")\n LOCK_STAR: IconName = IconName(\"lock-star\")\n LOCK_UP: IconName = IconName(\"lock-up\")\n LOCK_X: IconName = IconName(\"lock-x\")\n LOGIC_AND: IconName = IconName(\"logic-and\")\n LOGIC_BUFFER: IconName = IconName(\"logic-buffer\")\n LOGIC_NAND: IconName = IconName(\"logic-nand\")\n LOGIC_NOR: IconName = IconName(\"logic-nor\")\n LOGIC_NOT: IconName = IconName(\"logic-not\")\n LOGIC_OR: IconName = IconName(\"logic-or\")\n LOGIC_XNOR: IconName = IconName(\"logic-xnor\")\n LOGIC_XOR: IconName = IconName(\"logic-xor\")\n LOGIN: IconName = IconName(\"login\")\n LOGOUT: IconName = IconName(\"logout\")\n LOGOUT_2: IconName = IconName(\"logout-2\")\n LOLLIPOP: IconName = IconName(\"lollipop\")\n LOLLIPOP_OFF: IconName = IconName(\"lollipop-off\")\n LUGGAGE: IconName = IconName(\"luggage\")\n LUGGAGE_OFF: IconName = IconName(\"luggage-off\")\n LUNGS: IconName = IconName(\"lungs\")\n LUNGS_OFF: IconName = IconName(\"lungs-off\")\n MACRO: IconName = IconName(\"macro\")\n MACRO_OFF: IconName = IconName(\"macro-off\")\n MAGNET: IconName = IconName(\"magnet\")\n MAGNET_OFF: IconName = IconName(\"magnet-off\")\n MAIL: IconName = IconName(\"mail\")\n MAIL_AI: IconName = IconName(\"mail-ai\")\n MAIL_BOLT: IconName = IconName(\"mail-bolt\")\n MAIL_CANCEL: IconName = IconName(\"mail-cancel\")\n MAIL_CHECK: IconName = IconName(\"mail-check\")\n MAIL_CODE: IconName = IconName(\"mail-code\")\n MAIL_COG: IconName = IconName(\"mail-cog\")\n MAIL_DOLLAR: IconName = IconName(\"mail-dollar\")\n MAIL_DOWN: IconName = IconName(\"mail-down\")\n MAIL_EXCLAMATION: IconName = IconName(\"mail-exclamation\")\n MAIL_FAST: IconName = IconName(\"mail-fast\")\n MAIL_FILLED: IconName = IconName(\"mail-filled\")\n MAIL_FORWARD: IconName = IconName(\"mail-forward\")\n MAIL_HEART: IconName = IconName(\"mail-heart\")\n MAIL_MINUS: IconName = IconName(\"mail-minus\")\n MAIL_OFF: IconName = IconName(\"mail-off\")\n MAIL_OPENED: IconName = IconName(\"mail-opened\")\n MAIL_OPENED_FILLED: IconName = IconName(\"mail-opened-filled\")\n MAIL_PAUSE: IconName = IconName(\"mail-pause\")\n MAIL_PIN: IconName = IconName(\"mail-pin\")\n MAIL_PLUS: IconName = IconName(\"mail-plus\")\n MAIL_QUESTION: IconName = IconName(\"mail-question\")\n MAIL_SEARCH: IconName = IconName(\"mail-search\")\n MAIL_SHARE: IconName = IconName(\"mail-share\")\n MAIL_STAR: IconName = IconName(\"mail-star\")\n MAIL_UP: IconName = IconName(\"mail-up\")\n MAIL_X: IconName = IconName(\"mail-x\")\n MAILBOX: IconName = IconName(\"mailbox\")\n MAILBOX_OFF: IconName = IconName(\"mailbox-off\")\n MAN: IconName = IconName(\"man\")\n MANUAL_GEARBOX: IconName = IconName(\"manual-gearbox\")\n MAP: IconName = IconName(\"map\")\n MAP_2: IconName = IconName(\"map-2\")\n MAP_OFF: IconName = IconName(\"map-off\")\n MAP_PIN: IconName = IconName(\"map-pin\")\n MAP_PIN_BOLT: IconName = IconName(\"map-pin-bolt\")\n MAP_PIN_CANCEL: IconName = IconName(\"map-pin-cancel\")\n MAP_PIN_CHECK: IconName = IconName(\"map-pin-check\")\n MAP_PIN_CODE: IconName = IconName(\"map-pin-code\")\n MAP_PIN_COG: IconName = IconName(\"map-pin-cog\")\n MAP_PIN_DOLLAR: IconName = IconName(\"map-pin-dollar\")\n MAP_PIN_DOWN: IconName = IconName(\"map-pin-down\")\n MAP_PIN_EXCLAMATION: IconName = IconName(\"map-pin-exclamation\")\n MAP_PIN_FILLED: IconName = IconName(\"map-pin-filled\")\n MAP_PIN_HEART: IconName = IconName(\"map-pin-heart\")\n MAP_PIN_MINUS: IconName = IconName(\"map-pin-minus\")\n MAP_PIN_OFF: IconName = IconName(\"map-pin-off\")\n MAP_PIN_PAUSE: IconName = IconName(\"map-pin-pause\")\n MAP_PIN_PIN: IconName = IconName(\"map-pin-pin\")\n MAP_PIN_PLUS: IconName = IconName(\"map-pin-plus\")\n MAP_PIN_QUESTION: IconName = IconName(\"map-pin-question\")\n MAP_PIN_SEARCH: IconName = IconName(\"map-pin-search\")\n MAP_PIN_SHARE: IconName = IconName(\"map-pin-share\")\n MAP_PIN_STAR: IconName = IconName(\"map-pin-star\")\n MAP_PIN_UP: IconName = IconName(\"map-pin-up\")\n MAP_PIN_X: IconName = IconName(\"map-pin-x\")\n MAP_PINS: IconName = IconName(\"map-pins\")\n MAP_SEARCH: IconName = IconName(\"map-search\")\n MARKDOWN: IconName = IconName(\"markdown\")\n MARKDOWN_OFF: IconName = IconName(\"markdown-off\")\n MARQUEE: IconName = IconName(\"marquee\")\n MARQUEE_2: IconName = IconName(\"marquee-2\")\n MARQUEE_OFF: IconName = IconName(\"marquee-off\")\n MARS: IconName = IconName(\"mars\")\n MASK: IconName = IconName(\"mask\")\n MASK_OFF: IconName = IconName(\"mask-off\")\n MASKS_THEATER: IconName = IconName(\"masks-theater\")\n MASKS_THEATER_OFF: IconName = IconName(\"masks-theater-off\")\n MASSAGE: IconName = IconName(\"massage\")\n MATCHSTICK: IconName = IconName(\"matchstick\")\n MATH: IconName = IconName(\"math\")\n MATH_1_DIVIDE_2: IconName = IconName(\"math-1-divide-2\")\n MATH_1_DIVIDE_3: IconName = IconName(\"math-1-divide-3\")\n MATH_AVG: IconName = IconName(\"math-avg\")\n MATH_EQUAL_GREATER: IconName = IconName(\"math-equal-greater\")\n MATH_EQUAL_LOWER: IconName = IconName(\"math-equal-lower\")\n MATH_FUNCTION: IconName = IconName(\"math-function\")\n MATH_FUNCTION_OFF: IconName = IconName(\"math-function-off\")\n MATH_FUNCTION_Y: IconName = IconName(\"math-function-y\")\n MATH_GREATER: IconName = IconName(\"math-greater\")\n MATH_INTEGRAL: IconName = IconName(\"math-integral\")\n MATH_INTEGRAL_X: IconName = IconName(\"math-integral-x\")\n MATH_INTEGRALS: IconName = IconName(\"math-integrals\")\n MATH_LOWER: IconName = IconName(\"math-lower\")\n MATH_MAX: IconName = IconName(\"math-max\")\n MATH_MIN: IconName = IconName(\"math-min\")\n MATH_NOT: IconName = IconName(\"math-not\")\n MATH_OFF: IconName = IconName(\"math-off\")\n MATH_PI: IconName = IconName(\"math-pi\")\n MATH_PI_DIVIDE_2: IconName = IconName(\"math-pi-divide-2\")\n MATH_SYMBOLS: IconName = IconName(\"math-symbols\")\n MATH_X_DIVIDE_2: IconName = IconName(\"math-x-divide-2\")\n MATH_X_DIVIDE_Y: IconName = IconName(\"math-x-divide-y\")\n MATH_X_DIVIDE_Y_2: IconName = IconName(\"math-x-divide-y-2\")\n MATH_X_MINUS_X: IconName = IconName(\"math-x-minus-x\")\n MATH_X_MINUS_Y: IconName = IconName(\"math-x-minus-y\")\n MATH_X_PLUS_X: IconName = IconName(\"math-x-plus-x\")\n MATH_X_PLUS_Y: IconName = IconName(\"math-x-plus-y\")\n MATH_XY: IconName = IconName(\"math-xy\")\n MATH_Y_MINUS_Y: IconName = IconName(\"math-y-minus-y\")\n MATH_Y_PLUS_Y: IconName = IconName(\"math-y-plus-y\")\n MAXIMIZE: IconName = IconName(\"maximize\")\n MAXIMIZE_OFF: IconName = IconName(\"maximize-off\")\n MEAT: IconName = IconName(\"meat\")\n MEAT_OFF: IconName = IconName(\"meat-off\")\n MEDAL: IconName = IconName(\"medal\")\n MEDAL_2: IconName = IconName(\"medal-2\")\n MEDICAL_CROSS: IconName = IconName(\"medical-cross\")\n MEDICAL_CROSS_CIRCLE: IconName = IconName(\"medical-cross-circle\")\n MEDICAL_CROSS_FILLED: IconName = IconName(\"medical-cross-filled\")\n MEDICAL_CROSS_OFF: IconName = IconName(\"medical-cross-off\")\n MEDICINE_SYRUP: IconName = IconName(\"medicine-syrup\")\n MEEPLE: IconName = IconName(\"meeple\")\n MENORAH: IconName = IconName(\"menorah\")\n MENU: IconName = IconName(\"menu\")\n MENU_2: IconName = IconName(\"menu-2\")\n MENU_DEEP: IconName = IconName(\"menu-deep\")\n MENU_ORDER: IconName = IconName(\"menu-order\")\n MESSAGE: IconName = IconName(\"message\")\n MESSAGE_2: IconName = IconName(\"message-2\")\n MESSAGE_2_BOLT: IconName = IconName(\"message-2-bolt\")\n MESSAGE_2_CANCEL: IconName = IconName(\"message-2-cancel\")\n MESSAGE_2_CHECK: IconName = IconName(\"message-2-check\")\n MESSAGE_2_CODE: IconName = IconName(\"message-2-code\")\n MESSAGE_2_COG: IconName = IconName(\"message-2-cog\")\n MESSAGE_2_DOLLAR: IconName = IconName(\"message-2-dollar\")\n MESSAGE_2_DOWN: IconName = IconName(\"message-2-down\")\n MESSAGE_2_EXCLAMATION: IconName = IconName(\"message-2-exclamation\")\n MESSAGE_2_HEART: IconName = IconName(\"message-2-heart\")\n MESSAGE_2_MINUS: IconName = IconName(\"message-2-minus\")\n MESSAGE_2_OFF: IconName = IconName(\"message-2-off\")\n MESSAGE_2_PAUSE: IconName = IconName(\"message-2-pause\")\n MESSAGE_2_PIN: IconName = IconName(\"message-2-pin\")\n MESSAGE_2_PLUS: IconName = IconName(\"message-2-plus\")\n MESSAGE_2_QUESTION: IconName = IconName(\"message-2-question\")\n MESSAGE_2_SEARCH: IconName = IconName(\"message-2-search\")\n MESSAGE_2_SHARE: IconName = IconName(\"message-2-share\")\n MESSAGE_2_STAR: IconName = IconName(\"message-2-star\")\n MESSAGE_2_UP: IconName = IconName(\"message-2-up\")\n MESSAGE_2_X: IconName = IconName(\"message-2-x\")\n MESSAGE_BOLT: IconName = IconName(\"message-bolt\")\n MESSAGE_CANCEL: IconName = IconName(\"message-cancel\")\n MESSAGE_CHATBOT: IconName = IconName(\"message-chatbot\")\n MESSAGE_CHECK: IconName = IconName(\"message-check\")\n MESSAGE_CIRCLE: IconName = IconName(\"message-circle\")\n MESSAGE_CIRCLE_2: IconName = IconName(\"message-circle-2\")\n MESSAGE_CIRCLE_2_FILLED: IconName = IconName(\"message-circle-2-filled\")\n MESSAGE_CIRCLE_BOLT: IconName = IconName(\"message-circle-bolt\")\n MESSAGE_CIRCLE_CANCEL: IconName = IconName(\"message-circle-cancel\")\n MESSAGE_CIRCLE_CHECK: IconName = IconName(\"message-circle-check\")\n MESSAGE_CIRCLE_CODE: IconName = IconName(\"message-circle-code\")\n MESSAGE_CIRCLE_COG: IconName = IconName(\"message-circle-cog\")\n MESSAGE_CIRCLE_DOLLAR: IconName = IconName(\"message-circle-dollar\")\n MESSAGE_CIRCLE_DOWN: IconName = IconName(\"message-circle-down\")\n MESSAGE_CIRCLE_EXCLAMATION: IconName = IconName(\"message-circle-exclamation\")\n MESSAGE_CIRCLE_HEART: IconName = IconName(\"message-circle-heart\")\n MESSAGE_CIRCLE_MINUS: IconName = IconName(\"message-circle-minus\")\n MESSAGE_CIRCLE_OFF: IconName = IconName(\"message-circle-off\")\n MESSAGE_CIRCLE_PAUSE: IconName = IconName(\"message-circle-pause\")\n MESSAGE_CIRCLE_PIN: IconName = IconName(\"message-circle-pin\")\n MESSAGE_CIRCLE_PLUS: IconName = IconName(\"message-circle-plus\")\n MESSAGE_CIRCLE_QUESTION: IconName = IconName(\"message-circle-question\")\n MESSAGE_CIRCLE_SEARCH: IconName = IconName(\"message-circle-search\")\n MESSAGE_CIRCLE_SHARE: IconName = IconName(\"message-circle-share\")\n MESSAGE_CIRCLE_STAR: IconName = IconName(\"message-circle-star\")\n MESSAGE_CIRCLE_UP: IconName = IconName(\"message-circle-up\")\n MESSAGE_CIRCLE_X: IconName = IconName(\"message-circle-x\")\n MESSAGE_CODE: IconName = IconName(\"message-code\")\n MESSAGE_COG: IconName = IconName(\"message-cog\")\n MESSAGE_DOLLAR: IconName = IconName(\"message-dollar\")\n MESSAGE_DOTS: IconName = IconName(\"message-dots\")\n MESSAGE_DOWN: IconName = IconName(\"message-down\")\n MESSAGE_EXCLAMATION: IconName = IconName(\"message-exclamation\")\n MESSAGE_FORWARD: IconName = IconName(\"message-forward\")\n MESSAGE_HEART: IconName = IconName(\"message-heart\")\n MESSAGE_LANGUAGE: IconName = IconName(\"message-language\")\n MESSAGE_MINUS: IconName = IconName(\"message-minus\")\n MESSAGE_OFF: IconName = IconName(\"message-off\")\n MESSAGE_PAUSE: IconName = IconName(\"message-pause\")\n MESSAGE_PIN: IconName = IconName(\"message-pin\")\n MESSAGE_PLUS: IconName = IconName(\"message-plus\")\n MESSAGE_QUESTION: IconName = IconName(\"message-question\")\n MESSAGE_REPORT: IconName = IconName(\"message-report\")\n MESSAGE_SEARCH: IconName = IconName(\"message-search\")\n MESSAGE_SHARE: IconName = IconName(\"message-share\")\n MESSAGE_STAR: IconName = IconName(\"message-star\")\n MESSAGE_UP: IconName = IconName(\"message-up\")\n MESSAGE_X: IconName = IconName(\"message-x\")\n MESSAGES: IconName = IconName(\"messages\")\n MESSAGES_OFF: IconName = IconName(\"messages-off\")\n METEOR: IconName = IconName(\"meteor\")\n METEOR_OFF: IconName = IconName(\"meteor-off\")\n MICHELIN_BIB_GOURMAND: IconName = IconName(\"michelin-bib-gourmand\")\n MICHELIN_STAR: IconName = IconName(\"michelin-star\")\n MICHELIN_STAR_GREEN: IconName = IconName(\"michelin-star-green\")\n MICKEY: IconName = IconName(\"mickey\")\n MICKEY_FILLED: IconName = IconName(\"mickey-filled\")\n MICROPHONE: IconName = IconName(\"microphone\")\n MICROPHONE_2: IconName = IconName(\"microphone-2\")\n MICROPHONE_2_OFF: IconName = IconName(\"microphone-2-off\")\n MICROPHONE_OFF: IconName = IconName(\"microphone-off\")\n MICROSCOPE: IconName = IconName(\"microscope\")\n MICROSCOPE_OFF: IconName = IconName(\"microscope-off\")\n MICROWAVE: IconName = IconName(\"microwave\")\n MICROWAVE_OFF: IconName = IconName(\"microwave-off\")\n MILITARY_AWARD: IconName = IconName(\"military-award\")\n MILITARY_RANK: IconName = IconName(\"military-rank\")\n MILK: IconName = IconName(\"milk\")\n MILK_OFF: IconName = IconName(\"milk-off\")\n MILKSHAKE: IconName = IconName(\"milkshake\")\n MINIMIZE: IconName = IconName(\"minimize\")\n MINUS: IconName = IconName(\"minus\")\n MINUS_VERTICAL: IconName = IconName(\"minus-vertical\")\n MIST: IconName = IconName(\"mist\")\n MIST_OFF: IconName = IconName(\"mist-off\")\n MOBILEDATA: IconName = IconName(\"mobiledata\")\n MOBILEDATA_OFF: IconName = IconName(\"mobiledata-off\")\n MONEYBAG: IconName = IconName(\"moneybag\")\n MOOD_ANGRY: IconName = IconName(\"mood-angry\")\n MOOD_ANNOYED: IconName = IconName(\"mood-annoyed\")\n MOOD_ANNOYED_2: IconName = IconName(\"mood-annoyed-2\")\n MOOD_BOY: IconName = IconName(\"mood-boy\")\n MOOD_CHECK: IconName = IconName(\"mood-check\")\n MOOD_COG: IconName = IconName(\"mood-cog\")\n MOOD_CONFUZED: IconName = IconName(\"mood-confuzed\")\n MOOD_CONFUZED_FILLED: IconName = IconName(\"mood-confuzed-filled\")\n MOOD_CRAZY_HAPPY: IconName = IconName(\"mood-crazy-happy\")\n MOOD_CRY: IconName = IconName(\"mood-cry\")\n MOOD_DOLLAR: IconName = IconName(\"mood-dollar\")\n MOOD_EDIT: IconName = IconName(\"mood-edit\")\n MOOD_EMPTY: IconName = IconName(\"mood-empty\")\n MOOD_EMPTY_FILLED: IconName = IconName(\"mood-empty-filled\")\n MOOD_HAPPY: IconName = IconName(\"mood-happy\")\n MOOD_HAPPY_FILLED: IconName = IconName(\"mood-happy-filled\")\n MOOD_HEART: IconName = IconName(\"mood-heart\")\n MOOD_KID: IconName = IconName(\"mood-kid\")\n MOOD_KID_FILLED: IconName = IconName(\"mood-kid-filled\")\n MOOD_LOOK_LEFT: IconName = IconName(\"mood-look-left\")\n MOOD_LOOK_RIGHT: IconName = IconName(\"mood-look-right\")\n MOOD_MINUS: IconName = IconName(\"mood-minus\")\n MOOD_NERD: IconName = IconName(\"mood-nerd\")\n MOOD_NERVOUS: IconName = IconName(\"mood-nervous\")\n MOOD_NEUTRAL: IconName = IconName(\"mood-neutral\")\n MOOD_NEUTRAL_FILLED: IconName = IconName(\"mood-neutral-filled\")\n MOOD_OFF: IconName = IconName(\"mood-off\")\n MOOD_PIN: IconName = IconName(\"mood-pin\")\n MOOD_PLUS: IconName = IconName(\"mood-plus\")\n MOOD_SAD: IconName = IconName(\"mood-sad\")\n MOOD_SAD_2: IconName = IconName(\"mood-sad-2\")\n MOOD_SAD_DIZZY: IconName = IconName(\"mood-sad-dizzy\")\n MOOD_SAD_FILLED: IconName = IconName(\"mood-sad-filled\")\n MOOD_SAD_SQUINT: IconName = IconName(\"mood-sad-squint\")\n MOOD_SEARCH: IconName = IconName(\"mood-search\")\n MOOD_SHARE: IconName = IconName(\"mood-share\")\n MOOD_SICK: IconName = IconName(\"mood-sick\")\n MOOD_SILENCE: IconName = IconName(\"mood-silence\")\n MOOD_SING: IconName = IconName(\"mood-sing\")\n MOOD_SMILE: IconName = IconName(\"mood-smile\")\n MOOD_SMILE_BEAM: IconName = IconName(\"mood-smile-beam\")\n MOOD_SMILE_DIZZY: IconName = IconName(\"mood-smile-dizzy\")\n MOOD_SMILE_FILLED: IconName = IconName(\"mood-smile-filled\")\n MOOD_SUPRISED: IconName = IconName(\"mood-suprised\")\n MOOD_TONGUE: IconName = IconName(\"mood-tongue\")\n MOOD_TONGUE_WINK: IconName = IconName(\"mood-tongue-wink\")\n MOOD_TONGUE_WINK_2: IconName = IconName(\"mood-tongue-wink-2\")\n MOOD_UNAMUSED: IconName = IconName(\"mood-unamused\")\n MOOD_UP: IconName = IconName(\"mood-up\")\n MOOD_WINK: IconName = IconName(\"mood-wink\")\n MOOD_WINK_2: IconName = IconName(\"mood-wink-2\")\n MOOD_WRRR: IconName = IconName(\"mood-wrrr\")\n MOOD_X: IconName = IconName(\"mood-x\")\n MOOD_XD: IconName = IconName(\"mood-xd\")\n MOON: IconName = IconName(\"moon\")\n MOON_2: IconName = IconName(\"moon-2\")\n MOON_FILLED: IconName = IconName(\"moon-filled\")\n MOON_OFF: IconName = IconName(\"moon-off\")\n MOON_STARS: IconName = IconName(\"moon-stars\")\n MOPED: IconName = IconName(\"moped\")\n MOTORBIKE: IconName = IconName(\"motorbike\")\n MOUNTAIN: IconName = IconName(\"mountain\")\n MOUNTAIN_OFF: IconName = IconName(\"mountain-off\")\n MOUSE: IconName = IconName(\"mouse\")\n MOUSE_2: IconName = IconName(\"mouse-2\")\n MOUSE_OFF: IconName = IconName(\"mouse-off\")\n MOUSTACHE: IconName = IconName(\"moustache\")\n MOVIE: IconName = IconName(\"movie\")\n MOVIE_OFF: IconName = IconName(\"movie-off\")\n MUG: IconName = IconName(\"mug\")\n MUG_OFF: IconName = IconName(\"mug-off\")\n MULTIPLIER_0_5X: IconName = IconName(\"multiplier-0-5x\")\n MULTIPLIER_1_5X: IconName = IconName(\"multiplier-1-5x\")\n MULTIPLIER_1X: IconName = IconName(\"multiplier-1x\")\n MULTIPLIER_2X: IconName = IconName(\"multiplier-2x\")\n MUSHROOM: IconName = IconName(\"mushroom\")\n MUSHROOM_FILLED: IconName = IconName(\"mushroom-filled\")\n MUSHROOM_OFF: IconName = IconName(\"mushroom-off\")\n MUSIC: IconName = IconName(\"music\")\n MUSIC_OFF: IconName = IconName(\"music-off\")\n NAVIGATION: IconName = IconName(\"navigation\")\n NAVIGATION_FILLED: IconName = IconName(\"navigation-filled\")\n NAVIGATION_NORTH: IconName = IconName(\"navigation-north\")\n NAVIGATION_OFF: IconName = IconName(\"navigation-off\")\n NEEDLE: IconName = IconName(\"needle\")\n NEEDLE_THREAD: IconName = IconName(\"needle-thread\")\n NETWORK: IconName = IconName(\"network\")\n NETWORK_OFF: IconName = IconName(\"network-off\")\n NEW_SECTION: IconName = IconName(\"new-section\")\n NEWS: IconName = IconName(\"news\")\n NEWS_OFF: IconName = IconName(\"news-off\")\n NFC: IconName = IconName(\"nfc\")\n NFC_OFF: IconName = IconName(\"nfc-off\")\n NO_COPYRIGHT: IconName = IconName(\"no-copyright\")\n NO_CREATIVE_COMMONS: IconName = IconName(\"no-creative-commons\")\n NO_DERIVATIVES: IconName = IconName(\"no-derivatives\")\n NORTH_STAR: IconName = IconName(\"north-star\")\n NOTE: IconName = IconName(\"note\")\n NOTE_OFF: IconName = IconName(\"note-off\")\n NOTEBOOK: IconName = IconName(\"notebook\")\n NOTEBOOK_OFF: IconName = IconName(\"notebook-off\")\n NOTES: IconName = IconName(\"notes\")\n NOTES_OFF: IconName = IconName(\"notes-off\")\n NOTIFICATION: IconName = IconName(\"notification\")\n NOTIFICATION_OFF: IconName = IconName(\"notification-off\")\n NUMBER: IconName = IconName(\"number\")\n NUMBER_0: IconName = IconName(\"number-0\")\n NUMBER_1: IconName = IconName(\"number-1\")\n NUMBER_2: IconName = IconName(\"number-2\")\n NUMBER_3: IconName = IconName(\"number-3\")\n NUMBER_4: IconName = IconName(\"number-4\")\n NUMBER_5: IconName = IconName(\"number-5\")\n NUMBER_6: IconName = IconName(\"number-6\")\n NUMBER_7: IconName = IconName(\"number-7\")\n NUMBER_8: IconName = IconName(\"number-8\")\n NUMBER_9: IconName = IconName(\"number-9\")\n NUMBERS: IconName = IconName(\"numbers\")\n NURSE: IconName = IconName(\"nurse\")\n OCTAGON: IconName = IconName(\"octagon\")\n OCTAGON_FILLED: IconName = IconName(\"octagon-filled\")\n OCTAGON_OFF: IconName = IconName(\"octagon-off\")\n OCTAHEDRON: IconName = IconName(\"octahedron\")\n OCTAHEDRON_OFF: IconName = IconName(\"octahedron-off\")\n OCTAHEDRON_PLUS: IconName = IconName(\"octahedron-plus\")\n OLD: IconName = IconName(\"old\")\n OLYMPICS: IconName = IconName(\"olympics\")\n OLYMPICS_OFF: IconName = IconName(\"olympics-off\")\n OM: IconName = IconName(\"om\")\n OMEGA: IconName = IconName(\"omega\")\n OUTBOUND: IconName = IconName(\"outbound\")\n OUTLET: IconName = IconName(\"outlet\")\n OVAL: IconName = IconName(\"oval\")\n OVAL_FILLED: IconName = IconName(\"oval-filled\")\n OVAL_VERTICAL: IconName = IconName(\"oval-vertical\")\n OVAL_VERTICAL_FILLED: IconName = IconName(\"oval-vertical-filled\")\n OVERLINE: IconName = IconName(\"overline\")\n PACKAGE: IconName = IconName(\"package\")\n PACKAGE_EXPORT: IconName = IconName(\"package-export\")\n PACKAGE_IMPORT: IconName = IconName(\"package-import\")\n PACKAGE_OFF: IconName = IconName(\"package-off\")\n PACKAGES: IconName = IconName(\"packages\")\n PACMAN: IconName = IconName(\"pacman\")\n PAGE_BREAK: IconName = IconName(\"page-break\")\n PAINT: IconName = IconName(\"paint\")\n PAINT_FILLED: IconName = IconName(\"paint-filled\")\n PAINT_OFF: IconName = IconName(\"paint-off\")\n PALETTE: IconName = IconName(\"palette\")\n PALETTE_OFF: IconName = IconName(\"palette-off\")\n PANORAMA_HORIZONTAL: IconName = IconName(\"panorama-horizontal\")\n PANORAMA_HORIZONTAL_OFF: IconName = IconName(\"panorama-horizontal-off\")\n PANORAMA_VERTICAL: IconName = IconName(\"panorama-vertical\")\n PANORAMA_VERTICAL_OFF: IconName = IconName(\"panorama-vertical-off\")\n PAPER_BAG: IconName = IconName(\"paper-bag\")\n PAPER_BAG_OFF: IconName = IconName(\"paper-bag-off\")\n PAPERCLIP: IconName = IconName(\"paperclip\")\n PARACHUTE: IconName = IconName(\"parachute\")\n PARACHUTE_OFF: IconName = IconName(\"parachute-off\")\n PARENTHESES: IconName = IconName(\"parentheses\")\n PARENTHESES_OFF: IconName = IconName(\"parentheses-off\")\n PARKING: IconName = IconName(\"parking\")\n PARKING_OFF: IconName = IconName(\"parking-off\")\n PASSWORD: IconName = IconName(\"password\")\n PAW: IconName = IconName(\"paw\")\n PAW_FILLED: IconName = IconName(\"paw-filled\")\n PAW_OFF: IconName = IconName(\"paw-off\")\n PDF: IconName = IconName(\"pdf\")\n PEACE: IconName = IconName(\"peace\")\n PENCIL: IconName = IconName(\"pencil\")\n PENCIL_MINUS: IconName = IconName(\"pencil-minus\")\n PENCIL_OFF: IconName = IconName(\"pencil-off\")\n PENCIL_PLUS: IconName = IconName(\"pencil-plus\")\n PENNANT: IconName = IconName(\"pennant\")\n PENNANT_2: IconName = IconName(\"pennant-2\")\n PENNANT_2_FILLED: IconName = IconName(\"pennant-2-filled\")\n PENNANT_FILLED: IconName = IconName(\"pennant-filled\")\n PENNANT_OFF: IconName = IconName(\"pennant-off\")\n PENTAGON: IconName = IconName(\"pentagon\")\n PENTAGON_FILLED: IconName = IconName(\"pentagon-filled\")\n PENTAGON_OFF: IconName = IconName(\"pentagon-off\")\n PENTAGRAM: IconName = IconName(\"pentagram\")\n PEPPER: IconName = IconName(\"pepper\")\n PEPPER_OFF: IconName = IconName(\"pepper-off\")\n PERCENTAGE: IconName = IconName(\"percentage\")\n PERFUME: IconName = IconName(\"perfume\")\n PERSPECTIVE: IconName = IconName(\"perspective\")\n PERSPECTIVE_OFF: IconName = IconName(\"perspective-off\")\n PHONE: IconName = IconName(\"phone\")\n PHONE_CALL: IconName = IconName(\"phone-call\")\n PHONE_CALLING: IconName = IconName(\"phone-calling\")\n PHONE_CHECK: IconName = IconName(\"phone-check\")\n PHONE_FILLED: IconName = IconName(\"phone-filled\")\n PHONE_INCOMING: IconName = IconName(\"phone-incoming\")\n PHONE_OFF: IconName = IconName(\"phone-off\")\n PHONE_OUTGOING: IconName = IconName(\"phone-outgoing\")\n PHONE_PAUSE: IconName = IconName(\"phone-pause\")\n PHONE_PLUS: IconName = IconName(\"phone-plus\")\n PHONE_X: IconName = IconName(\"phone-x\")\n PHOTO: IconName = IconName(\"photo\")\n PHOTO_AI: IconName = IconName(\"photo-ai\")\n PHOTO_BOLT: IconName = IconName(\"photo-bolt\")\n PHOTO_CANCEL: IconName = IconName(\"photo-cancel\")\n PHOTO_CHECK: IconName = IconName(\"photo-check\")\n PHOTO_CODE: IconName = IconName(\"photo-code\")\n PHOTO_COG: IconName = IconName(\"photo-cog\")\n PHOTO_DOLLAR: IconName = IconName(\"photo-dollar\")\n PHOTO_DOWN: IconName = IconName(\"photo-down\")\n PHOTO_EDIT: IconName = IconName(\"photo-edit\")\n PHOTO_EXCLAMATION: IconName = IconName(\"photo-exclamation\")\n PHOTO_FILLED: IconName = IconName(\"photo-filled\")\n PHOTO_HEART: IconName = IconName(\"photo-heart\")\n PHOTO_MINUS: IconName = IconName(\"photo-minus\")\n PHOTO_OFF: IconName = IconName(\"photo-off\")\n PHOTO_PAUSE: IconName = IconName(\"photo-pause\")\n PHOTO_PIN: IconName = IconName(\"photo-pin\")\n PHOTO_PLUS: IconName = IconName(\"photo-plus\")\n PHOTO_QUESTION: IconName = IconName(\"photo-question\")\n PHOTO_SEARCH: IconName = IconName(\"photo-search\")\n PHOTO_SENSOR: IconName = IconName(\"photo-sensor\")\n PHOTO_SENSOR_2: IconName = IconName(\"photo-sensor-2\")\n PHOTO_SENSOR_3: IconName = IconName(\"photo-sensor-3\")\n PHOTO_SHARE: IconName = IconName(\"photo-share\")\n PHOTO_SHIELD: IconName = IconName(\"photo-shield\")\n PHOTO_STAR: IconName = IconName(\"photo-star\")\n PHOTO_UP: IconName = IconName(\"photo-up\")\n PHOTO_X: IconName = IconName(\"photo-x\")\n PHYSOTHERAPIST: IconName = IconName(\"physotherapist\")\n PIANO: IconName = IconName(\"piano\")\n PICK: IconName = IconName(\"pick\")\n PICTURE_IN_PICTURE: IconName = IconName(\"picture-in-picture\")\n PICTURE_IN_PICTURE_OFF: IconName = IconName(\"picture-in-picture-off\")\n PICTURE_IN_PICTURE_ON: IconName = IconName(\"picture-in-picture-on\")\n PICTURE_IN_PICTURE_TOP: IconName = IconName(\"picture-in-picture-top\")\n PIG: IconName = IconName(\"pig\")\n PIG_MONEY: IconName = IconName(\"pig-money\")\n PIG_OFF: IconName = IconName(\"pig-off\")\n PILCROW: IconName = IconName(\"pilcrow\")\n PILL: IconName = IconName(\"pill\")\n PILL_OFF: IconName = IconName(\"pill-off\")\n PILLS: IconName = IconName(\"pills\")\n PIN: IconName = IconName(\"pin\")\n PIN_FILLED: IconName = IconName(\"pin-filled\")\n PING_PONG: IconName = IconName(\"ping-pong\")\n PINNED: IconName = IconName(\"pinned\")\n PINNED_FILLED: IconName = IconName(\"pinned-filled\")\n PINNED_OFF: IconName = IconName(\"pinned-off\")\n PIZZA: IconName = IconName(\"pizza\")\n PIZZA_OFF: IconName = IconName(\"pizza-off\")\n PLACEHOLDER: IconName = IconName(\"placeholder\")\n PLANE: IconName = IconName(\"plane\")\n PLANE_ARRIVAL: IconName = IconName(\"plane-arrival\")\n PLANE_DEPARTURE: IconName = IconName(\"plane-departure\")\n PLANE_INFLIGHT: IconName = IconName(\"plane-inflight\")\n PLANE_OFF: IconName = IconName(\"plane-off\")\n PLANE_TILT: IconName = IconName(\"plane-tilt\")\n PLANET: IconName = IconName(\"planet\")\n PLANET_OFF: IconName = IconName(\"planet-off\")\n PLANT: IconName = IconName(\"plant\")\n PLANT_2: IconName = IconName(\"plant-2\")\n PLANT_2_OFF: IconName = IconName(\"plant-2-off\")\n PLANT_OFF: IconName = IconName(\"plant-off\")\n PLAY_BASKETBALL: IconName = IconName(\"play-basketball\")\n PLAY_CARD: IconName = IconName(\"play-card\")\n PLAY_CARD_OFF: IconName = IconName(\"play-card-off\")\n PLAY_FOOTBALL: IconName = IconName(\"play-football\")\n PLAY_HANDBALL: IconName = IconName(\"play-handball\")\n PLAY_VOLLEYBALL: IconName = IconName(\"play-volleyball\")\n PLAYER_EJECT: IconName = IconName(\"player-eject\")\n PLAYER_EJECT_FILLED: IconName = IconName(\"player-eject-filled\")\n PLAYER_PAUSE: IconName = IconName(\"player-pause\")\n PLAYER_PAUSE_FILLED: IconName = IconName(\"player-pause-filled\")\n PLAYER_PLAY: IconName = IconName(\"player-play\")\n PLAYER_PLAY_FILLED: IconName = IconName(\"player-play-filled\")\n PLAYER_RECORD: IconName = IconName(\"player-record\")\n PLAYER_RECORD_FILLED: IconName = IconName(\"player-record-filled\")\n PLAYER_SKIP_BACK: IconName = IconName(\"player-skip-back\")\n PLAYER_SKIP_BACK_FILLED: IconName = IconName(\"player-skip-back-filled\")\n PLAYER_SKIP_FORWARD: IconName = IconName(\"player-skip-forward\")\n PLAYER_SKIP_FORWARD_FILLED: IconName = IconName(\"player-skip-forward-filled\")\n PLAYER_STOP: IconName = IconName(\"player-stop\")\n PLAYER_STOP_FILLED: IconName = IconName(\"player-stop-filled\")\n PLAYER_TRACK_NEXT: IconName = IconName(\"player-track-next\")\n PLAYER_TRACK_NEXT_FILLED: IconName = IconName(\"player-track-next-filled\")\n PLAYER_TRACK_PREV: IconName = IconName(\"player-track-prev\")\n PLAYER_TRACK_PREV_FILLED: IconName = IconName(\"player-track-prev-filled\")\n PLAYLIST: IconName = IconName(\"playlist\")\n PLAYLIST_ADD: IconName = IconName(\"playlist-add\")\n PLAYLIST_OFF: IconName = IconName(\"playlist-off\")\n PLAYLIST_X: IconName = IconName(\"playlist-x\")\n PLAYSTATION_CIRCLE: IconName = IconName(\"playstation-circle\")\n PLAYSTATION_SQUARE: IconName = IconName(\"playstation-square\")\n PLAYSTATION_TRIANGLE: IconName = IconName(\"playstation-triangle\")\n PLAYSTATION_X: IconName = IconName(\"playstation-x\")\n PLUG: IconName = IconName(\"plug\")\n PLUG_CONNECTED: IconName = IconName(\"plug-connected\")\n PLUG_CONNECTED_X: IconName = IconName(\"plug-connected-x\")\n PLUG_OFF: IconName = IconName(\"plug-off\")\n PLUG_X: IconName = IconName(\"plug-x\")\n PLUS: IconName = IconName(\"plus\")\n PLUS_EQUAL: IconName = IconName(\"plus-equal\")\n PLUS_MINUS: IconName = IconName(\"plus-minus\")\n PNG: IconName = IconName(\"png\")\n PODIUM: IconName = IconName(\"podium\")\n PODIUM_OFF: IconName = IconName(\"podium-off\")\n POINT: IconName = IconName(\"point\")\n POINT_FILLED: IconName = IconName(\"point-filled\")\n POINT_OFF: IconName = IconName(\"point-off\")\n POINTER: IconName = IconName(\"pointer\")\n POINTER_BOLT: IconName = IconName(\"pointer-bolt\")\n POINTER_CANCEL: IconName = IconName(\"pointer-cancel\")\n POINTER_CHECK: IconName = IconName(\"pointer-check\")\n POINTER_CODE: IconName = IconName(\"pointer-code\")\n POINTER_COG: IconName = IconName(\"pointer-cog\")\n POINTER_DOLLAR: IconName = IconName(\"pointer-dollar\")\n POINTER_DOWN: IconName = IconName(\"pointer-down\")\n POINTER_EXCLAMATION: IconName = IconName(\"pointer-exclamation\")\n POINTER_HEART: IconName = IconName(\"pointer-heart\")\n POINTER_MINUS: IconName = IconName(\"pointer-minus\")\n POINTER_OFF: IconName = IconName(\"pointer-off\")\n POINTER_PAUSE: IconName = IconName(\"pointer-pause\")\n POINTER_PIN: IconName = IconName(\"pointer-pin\")\n POINTER_PLUS: IconName = IconName(\"pointer-plus\")\n POINTER_QUESTION: IconName = IconName(\"pointer-question\")\n POINTER_SEARCH: IconName = IconName(\"pointer-search\")\n POINTER_SHARE: IconName = IconName(\"pointer-share\")\n POINTER_STAR: IconName = IconName(\"pointer-star\")\n POINTER_UP: IconName = IconName(\"pointer-up\")\n POINTER_X: IconName = IconName(\"pointer-x\")\n POKEBALL: IconName = IconName(\"pokeball\")\n POKEBALL_OFF: IconName = IconName(\"pokeball-off\")\n POKER_CHIP: IconName = IconName(\"poker-chip\")\n POLAROID: IconName = IconName(\"polaroid\")\n POLAROID_FILLED: IconName = IconName(\"polaroid-filled\")\n POLYGON: IconName = IconName(\"polygon\")\n POLYGON_OFF: IconName = IconName(\"polygon-off\")\n POO: IconName = IconName(\"poo\")\n POOL: IconName = IconName(\"pool\")\n POOL_OFF: IconName = IconName(\"pool-off\")\n POWER: IconName = IconName(\"power\")\n PRAY: IconName = IconName(\"pray\")\n PREMIUM_RIGHTS: IconName = IconName(\"premium-rights\")\n PRESCRIPTION: IconName = IconName(\"prescription\")\n PRESENTATION: IconName = IconName(\"presentation\")\n PRESENTATION_ANALYTICS: IconName = IconName(\"presentation-analytics\")\n PRESENTATION_OFF: IconName = IconName(\"presentation-off\")\n PRINTER: IconName = IconName(\"printer\")\n PRINTER_OFF: IconName = IconName(\"printer-off\")\n PRISM: IconName = IconName(\"prism\")\n PRISM_OFF: IconName = IconName(\"prism-off\")\n PRISM_PLUS: IconName = IconName(\"prism-plus\")\n PRISON: IconName = IconName(\"prison\")\n PROGRESS: IconName = IconName(\"progress\")\n PROGRESS_ALERT: IconName = IconName(\"progress-alert\")\n PROGRESS_BOLT: IconName = IconName(\"progress-bolt\")\n PROGRESS_CHECK: IconName = IconName(\"progress-check\")\n PROGRESS_DOWN: IconName = IconName(\"progress-down\")\n PROGRESS_HELP: IconName = IconName(\"progress-help\")\n PROGRESS_X: IconName = IconName(\"progress-x\")\n PROMPT: IconName = IconName(\"prompt\")\n PROPELLER: IconName = IconName(\"propeller\")\n PROPELLER_OFF: IconName = IconName(\"propeller-off\")\n PUMPKIN_SCARY: IconName = IconName(\"pumpkin-scary\")\n PUZZLE: IconName = IconName(\"puzzle\")\n PUZZLE_2: IconName = IconName(\"puzzle-2\")\n PUZZLE_FILLED: IconName = IconName(\"puzzle-filled\")\n PUZZLE_OFF: IconName = IconName(\"puzzle-off\")\n PYRAMID: IconName = IconName(\"pyramid\")\n PYRAMID_OFF: IconName = IconName(\"pyramid-off\")\n PYRAMID_PLUS: IconName = IconName(\"pyramid-plus\")\n QRCODE: IconName = IconName(\"qrcode\")\n QRCODE_OFF: IconName = IconName(\"qrcode-off\")\n QUESTION_MARK: IconName = IconName(\"question-mark\")\n QUOTE: IconName = IconName(\"quote\")\n QUOTE_OFF: IconName = IconName(\"quote-off\")\n RADAR: IconName = IconName(\"radar\")\n RADAR_2: IconName = IconName(\"radar-2\")\n RADAR_OFF: IconName = IconName(\"radar-off\")\n RADIO: IconName = IconName(\"radio\")\n RADIO_OFF: IconName = IconName(\"radio-off\")\n RADIOACTIVE: IconName = IconName(\"radioactive\")\n RADIOACTIVE_FILLED: IconName = IconName(\"radioactive-filled\")\n RADIOACTIVE_OFF: IconName = IconName(\"radioactive-off\")\n RADIUS_BOTTOM_LEFT: IconName = IconName(\"radius-bottom-left\")\n RADIUS_BOTTOM_RIGHT: IconName = IconName(\"radius-bottom-right\")\n RADIUS_TOP_LEFT: IconName = IconName(\"radius-top-left\")\n RADIUS_TOP_RIGHT: IconName = IconName(\"radius-top-right\")\n RAINBOW: IconName = IconName(\"rainbow\")\n RAINBOW_OFF: IconName = IconName(\"rainbow-off\")\n RATING_12_PLUS: IconName = IconName(\"rating-12-plus\")\n RATING_14_PLUS: IconName = IconName(\"rating-14-plus\")\n RATING_16_PLUS: IconName = IconName(\"rating-16-plus\")\n RATING_18_PLUS: IconName = IconName(\"rating-18-plus\")\n RATING_21_PLUS: IconName = IconName(\"rating-21-plus\")\n RAZOR: IconName = IconName(\"razor\")\n RAZOR_ELECTRIC: IconName = IconName(\"razor-electric\")\n RECEIPT: IconName = IconName(\"receipt\")\n RECEIPT_2: IconName = IconName(\"receipt-2\")\n RECEIPT_OFF: IconName = IconName(\"receipt-off\")\n RECEIPT_REFUND: IconName = IconName(\"receipt-refund\")\n RECEIPT_TAX: IconName = IconName(\"receipt-tax\")\n RECHARGING: IconName = IconName(\"recharging\")\n RECORD_MAIL: IconName = IconName(\"record-mail\")\n RECORD_MAIL_OFF: IconName = IconName(\"record-mail-off\")\n RECTANGLE: IconName = IconName(\"rectangle\")\n RECTANGLE_FILLED: IconName = IconName(\"rectangle-filled\")\n RECTANGLE_ROUNDED_BOTTOM: IconName = IconName(\"rectangle-rounded-bottom\")\n RECTANGLE_ROUNDED_TOP: IconName = IconName(\"rectangle-rounded-top\")\n RECTANGLE_VERTICAL: IconName = IconName(\"rectangle-vertical\")\n RECTANGLE_VERTICAL_FILLED: IconName = IconName(\"rectangle-vertical-filled\")\n RECTANGULAR_PRISM: IconName = IconName(\"rectangular-prism\")\n RECTANGULAR_PRISM_OFF: IconName = IconName(\"rectangular-prism-off\")\n RECTANGULAR_PRISM_PLUS: IconName = IconName(\"rectangular-prism-plus\")\n RECYCLE: IconName = IconName(\"recycle\")\n RECYCLE_OFF: IconName = IconName(\"recycle-off\")\n REFRESH: IconName = IconName(\"refresh\")\n REFRESH_ALERT: IconName = IconName(\"refresh-alert\")\n REFRESH_DOT: IconName = IconName(\"refresh-dot\")\n REFRESH_OFF: IconName = IconName(\"refresh-off\")\n REGEX: IconName = IconName(\"regex\")\n REGEX_OFF: IconName = IconName(\"regex-off\")\n REGISTERED: IconName = IconName(\"registered\")\n RELATION_MANY_TO_MANY: IconName = IconName(\"relation-many-to-many\")\n RELATION_ONE_TO_MANY: IconName = IconName(\"relation-one-to-many\")\n RELATION_ONE_TO_ONE: IconName = IconName(\"relation-one-to-one\")\n RELOAD: IconName = IconName(\"reload\")\n REPEAT: IconName = IconName(\"repeat\")\n REPEAT_OFF: IconName = IconName(\"repeat-off\")\n REPEAT_ONCE: IconName = IconName(\"repeat-once\")\n REPLACE: IconName = IconName(\"replace\")\n REPLACE_FILLED: IconName = IconName(\"replace-filled\")\n REPLACE_OFF: IconName = IconName(\"replace-off\")\n REPORT: IconName = IconName(\"report\")\n REPORT_ANALYTICS: IconName = IconName(\"report-analytics\")\n REPORT_MEDICAL: IconName = IconName(\"report-medical\")\n REPORT_MONEY: IconName = IconName(\"report-money\")\n REPORT_OFF: IconName = IconName(\"report-off\")\n REPORT_SEARCH: IconName = IconName(\"report-search\")\n RESERVED_LINE: IconName = IconName(\"reserved-line\")\n RESIZE: IconName = IconName(\"resize\")\n RESTORE: IconName = IconName(\"restore\")\n REWIND_BACKWARD_10: IconName = IconName(\"rewind-backward-10\")\n REWIND_BACKWARD_15: IconName = IconName(\"rewind-backward-15\")\n REWIND_BACKWARD_20: IconName = IconName(\"rewind-backward-20\")\n REWIND_BACKWARD_30: IconName = IconName(\"rewind-backward-30\")\n REWIND_BACKWARD_40: IconName = IconName(\"rewind-backward-40\")\n REWIND_BACKWARD_5: IconName = IconName(\"rewind-backward-5\")\n REWIND_BACKWARD_50: IconName = IconName(\"rewind-backward-50\")\n REWIND_BACKWARD_60: IconName = IconName(\"rewind-backward-60\")\n REWIND_FORWARD_10: IconName = IconName(\"rewind-forward-10\")\n REWIND_FORWARD_15: IconName = IconName(\"rewind-forward-15\")\n REWIND_FORWARD_20: IconName = IconName(\"rewind-forward-20\")\n REWIND_FORWARD_30: IconName = IconName(\"rewind-forward-30\")\n REWIND_FORWARD_40: IconName = IconName(\"rewind-forward-40\")\n REWIND_FORWARD_5: IconName = IconName(\"rewind-forward-5\")\n REWIND_FORWARD_50: IconName = IconName(\"rewind-forward-50\")\n REWIND_FORWARD_60: IconName = IconName(\"rewind-forward-60\")\n RIBBON_HEALTH: IconName = IconName(\"ribbon-health\")\n RINGS: IconName = IconName(\"rings\")\n RIPPLE: IconName = IconName(\"ripple\")\n RIPPLE_OFF: IconName = IconName(\"ripple-off\")\n ROAD: IconName = IconName(\"road\")\n ROAD_OFF: IconName = IconName(\"road-off\")\n ROAD_SIGN: IconName = IconName(\"road-sign\")\n ROBOT: IconName = IconName(\"robot\")\n ROBOT_OFF: IconName = IconName(\"robot-off\")\n ROCKET: IconName = IconName(\"rocket\")\n ROCKET_OFF: IconName = IconName(\"rocket-off\")\n ROLLER_SKATING: IconName = IconName(\"roller-skating\")\n ROLLERCOASTER: IconName = IconName(\"rollercoaster\")\n ROLLERCOASTER_OFF: IconName = IconName(\"rollercoaster-off\")\n ROSETTE: IconName = IconName(\"rosette\")\n ROSETTE_FILLED: IconName = IconName(\"rosette-filled\")\n ROSETTE_NUMBER_0: IconName = IconName(\"rosette-number-0\")\n ROSETTE_NUMBER_1: IconName = IconName(\"rosette-number-1\")\n ROSETTE_NUMBER_2: IconName = IconName(\"rosette-number-2\")\n ROSETTE_NUMBER_3: IconName = IconName(\"rosette-number-3\")\n ROSETTE_NUMBER_4: IconName = IconName(\"rosette-number-4\")\n ROSETTE_NUMBER_5: IconName = IconName(\"rosette-number-5\")\n ROSETTE_NUMBER_6: IconName = IconName(\"rosette-number-6\")\n ROSETTE_NUMBER_7: IconName = IconName(\"rosette-number-7\")\n ROSETTE_NUMBER_8: IconName = IconName(\"rosette-number-8\")\n ROSETTE_NUMBER_9: IconName = IconName(\"rosette-number-9\")\n ROTATE: IconName = IconName(\"rotate\")\n ROTATE_2: IconName = IconName(\"rotate-2\")\n ROTATE_360: IconName = IconName(\"rotate-360\")\n ROTATE_CLOCKWISE: IconName = IconName(\"rotate-clockwise\")\n ROTATE_CLOCKWISE_2: IconName = IconName(\"rotate-clockwise-2\")\n ROTATE_DOT: IconName = IconName(\"rotate-dot\")\n ROTATE_RECTANGLE: IconName = IconName(\"rotate-rectangle\")\n ROUTE: IconName = IconName(\"route\")\n ROUTE_2: IconName = IconName(\"route-2\")\n ROUTE_OFF: IconName = IconName(\"route-off\")\n ROUTER: IconName = IconName(\"router\")\n ROUTER_OFF: IconName = IconName(\"router-off\")\n ROW_INSERT_BOTTOM: IconName = IconName(\"row-insert-bottom\")\n ROW_INSERT_TOP: IconName = IconName(\"row-insert-top\")\n ROW_REMOVE: IconName = IconName(\"row-remove\")\n RSS: IconName = IconName(\"rss\")\n RUBBER_STAMP: IconName = IconName(\"rubber-stamp\")\n RUBBER_STAMP_OFF: IconName = IconName(\"rubber-stamp-off\")\n RULER: IconName = IconName(\"ruler\")\n RULER_2: IconName = IconName(\"ruler-2\")\n RULER_2_OFF: IconName = IconName(\"ruler-2-off\")\n RULER_3: IconName = IconName(\"ruler-3\")\n RULER_MEASURE: IconName = IconName(\"ruler-measure\")\n RULER_OFF: IconName = IconName(\"ruler-off\")\n RUN: IconName = IconName(\"run\")\n S_TURN_DOWN: IconName = IconName(\"s-turn-down\")\n S_TURN_LEFT: IconName = IconName(\"s-turn-left\")\n S_TURN_RIGHT: IconName = IconName(\"s-turn-right\")\n S_TURN_UP: IconName = IconName(\"s-turn-up\")\n SAILBOAT: IconName = IconName(\"sailboat\")\n SAILBOAT_2: IconName = IconName(\"sailboat-2\")\n SAILBOAT_OFF: IconName = IconName(\"sailboat-off\")\n SALAD: IconName = IconName(\"salad\")\n SALT: IconName = IconName(\"salt\")\n SATELLITE: IconName = IconName(\"satellite\")\n SATELLITE_OFF: IconName = IconName(\"satellite-off\")\n SAUSAGE: IconName = IconName(\"sausage\")\n SCALE: IconName = IconName(\"scale\")\n SCALE_OFF: IconName = IconName(\"scale-off\")\n SCALE_OUTLINE: IconName = IconName(\"scale-outline\")\n SCALE_OUTLINE_OFF: IconName = IconName(\"scale-outline-off\")\n SCAN: IconName = IconName(\"scan\")\n SCAN_EYE: IconName = IconName(\"scan-eye\")\n SCHEMA: IconName = IconName(\"schema\")\n SCHEMA_OFF: IconName = IconName(\"schema-off\")\n SCHOOL: IconName = IconName(\"school\")\n SCHOOL_BELL: IconName = IconName(\"school-bell\")\n SCHOOL_OFF: IconName = IconName(\"school-off\")\n SCISSORS: IconName = IconName(\"scissors\")\n SCISSORS_OFF: IconName = IconName(\"scissors-off\")\n SCOOTER: IconName = IconName(\"scooter\")\n SCOOTER_ELECTRIC: IconName = IconName(\"scooter-electric\")\n SCOREBOARD: IconName = IconName(\"scoreboard\")\n SCREEN_SHARE: IconName = IconName(\"screen-share\")\n SCREEN_SHARE_OFF: IconName = IconName(\"screen-share-off\")\n SCREENSHOT: IconName = IconName(\"screenshot\")\n SCRIBBLE: IconName = IconName(\"scribble\")\n SCRIBBLE_OFF: IconName = IconName(\"scribble-off\")\n SCRIPT: IconName = IconName(\"script\")\n SCRIPT_MINUS: IconName = IconName(\"script-minus\")\n SCRIPT_PLUS: IconName = IconName(\"script-plus\")\n SCRIPT_X: IconName = IconName(\"script-x\")\n SCUBA_MASK: IconName = IconName(\"scuba-mask\")\n SCUBA_MASK_OFF: IconName = IconName(\"scuba-mask-off\")\n SDK: IconName = IconName(\"sdk\")\n SEARCH: IconName = IconName(\"search\")\n SEARCH_OFF: IconName = IconName(\"search-off\")\n SECTION: IconName = IconName(\"section\")\n SECTION_SIGN: IconName = IconName(\"section-sign\")\n SEEDING: IconName = IconName(\"seeding\")\n SEEDING_OFF: IconName = IconName(\"seeding-off\")\n SELECT: IconName = IconName(\"select\")\n SELECT_ALL: IconName = IconName(\"select-all\")\n SELECTOR: IconName = IconName(\"selector\")\n SEND: IconName = IconName(\"send\")\n SEND_OFF: IconName = IconName(\"send-off\")\n SEO: IconName = IconName(\"seo\")\n SEPARATOR: IconName = IconName(\"separator\")\n SEPARATOR_HORIZONTAL: IconName = IconName(\"separator-horizontal\")\n SEPARATOR_VERTICAL: IconName = IconName(\"separator-vertical\")\n SERVER: IconName = IconName(\"server\")\n SERVER_2: IconName = IconName(\"server-2\")\n SERVER_BOLT: IconName = IconName(\"server-bolt\")\n SERVER_COG: IconName = IconName(\"server-cog\")\n SERVER_OFF: IconName = IconName(\"server-off\")\n SERVICEMARK: IconName = IconName(\"servicemark\")\n SETTINGS: IconName = IconName(\"settings\")\n SETTINGS_2: IconName = IconName(\"settings-2\")\n SETTINGS_AUTOMATION: IconName = IconName(\"settings-automation\")\n SETTINGS_BOLT: IconName = IconName(\"settings-bolt\")\n SETTINGS_CANCEL: IconName = IconName(\"settings-cancel\")\n SETTINGS_CHECK: IconName = IconName(\"settings-check\")\n SETTINGS_CODE: IconName = IconName(\"settings-code\")\n SETTINGS_COG: IconName = IconName(\"settings-cog\")\n SETTINGS_DOLLAR: IconName = IconName(\"settings-dollar\")\n SETTINGS_DOWN: IconName = IconName(\"settings-down\")\n SETTINGS_EXCLAMATION: IconName = IconName(\"settings-exclamation\")\n SETTINGS_FILLED: IconName = IconName(\"settings-filled\")\n SETTINGS_HEART: IconName = IconName(\"settings-heart\")\n SETTINGS_MINUS: IconName = IconName(\"settings-minus\")\n SETTINGS_OFF: IconName = IconName(\"settings-off\")\n SETTINGS_PAUSE: IconName = IconName(\"settings-pause\")\n SETTINGS_PIN: IconName = IconName(\"settings-pin\")\n SETTINGS_PLUS: IconName = IconName(\"settings-plus\")\n SETTINGS_QUESTION: IconName = IconName(\"settings-question\")\n SETTINGS_SEARCH: IconName = IconName(\"settings-search\")\n SETTINGS_SHARE: IconName = IconName(\"settings-share\")\n SETTINGS_STAR: IconName = IconName(\"settings-star\")\n SETTINGS_UP: IconName = IconName(\"settings-up\")\n SETTINGS_X: IconName = IconName(\"settings-x\")\n SHADOW: IconName = IconName(\"shadow\")\n SHADOW_OFF: IconName = IconName(\"shadow-off\")\n SHAPE: IconName = IconName(\"shape\")\n SHAPE_2: IconName = IconName(\"shape-2\")\n SHAPE_3: IconName = IconName(\"shape-3\")\n SHAPE_OFF: IconName = IconName(\"shape-off\")\n SHARE: IconName = IconName(\"share\")\n SHARE_2: IconName = IconName(\"share-2\")\n SHARE_3: IconName = IconName(\"share-3\")\n SHARE_OFF: IconName = IconName(\"share-off\")\n SHI_JUMPING: IconName = IconName(\"shi-jumping\")\n SHIELD: IconName = IconName(\"shield\")\n SHIELD_BOLT: IconName = IconName(\"shield-bolt\")\n SHIELD_CANCEL: IconName = IconName(\"shield-cancel\")\n SHIELD_CHECK: IconName = IconName(\"shield-check\")\n SHIELD_CHECK_FILLED: IconName = IconName(\"shield-check-filled\")\n SHIELD_CHECKERED: IconName = IconName(\"shield-checkered\")\n SHIELD_CHECKERED_FILLED: IconName = IconName(\"shield-checkered-filled\")\n SHIELD_CHEVRON: IconName = IconName(\"shield-chevron\")\n SHIELD_CODE: IconName = IconName(\"shield-code\")\n SHIELD_COG: IconName = IconName(\"shield-cog\")\n SHIELD_DOLLAR: IconName = IconName(\"shield-dollar\")\n SHIELD_DOWN: IconName = IconName(\"shield-down\")\n SHIELD_EXCLAMATION: IconName = IconName(\"shield-exclamation\")\n SHIELD_FILLED: IconName = IconName(\"shield-filled\")\n SHIELD_HALF: IconName = IconName(\"shield-half\")\n SHIELD_HALF_FILLED: IconName = IconName(\"shield-half-filled\")\n SHIELD_HEART: IconName = IconName(\"shield-heart\")\n SHIELD_LOCK: IconName = IconName(\"shield-lock\")\n SHIELD_LOCK_FILLED: IconName = IconName(\"shield-lock-filled\")\n SHIELD_MINUS: IconName = IconName(\"shield-minus\")\n SHIELD_OFF: IconName = IconName(\"shield-off\")\n SHIELD_PAUSE: IconName = IconName(\"shield-pause\")\n SHIELD_PIN: IconName = IconName(\"shield-pin\")\n SHIELD_PLUS: IconName = IconName(\"shield-plus\")\n SHIELD_QUESTION: IconName = IconName(\"shield-question\")\n SHIELD_SEARCH: IconName = IconName(\"shield-search\")\n SHIELD_SHARE: IconName = IconName(\"shield-share\")\n SHIELD_STAR: IconName = IconName(\"shield-star\")\n SHIELD_UP: IconName = IconName(\"shield-up\")\n SHIELD_X: IconName = IconName(\"shield-x\")\n SHIP: IconName = IconName(\"ship\")\n SHIP_OFF: IconName = IconName(\"ship-off\")\n SHIRT: IconName = IconName(\"shirt\")\n SHIRT_FILLED: IconName = IconName(\"shirt-filled\")\n SHIRT_OFF: IconName = IconName(\"shirt-off\")\n SHIRT_SPORT: IconName = IconName(\"shirt-sport\")\n SHOE: IconName = IconName(\"shoe\")\n SHOE_OFF: IconName = IconName(\"shoe-off\")\n SHOPPING_BAG: IconName = IconName(\"shopping-bag\")\n SHOPPING_CART: IconName = IconName(\"shopping-cart\")\n SHOPPING_CART_DISCOUNT: IconName = IconName(\"shopping-cart-discount\")\n SHOPPING_CART_OFF: IconName = IconName(\"shopping-cart-off\")\n SHOPPING_CART_PLUS: IconName = IconName(\"shopping-cart-plus\")\n SHOPPING_CART_X: IconName = IconName(\"shopping-cart-x\")\n SHOVEL: IconName = IconName(\"shovel\")\n SHREDDER: IconName = IconName(\"shredder\")\n SIGN_LEFT: IconName = IconName(\"sign-left\")\n SIGN_LEFT_FILLED: IconName = IconName(\"sign-left-filled\")\n SIGN_RIGHT: IconName = IconName(\"sign-right\")\n SIGN_RIGHT_FILLED: IconName = IconName(\"sign-right-filled\")\n SIGNAL_2G: IconName = IconName(\"signal-2g\")\n SIGNAL_3G: IconName = IconName(\"signal-3g\")\n SIGNAL_4G: IconName = IconName(\"signal-4g\")\n SIGNAL_4G_PLUS: IconName = IconName(\"signal-4g-plus\")\n SIGNAL_5G: IconName = IconName(\"signal-5g\")\n SIGNAL_6G: IconName = IconName(\"signal-6g\")\n SIGNAL_E: IconName = IconName(\"signal-e\")\n SIGNAL_G: IconName = IconName(\"signal-g\")\n SIGNAL_H: IconName = IconName(\"signal-h\")\n SIGNAL_H_PLUS: IconName = IconName(\"signal-h-plus\")\n SIGNAL_LTE: IconName = IconName(\"signal-lte\")\n SIGNATURE: IconName = IconName(\"signature\")\n SIGNATURE_OFF: IconName = IconName(\"signature-off\")\n SITEMAP: IconName = IconName(\"sitemap\")\n SITEMAP_OFF: IconName = IconName(\"sitemap-off\")\n SKATEBOARD: IconName = IconName(\"skateboard\")\n SKATEBOARD_OFF: IconName = IconName(\"skateboard-off\")\n SKATEBOARDING: IconName = IconName(\"skateboarding\")\n SKULL: IconName = IconName(\"skull\")\n SLASH: IconName = IconName(\"slash\")\n SLASHES: IconName = IconName(\"slashes\")\n SLEIGH: IconName = IconName(\"sleigh\")\n SLICE: IconName = IconName(\"slice\")\n SLIDESHOW: IconName = IconName(\"slideshow\")\n SMART_HOME: IconName = IconName(\"smart-home\")\n SMART_HOME_OFF: IconName = IconName(\"smart-home-off\")\n SMOKING: IconName = IconName(\"smoking\")\n SMOKING_NO: IconName = IconName(\"smoking-no\")\n SNOWFLAKE: IconName = IconName(\"snowflake\")\n SNOWFLAKE_OFF: IconName = IconName(\"snowflake-off\")\n SNOWMAN: IconName = IconName(\"snowman\")\n SOCCER_FIELD: IconName = IconName(\"soccer-field\")\n SOCIAL: IconName = IconName(\"social\")\n SOCIAL_OFF: IconName = IconName(\"social-off\")\n SOCK: IconName = IconName(\"sock\")\n SOFA: IconName = IconName(\"sofa\")\n SOFA_OFF: IconName = IconName(\"sofa-off\")\n SOLAR_PANEL: IconName = IconName(\"solar-panel\")\n SOLAR_PANEL_2: IconName = IconName(\"solar-panel-2\")\n SORT_0_9: IconName = IconName(\"sort-0-9\")\n SORT_9_0: IconName = IconName(\"sort-9-0\")\n SORT_A_Z: IconName = IconName(\"sort-a-z\")\n SORT_ASCENDING: IconName = IconName(\"sort-ascending\")\n SORT_ASCENDING_2: IconName = IconName(\"sort-ascending-2\")\n SORT_ASCENDING_LETTERS: IconName = IconName(\"sort-ascending-letters\")\n SORT_ASCENDING_NUMBERS: IconName = IconName(\"sort-ascending-numbers\")\n SORT_DESCENDING: IconName = IconName(\"sort-descending\")\n SORT_DESCENDING_2: IconName = IconName(\"sort-descending-2\")\n SORT_DESCENDING_LETTERS: IconName = IconName(\"sort-descending-letters\")\n SORT_DESCENDING_NUMBERS: IconName = IconName(\"sort-descending-numbers\")\n SORT_Z_A: IconName = IconName(\"sort-z-a\")\n SOS: IconName = IconName(\"sos\")\n SOUP: IconName = IconName(\"soup\")\n SOUP_OFF: IconName = IconName(\"soup-off\")\n SOURCE_CODE: IconName = IconName(\"source-code\")\n SPACE: IconName = IconName(\"space\")\n SPACE_OFF: IconName = IconName(\"space-off\")\n SPACING_HORIZONTAL: IconName = IconName(\"spacing-horizontal\")\n SPACING_VERTICAL: IconName = IconName(\"spacing-vertical\")\n SPADE: IconName = IconName(\"spade\")\n SPADE_FILLED: IconName = IconName(\"spade-filled\")\n SPARKLES: IconName = IconName(\"sparkles\")\n SPEAKERPHONE: IconName = IconName(\"speakerphone\")\n SPEEDBOAT: IconName = IconName(\"speedboat\")\n SPHERE: IconName = IconName(\"sphere\")\n SPHERE_OFF: IconName = IconName(\"sphere-off\")\n SPHERE_PLUS: IconName = IconName(\"sphere-plus\")\n SPIDER: IconName = IconName(\"spider\")\n SPIRAL: IconName = IconName(\"spiral\")\n SPIRAL_OFF: IconName = IconName(\"spiral-off\")\n SPORT_BILLARD: IconName = IconName(\"sport-billard\")\n SPRAY: IconName = IconName(\"spray\")\n SPY: IconName = IconName(\"spy\")\n SPY_OFF: IconName = IconName(\"spy-off\")\n SQL: IconName = IconName(\"sql\")\n SQUARE: IconName = IconName(\"square\")\n SQUARE_0_FILLED: IconName = IconName(\"square-0-filled\")\n SQUARE_1_FILLED: IconName = IconName(\"square-1-filled\")\n SQUARE_2_FILLED: IconName = IconName(\"square-2-filled\")\n SQUARE_3_FILLED: IconName = IconName(\"square-3-filled\")\n SQUARE_4_FILLED: IconName = IconName(\"square-4-filled\")\n SQUARE_5_FILLED: IconName = IconName(\"square-5-filled\")\n SQUARE_6_FILLED: IconName = IconName(\"square-6-filled\")\n SQUARE_7_FILLED: IconName = IconName(\"square-7-filled\")\n SQUARE_8_FILLED: IconName = IconName(\"square-8-filled\")\n SQUARE_9_FILLED: IconName = IconName(\"square-9-filled\")\n SQUARE_ARROW_DOWN: IconName = IconName(\"square-arrow-down\")\n SQUARE_ARROW_LEFT: IconName = IconName(\"square-arrow-left\")\n SQUARE_ARROW_RIGHT: IconName = IconName(\"square-arrow-right\")\n SQUARE_ARROW_UP: IconName = IconName(\"square-arrow-up\")\n SQUARE_ASTERISK: IconName = IconName(\"square-asterisk\")\n SQUARE_CHECK: IconName = IconName(\"square-check\")\n SQUARE_CHECK_FILLED: IconName = IconName(\"square-check-filled\")\n SQUARE_CHEVRON_DOWN: IconName = IconName(\"square-chevron-down\")\n SQUARE_CHEVRON_LEFT: IconName = IconName(\"square-chevron-left\")\n SQUARE_CHEVRON_RIGHT: IconName = IconName(\"square-chevron-right\")\n SQUARE_CHEVRON_UP: IconName = IconName(\"square-chevron-up\")\n SQUARE_CHEVRONS_DOWN: IconName = IconName(\"square-chevrons-down\")\n SQUARE_CHEVRONS_LEFT: IconName = IconName(\"square-chevrons-left\")\n SQUARE_CHEVRONS_RIGHT: IconName = IconName(\"square-chevrons-right\")\n SQUARE_CHEVRONS_UP: IconName = IconName(\"square-chevrons-up\")\n SQUARE_DOT: IconName = IconName(\"square-dot\")\n SQUARE_F0: IconName = IconName(\"square-f0\")\n SQUARE_F0_FILLED: IconName = IconName(\"square-f0-filled\")\n SQUARE_F1: IconName = IconName(\"square-f1\")\n SQUARE_F1_FILLED: IconName = IconName(\"square-f1-filled\")\n SQUARE_F2: IconName = IconName(\"square-f2\")\n SQUARE_F2_FILLED: IconName = IconName(\"square-f2-filled\")\n SQUARE_F3: IconName = IconName(\"square-f3\")\n SQUARE_F3_FILLED: IconName = IconName(\"square-f3-filled\")\n SQUARE_F4: IconName = IconName(\"square-f4\")\n SQUARE_F4_FILLED: IconName = IconName(\"square-f4-filled\")\n SQUARE_F5: IconName = IconName(\"square-f5\")\n SQUARE_F5_FILLED: IconName = IconName(\"square-f5-filled\")\n SQUARE_F6: IconName = IconName(\"square-f6\")\n SQUARE_F6_FILLED: IconName = IconName(\"square-f6-filled\")\n SQUARE_F7: IconName = IconName(\"square-f7\")\n SQUARE_F7_FILLED: IconName = IconName(\"square-f7-filled\")\n SQUARE_F8: IconName = IconName(\"square-f8\")\n SQUARE_F8_FILLED: IconName = IconName(\"square-f8-filled\")\n SQUARE_F9: IconName = IconName(\"square-f9\")\n SQUARE_F9_FILLED: IconName = IconName(\"square-f9-filled\")\n SQUARE_FORBID: IconName = IconName(\"square-forbid\")\n SQUARE_FORBID_2: IconName = IconName(\"square-forbid-2\")\n SQUARE_HALF: IconName = IconName(\"square-half\")\n SQUARE_KEY: IconName = IconName(\"square-key\")\n SQUARE_LETTER_A: IconName = IconName(\"square-letter-a\")\n SQUARE_LETTER_B: IconName = IconName(\"square-letter-b\")\n SQUARE_LETTER_C: IconName = IconName(\"square-letter-c\")\n SQUARE_LETTER_D: IconName = IconName(\"square-letter-d\")\n SQUARE_LETTER_E: IconName = IconName(\"square-letter-e\")\n SQUARE_LETTER_F: IconName = IconName(\"square-letter-f\")\n SQUARE_LETTER_G: IconName = IconName(\"square-letter-g\")\n SQUARE_LETTER_H: IconName = IconName(\"square-letter-h\")\n SQUARE_LETTER_I: IconName = IconName(\"square-letter-i\")\n SQUARE_LETTER_J: IconName = IconName(\"square-letter-j\")\n SQUARE_LETTER_K: IconName = IconName(\"square-letter-k\")\n SQUARE_LETTER_L: IconName = IconName(\"square-letter-l\")\n SQUARE_LETTER_M: IconName = IconName(\"square-letter-m\")\n SQUARE_LETTER_N: IconName = IconName(\"square-letter-n\")\n SQUARE_LETTER_O: IconName = IconName(\"square-letter-o\")\n SQUARE_LETTER_P: IconName = IconName(\"square-letter-p\")\n SQUARE_LETTER_Q: IconName = IconName(\"square-letter-q\")\n SQUARE_LETTER_R: IconName = IconName(\"square-letter-r\")\n SQUARE_LETTER_S: IconName = IconName(\"square-letter-s\")\n SQUARE_LETTER_T: IconName = IconName(\"square-letter-t\")\n SQUARE_LETTER_U: IconName = IconName(\"square-letter-u\")\n SQUARE_LETTER_V: IconName = IconName(\"square-letter-v\")\n SQUARE_LETTER_W: IconName = IconName(\"square-letter-w\")\n SQUARE_LETTER_X: IconName = IconName(\"square-letter-x\")\n SQUARE_LETTER_Y: IconName = IconName(\"square-letter-y\")\n SQUARE_LETTER_Z: IconName = IconName(\"square-letter-z\")\n SQUARE_MINUS: IconName = IconName(\"square-minus\")\n SQUARE_NUMBER_0: IconName = IconName(\"square-number-0\")\n SQUARE_NUMBER_1: IconName = IconName(\"square-number-1\")\n SQUARE_NUMBER_2: IconName = IconName(\"square-number-2\")\n SQUARE_NUMBER_3: IconName = IconName(\"square-number-3\")\n SQUARE_NUMBER_4: IconName = IconName(\"square-number-4\")\n SQUARE_NUMBER_5: IconName = IconName(\"square-number-5\")\n SQUARE_NUMBER_6: IconName = IconName(\"square-number-6\")\n SQUARE_NUMBER_7: IconName = IconName(\"square-number-7\")\n SQUARE_NUMBER_8: IconName = IconName(\"square-number-8\")\n SQUARE_NUMBER_9: IconName = IconName(\"square-number-9\")\n SQUARE_OFF: IconName = IconName(\"square-off\")\n SQUARE_PLUS: IconName = IconName(\"square-plus\")\n SQUARE_ROOT: IconName = IconName(\"square-root\")\n SQUARE_ROOT_2: IconName = IconName(\"square-root-2\")\n SQUARE_ROTATED: IconName = IconName(\"square-rotated\")\n SQUARE_ROTATED_FILLED: IconName = IconName(\"square-rotated-filled\")\n SQUARE_ROTATED_FORBID: IconName = IconName(\"square-rotated-forbid\")\n SQUARE_ROTATED_FORBID_2: IconName = IconName(\"square-rotated-forbid-2\")\n SQUARE_ROTATED_OFF: IconName = IconName(\"square-rotated-off\")\n SQUARE_ROUNDED: IconName = IconName(\"square-rounded\")\n SQUARE_ROUNDED_ARROW_DOWN: IconName = IconName(\"square-rounded-arrow-down\")\n SQUARE_ROUNDED_ARROW_DOWN_FILLED: IconName = IconName(\n \"square-rounded-arrow-down-filled\"\n )\n SQUARE_ROUNDED_ARROW_LEFT: IconName = IconName(\"square-rounded-arrow-left\")\n SQUARE_ROUNDED_ARROW_LEFT_FILLED: IconName = IconName(\n \"square-rounded-arrow-left-filled\"\n )\n SQUARE_ROUNDED_ARROW_RIGHT: IconName = IconName(\"square-rounded-arrow-right\")\n SQUARE_ROUNDED_ARROW_RIGHT_FILLED: IconName = IconName(\n \"square-rounded-arrow-right-filled\"\n )\n SQUARE_ROUNDED_ARROW_UP: IconName = IconName(\"square-rounded-arrow-up\")\n SQUARE_ROUNDED_ARROW_UP_FILLED: IconName = IconName(\n \"square-rounded-arrow-up-filled\"\n )\n SQUARE_ROUNDED_CHECK: IconName = IconName(\"square-rounded-check\")\n SQUARE_ROUNDED_CHECK_FILLED: IconName = IconName(\"square-rounded-check-filled\")\n SQUARE_ROUNDED_CHEVRON_DOWN: IconName = IconName(\"square-rounded-chevron-down\")\n SQUARE_ROUNDED_CHEVRON_DOWN_FILLED: IconName = IconName(\n \"square-rounded-chevron-down-filled\"\n )\n SQUARE_ROUNDED_CHEVRON_LEFT: IconName = IconName(\"square-rounded-chevron-left\")\n SQUARE_ROUNDED_CHEVRON_LEFT_FILLED: IconName = IconName(\n \"square-rounded-chevron-left-filled\"\n )\n SQUARE_ROUNDED_CHEVRON_RIGHT: IconName = IconName(\"square-rounded-chevron-right\")\n SQUARE_ROUNDED_CHEVRON_RIGHT_FILLED: IconName = IconName(\n \"square-rounded-chevron-right-filled\"\n )\n SQUARE_ROUNDED_CHEVRON_UP: IconName = IconName(\"square-rounded-chevron-up\")\n SQUARE_ROUNDED_CHEVRON_UP_FILLED: IconName = IconName(\n \"square-rounded-chevron-up-filled\"\n )\n SQUARE_ROUNDED_CHEVRONS_DOWN: IconName = IconName(\"square-rounded-chevrons-down\")\n SQUARE_ROUNDED_CHEVRONS_DOWN_FILLED: IconName = IconName(\n \"square-rounded-chevrons-down-filled\"\n )\n SQUARE_ROUNDED_CHEVRONS_LEFT: IconName = IconName(\"square-rounded-chevrons-left\")\n SQUARE_ROUNDED_CHEVRONS_LEFT_FILLED: IconName = IconName(\n \"square-rounded-chevrons-left-filled\"\n )\n SQUARE_ROUNDED_CHEVRONS_RIGHT: IconName = IconName(\"square-rounded-chevrons-right\")\n SQUARE_ROUNDED_CHEVRONS_RIGHT_FILLED: IconName = IconName(\n \"square-rounded-chevrons-right-filled\"\n )\n SQUARE_ROUNDED_CHEVRONS_UP: IconName = IconName(\"square-rounded-chevrons-up\")\n SQUARE_ROUNDED_CHEVRONS_UP_FILLED: IconName = IconName(\n \"square-rounded-chevrons-up-filled\"\n )\n SQUARE_ROUNDED_FILLED: IconName = IconName(\"square-rounded-filled\")\n SQUARE_ROUNDED_LETTER_A: IconName = IconName(\"square-rounded-letter-a\")\n SQUARE_ROUNDED_LETTER_B: IconName = IconName(\"square-rounded-letter-b\")\n SQUARE_ROUNDED_LETTER_C: IconName = IconName(\"square-rounded-letter-c\")\n SQUARE_ROUNDED_LETTER_D: IconName = IconName(\"square-rounded-letter-d\")\n SQUARE_ROUNDED_LETTER_E: IconName = IconName(\"square-rounded-letter-e\")\n SQUARE_ROUNDED_LETTER_F: IconName = IconName(\"square-rounded-letter-f\")\n SQUARE_ROUNDED_LETTER_G: IconName = IconName(\"square-rounded-letter-g\")\n SQUARE_ROUNDED_LETTER_H: IconName = IconName(\"square-rounded-letter-h\")\n SQUARE_ROUNDED_LETTER_I: IconName = IconName(\"square-rounded-letter-i\")\n SQUARE_ROUNDED_LETTER_J: IconName = IconName(\"square-rounded-letter-j\")\n SQUARE_ROUNDED_LETTER_K: IconName = IconName(\"square-rounded-letter-k\")\n SQUARE_ROUNDED_LETTER_L: IconName = IconName(\"square-rounded-letter-l\")\n SQUARE_ROUNDED_LETTER_M: IconName = IconName(\"square-rounded-letter-m\")\n SQUARE_ROUNDED_LETTER_N: IconName = IconName(\"square-rounded-letter-n\")\n SQUARE_ROUNDED_LETTER_O: IconName = IconName(\"square-rounded-letter-o\")\n SQUARE_ROUNDED_LETTER_P: IconName = IconName(\"square-rounded-letter-p\")\n SQUARE_ROUNDED_LETTER_Q: IconName = IconName(\"square-rounded-letter-q\")\n SQUARE_ROUNDED_LETTER_R: IconName = IconName(\"square-rounded-letter-r\")\n SQUARE_ROUNDED_LETTER_S: IconName = IconName(\"square-rounded-letter-s\")\n SQUARE_ROUNDED_LETTER_T: IconName = IconName(\"square-rounded-letter-t\")\n SQUARE_ROUNDED_LETTER_U: IconName = IconName(\"square-rounded-letter-u\")\n SQUARE_ROUNDED_LETTER_V: IconName = IconName(\"square-rounded-letter-v\")\n SQUARE_ROUNDED_LETTER_W: IconName = IconName(\"square-rounded-letter-w\")\n SQUARE_ROUNDED_LETTER_X: IconName = IconName(\"square-rounded-letter-x\")\n SQUARE_ROUNDED_LETTER_Y: IconName = IconName(\"square-rounded-letter-y\")\n SQUARE_ROUNDED_LETTER_Z: IconName = IconName(\"square-rounded-letter-z\")\n SQUARE_ROUNDED_MINUS: IconName = IconName(\"square-rounded-minus\")\n SQUARE_ROUNDED_NUMBER_0: IconName = IconName(\"square-rounded-number-0\")\n SQUARE_ROUNDED_NUMBER_0_FILLED: IconName = IconName(\n \"square-rounded-number-0-filled\"\n )\n SQUARE_ROUNDED_NUMBER_1: IconName = IconName(\"square-rounded-number-1\")\n SQUARE_ROUNDED_NUMBER_1_FILLED: IconName = IconName(\n \"square-rounded-number-1-filled\"\n )\n SQUARE_ROUNDED_NUMBER_2: IconName = IconName(\"square-rounded-number-2\")\n SQUARE_ROUNDED_NUMBER_2_FILLED: IconName = IconName(\n \"square-rounded-number-2-filled\"\n )\n SQUARE_ROUNDED_NUMBER_3: IconName = IconName(\"square-rounded-number-3\")\n SQUARE_ROUNDED_NUMBER_3_FILLED: IconName = IconName(\n \"square-rounded-number-3-filled\"\n )\n SQUARE_ROUNDED_NUMBER_4: IconName = IconName(\"square-rounded-number-4\")\n SQUARE_ROUNDED_NUMBER_4_FILLED: IconName = IconName(\n \"square-rounded-number-4-filled\"\n )\n SQUARE_ROUNDED_NUMBER_5: IconName = IconName(\"square-rounded-number-5\")\n SQUARE_ROUNDED_NUMBER_5_FILLED: IconName = IconName(\n \"square-rounded-number-5-filled\"\n )\n SQUARE_ROUNDED_NUMBER_6: IconName = IconName(\"square-rounded-number-6\")\n SQUARE_ROUNDED_NUMBER_6_FILLED: IconName = IconName(\n \"square-rounded-number-6-filled\"\n )\n SQUARE_ROUNDED_NUMBER_7: IconName = IconName(\"square-rounded-number-7\")\n SQUARE_ROUNDED_NUMBER_7_FILLED: IconName = IconName(\n \"square-rounded-number-7-filled\"\n )\n SQUARE_ROUNDED_NUMBER_8: IconName = IconName(\"square-rounded-number-8\")\n SQUARE_ROUNDED_NUMBER_8_FILLED: IconName = IconName(\n \"square-rounded-number-8-filled\"\n )\n SQUARE_ROUNDED_NUMBER_9: IconName = IconName(\"square-rounded-number-9\")\n SQUARE_ROUNDED_NUMBER_9_FILLED: IconName = IconName(\n \"square-rounded-number-9-filled\"\n )\n SQUARE_ROUNDED_PLUS: IconName = IconName(\"square-rounded-plus\")\n SQUARE_ROUNDED_PLUS_FILLED: IconName = IconName(\"square-rounded-plus-filled\")\n SQUARE_ROUNDED_X: IconName = IconName(\"square-rounded-x\")\n SQUARE_ROUNDED_X_FILLED: IconName = IconName(\"square-rounded-x-filled\")\n SQUARE_TOGGLE: IconName = IconName(\"square-toggle\")\n SQUARE_TOGGLE_HORIZONTAL: IconName = IconName(\"square-toggle-horizontal\")\n SQUARE_X: IconName = IconName(\"square-x\")\n SQUARES_DIAGONAL: IconName = IconName(\"squares-diagonal\")\n SQUARES_FILLED: IconName = IconName(\"squares-filled\")\n STACK: IconName = IconName(\"stack\")\n STACK_2: IconName = IconName(\"stack-2\")\n STACK_3: IconName = IconName(\"stack-3\")\n STACK_POP: IconName = IconName(\"stack-pop\")\n STACK_PUSH: IconName = IconName(\"stack-push\")\n STAIRS: IconName = IconName(\"stairs\")\n STAIRS_DOWN: IconName = IconName(\"stairs-down\")\n STAIRS_UP: IconName = IconName(\"stairs-up\")\n STAR: IconName = IconName(\"star\")\n STAR_FILLED: IconName = IconName(\"star-filled\")\n STAR_HALF: IconName = IconName(\"star-half\")\n STAR_HALF_FILLED: IconName = IconName(\"star-half-filled\")\n STAR_OFF: IconName = IconName(\"star-off\")\n STARS: IconName = IconName(\"stars\")\n STARS_FILLED: IconName = IconName(\"stars-filled\")\n STARS_OFF: IconName = IconName(\"stars-off\")\n STATUS_CHANGE: IconName = IconName(\"status-change\")\n STEAM: IconName = IconName(\"steam\")\n STEERING_WHEEL: IconName = IconName(\"steering-wheel\")\n STEERING_WHEEL_OFF: IconName = IconName(\"steering-wheel-off\")\n STEP_INTO: IconName = IconName(\"step-into\")\n STEP_OUT: IconName = IconName(\"step-out\")\n STEREO_GLASSES: IconName = IconName(\"stereo-glasses\")\n STETHOSCOPE: IconName = IconName(\"stethoscope\")\n STETHOSCOPE_OFF: IconName = IconName(\"stethoscope-off\")\n STICKER: IconName = IconName(\"sticker\")\n STORM: IconName = IconName(\"storm\")\n STORM_OFF: IconName = IconName(\"storm-off\")\n STRETCHING: IconName = IconName(\"stretching\")\n STRETCHING_2: IconName = IconName(\"stretching-2\")\n STRIKETHROUGH: IconName = IconName(\"strikethrough\")\n SUBMARINE: IconName = IconName(\"submarine\")\n SUBSCRIPT: IconName = IconName(\"subscript\")\n SUBTASK: IconName = IconName(\"subtask\")\n SUM: IconName = IconName(\"sum\")\n SUM_OFF: IconName = IconName(\"sum-off\")\n SUN: IconName = IconName(\"sun\")\n SUN_FILLED: IconName = IconName(\"sun-filled\")\n SUN_HIGH: IconName = IconName(\"sun-high\")\n SUN_LOW: IconName = IconName(\"sun-low\")\n SUN_MOON: IconName = IconName(\"sun-moon\")\n SUN_OFF: IconName = IconName(\"sun-off\")\n SUN_WIND: IconName = IconName(\"sun-wind\")\n SUNGLASSES: IconName = IconName(\"sunglasses\")\n SUNRISE: IconName = IconName(\"sunrise\")\n SUNSET: IconName = IconName(\"sunset\")\n SUNSET_2: IconName = IconName(\"sunset-2\")\n SUPERSCRIPT: IconName = IconName(\"superscript\")\n SVG: IconName = IconName(\"svg\")\n SWIMMING: IconName = IconName(\"swimming\")\n SWIPE: IconName = IconName(\"swipe\")\n SWITCH: IconName = IconName(\"switch\")\n SWITCH_2: IconName = IconName(\"switch-2\")\n SWITCH_3: IconName = IconName(\"switch-3\")\n SWITCH_HORIZONTAL: IconName = IconName(\"switch-horizontal\")\n SWITCH_VERTICAL: IconName = IconName(\"switch-vertical\")\n SWORD: IconName = IconName(\"sword\")\n SWORD_OFF: IconName = IconName(\"sword-off\")\n SWORDS: IconName = IconName(\"swords\")\n TABLE: IconName = IconName(\"table\")\n TABLE_ALIAS: IconName = IconName(\"table-alias\")\n TABLE_COLUMN: IconName = IconName(\"table-column\")\n TABLE_DOWN: IconName = IconName(\"table-down\")\n TABLE_EXPORT: IconName = IconName(\"table-export\")\n TABLE_FILLED: IconName = IconName(\"table-filled\")\n TABLE_HEART: IconName = IconName(\"table-heart\")\n TABLE_IMPORT: IconName = IconName(\"table-import\")\n TABLE_MINUS: IconName = IconName(\"table-minus\")\n TABLE_OFF: IconName = IconName(\"table-off\")\n TABLE_OPTIONS: IconName = IconName(\"table-options\")\n TABLE_PLUS: IconName = IconName(\"table-plus\")\n TABLE_ROW: IconName = IconName(\"table-row\")\n TABLE_SHARE: IconName = IconName(\"table-share\")\n TABLE_SHORTCUT: IconName = IconName(\"table-shortcut\")\n TAG: IconName = IconName(\"tag\")\n TAG_OFF: IconName = IconName(\"tag-off\")\n TAGS: IconName = IconName(\"tags\")\n TAGS_OFF: IconName = IconName(\"tags-off\")\n TALLYMARK_1: IconName = IconName(\"tallymark-1\")\n TALLYMARK_2: IconName = IconName(\"tallymark-2\")\n TALLYMARK_3: IconName = IconName(\"tallymark-3\")\n TALLYMARK_4: IconName = IconName(\"tallymark-4\")\n TALLYMARKS: IconName = IconName(\"tallymarks\")\n TANK: IconName = IconName(\"tank\")\n TARGET: IconName = IconName(\"target\")\n TARGET_ARROW: IconName = IconName(\"target-arrow\")\n TARGET_OFF: IconName = IconName(\"target-off\")\n TEAPOT: IconName = IconName(\"teapot\")\n TELESCOPE: IconName = IconName(\"telescope\")\n TELESCOPE_OFF: IconName = IconName(\"telescope-off\")\n TEMPERATURE: IconName = IconName(\"temperature\")\n TEMPERATURE_CELSIUS: IconName = IconName(\"temperature-celsius\")\n TEMPERATURE_FAHRENHEIT: IconName = IconName(\"temperature-fahrenheit\")\n TEMPERATURE_MINUS: IconName = IconName(\"temperature-minus\")\n TEMPERATURE_OFF: IconName = IconName(\"temperature-off\")\n TEMPERATURE_PLUS: IconName = IconName(\"temperature-plus\")\n TEMPLATE: IconName = IconName(\"template\")\n TEMPLATE_OFF: IconName = IconName(\"template-off\")\n TENT: IconName = IconName(\"tent\")\n TENT_OFF: IconName = IconName(\"tent-off\")\n TERMINAL: IconName = IconName(\"terminal\")\n TERMINAL_2: IconName = IconName(\"terminal-2\")\n TEST_PIPE: IconName = IconName(\"test-pipe\")\n TEST_PIPE_2: IconName = IconName(\"test-pipe-2\")\n TEST_PIPE_OFF: IconName = IconName(\"test-pipe-off\")\n TEX: IconName = IconName(\"tex\")\n TEXT_CAPTION: IconName = IconName(\"text-caption\")\n TEXT_COLOR: IconName = IconName(\"text-color\")\n TEXT_DECREASE: IconName = IconName(\"text-decrease\")\n TEXT_DIRECTION_LTR: IconName = IconName(\"text-direction-ltr\")\n TEXT_DIRECTION_RTL: IconName = IconName(\"text-direction-rtl\")\n TEXT_INCREASE: IconName = IconName(\"text-increase\")\n TEXT_ORIENTATION: IconName = IconName(\"text-orientation\")\n TEXT_PLUS: IconName = IconName(\"text-plus\")\n TEXT_RECOGNITION: IconName = IconName(\"text-recognition\")\n TEXT_RESIZE: IconName = IconName(\"text-resize\")\n TEXT_SIZE: IconName = IconName(\"text-size\")\n TEXT_SPELLCHECK: IconName = IconName(\"text-spellcheck\")\n TEXT_WRAP: IconName = IconName(\"text-wrap\")\n TEXT_WRAP_DISABLED: IconName = IconName(\"text-wrap-disabled\")\n TEXTURE: IconName = IconName(\"texture\")\n THEATER: IconName = IconName(\"theater\")\n THERMOMETER: IconName = IconName(\"thermometer\")\n THUMB_DOWN: IconName = IconName(\"thumb-down\")\n THUMB_DOWN_FILLED: IconName = IconName(\"thumb-down-filled\")\n THUMB_DOWN_OFF: IconName = IconName(\"thumb-down-off\")\n THUMB_UP: IconName = IconName(\"thumb-up\")\n THUMB_UP_FILLED: IconName = IconName(\"thumb-up-filled\")\n THUMB_UP_OFF: IconName = IconName(\"thumb-up-off\")\n TIC_TAC: IconName = IconName(\"tic-tac\")\n TICKET: IconName = IconName(\"ticket\")\n TICKET_OFF: IconName = IconName(\"ticket-off\")\n TIE: IconName = IconName(\"tie\")\n TILDE: IconName = IconName(\"tilde\")\n TILT_SHIFT: IconName = IconName(\"tilt-shift\")\n TILT_SHIFT_OFF: IconName = IconName(\"tilt-shift-off\")\n TIME_DURATION_0: IconName = IconName(\"time-duration-0\")\n TIME_DURATION_10: IconName = IconName(\"time-duration-10\")\n TIME_DURATION_15: IconName = IconName(\"time-duration-15\")\n TIME_DURATION_30: IconName = IconName(\"time-duration-30\")\n TIME_DURATION_45: IconName = IconName(\"time-duration-45\")\n TIME_DURATION_5: IconName = IconName(\"time-duration-5\")\n TIME_DURATION_60: IconName = IconName(\"time-duration-60\")\n TIME_DURATION_90: IconName = IconName(\"time-duration-90\")\n TIME_DURATION_OFF: IconName = IconName(\"time-duration-off\")\n TIMELINE: IconName = IconName(\"timeline\")\n TIMELINE_EVENT: IconName = IconName(\"timeline-event\")\n TIMELINE_EVENT_EXCLAMATION: IconName = IconName(\"timeline-event-exclamation\")\n TIMELINE_EVENT_MINUS: IconName = IconName(\"timeline-event-minus\")\n TIMELINE_EVENT_PLUS: IconName = IconName(\"timeline-event-plus\")\n TIMELINE_EVENT_TEXT: IconName = IconName(\"timeline-event-text\")\n TIMELINE_EVENT_X: IconName = IconName(\"timeline-event-x\")\n TIR: IconName = IconName(\"tir\")\n TOGGLE_LEFT: IconName = IconName(\"toggle-left\")\n TOGGLE_RIGHT: IconName = IconName(\"toggle-right\")\n TOILET_PAPER: IconName = IconName(\"toilet-paper\")\n TOILET_PAPER_OFF: IconName = IconName(\"toilet-paper-off\")\n TOML: IconName = IconName(\"toml\")\n TOOL: IconName = IconName(\"tool\")\n TOOLS: IconName = IconName(\"tools\")\n TOOLS_KITCHEN: IconName = IconName(\"tools-kitchen\")\n TOOLS_KITCHEN_2: IconName = IconName(\"tools-kitchen-2\")\n TOOLS_KITCHEN_2_OFF: IconName = IconName(\"tools-kitchen-2-off\")\n TOOLS_KITCHEN_OFF: IconName = IconName(\"tools-kitchen-off\")\n TOOLS_OFF: IconName = IconName(\"tools-off\")\n TOOLTIP: IconName = IconName(\"tooltip\")\n TOPOLOGY_BUS: IconName = IconName(\"topology-bus\")\n TOPOLOGY_COMPLEX: IconName = IconName(\"topology-complex\")\n TOPOLOGY_FULL: IconName = IconName(\"topology-full\")\n TOPOLOGY_FULL_HIERARCHY: IconName = IconName(\"topology-full-hierarchy\")\n TOPOLOGY_RING: IconName = IconName(\"topology-ring\")\n TOPOLOGY_RING_2: IconName = IconName(\"topology-ring-2\")\n TOPOLOGY_RING_3: IconName = IconName(\"topology-ring-3\")\n TOPOLOGY_STAR: IconName = IconName(\"topology-star\")\n TOPOLOGY_STAR_2: IconName = IconName(\"topology-star-2\")\n TOPOLOGY_STAR_3: IconName = IconName(\"topology-star-3\")\n TOPOLOGY_STAR_RING: IconName = IconName(\"topology-star-ring\")\n TOPOLOGY_STAR_RING_2: IconName = IconName(\"topology-star-ring-2\")\n TOPOLOGY_STAR_RING_3: IconName = IconName(\"topology-star-ring-3\")\n TORII: IconName = IconName(\"torii\")\n TORNADO: IconName = IconName(\"tornado\")\n TOURNAMENT: IconName = IconName(\"tournament\")\n TOWER: IconName = IconName(\"tower\")\n TOWER_OFF: IconName = IconName(\"tower-off\")\n TRACK: IconName = IconName(\"track\")\n TRACTOR: IconName = IconName(\"tractor\")\n TRADEMARK: IconName = IconName(\"trademark\")\n TRAFFIC_CONE: IconName = IconName(\"traffic-cone\")\n TRAFFIC_CONE_OFF: IconName = IconName(\"traffic-cone-off\")\n TRAFFIC_LIGHTS: IconName = IconName(\"traffic-lights\")\n TRAFFIC_LIGHTS_OFF: IconName = IconName(\"traffic-lights-off\")\n TRAIN: IconName = IconName(\"train\")\n TRANSFER_IN: IconName = IconName(\"transfer-in\")\n TRANSFER_OUT: IconName = IconName(\"transfer-out\")\n TRANSFORM: IconName = IconName(\"transform\")\n TRANSFORM_FILLED: IconName = IconName(\"transform-filled\")\n TRANSITION_BOTTOM: IconName = IconName(\"transition-bottom\")\n TRANSITION_LEFT: IconName = IconName(\"transition-left\")\n TRANSITION_RIGHT: IconName = IconName(\"transition-right\")\n TRANSITION_TOP: IconName = IconName(\"transition-top\")\n TRASH: IconName = IconName(\"trash\")\n TRASH_FILLED: IconName = IconName(\"trash-filled\")\n TRASH_OFF: IconName = IconName(\"trash-off\")\n TRASH_X: IconName = IconName(\"trash-x\")\n TRASH_X_FILLED: IconName = IconName(\"trash-x-filled\")\n TREADMILL: IconName = IconName(\"treadmill\")\n TREE: IconName = IconName(\"tree\")\n TREES: IconName = IconName(\"trees\")\n TREKKING: IconName = IconName(\"trekking\")\n TRENDING_DOWN: IconName = IconName(\"trending-down\")\n TRENDING_DOWN_2: IconName = IconName(\"trending-down-2\")\n TRENDING_DOWN_3: IconName = IconName(\"trending-down-3\")\n TRENDING_UP: IconName = IconName(\"trending-up\")\n TRENDING_UP_2: IconName = IconName(\"trending-up-2\")\n TRENDING_UP_3: IconName = IconName(\"trending-up-3\")\n TRIANGLE: IconName = IconName(\"triangle\")\n TRIANGLE_FILLED: IconName = IconName(\"triangle-filled\")\n TRIANGLE_INVERTED: IconName = IconName(\"triangle-inverted\")\n TRIANGLE_INVERTED_FILLED: IconName = IconName(\"triangle-inverted-filled\")\n TRIANGLE_OFF: IconName = IconName(\"triangle-off\")\n TRIANGLE_SQUARE_CIRCLE: IconName = IconName(\"triangle-square-circle\")\n TRIANGLES: IconName = IconName(\"triangles\")\n TRIDENT: IconName = IconName(\"trident\")\n TROLLEY: IconName = IconName(\"trolley\")\n TROPHY: IconName = IconName(\"trophy\")\n TROPHY_FILLED: IconName = IconName(\"trophy-filled\")\n TROPHY_OFF: IconName = IconName(\"trophy-off\")\n TROWEL: IconName = IconName(\"trowel\")\n TRUCK: IconName = IconName(\"truck\")\n TRUCK_DELIVERY: IconName = IconName(\"truck-delivery\")\n TRUCK_LOADING: IconName = IconName(\"truck-loading\")\n TRUCK_OFF: IconName = IconName(\"truck-off\")\n TRUCK_RETURN: IconName = IconName(\"truck-return\")\n TXT: IconName = IconName(\"txt\")\n TYPOGRAPHY: IconName = IconName(\"typography\")\n TYPOGRAPHY_OFF: IconName = IconName(\"typography-off\")\n UFO: IconName = IconName(\"ufo\")\n UFO_OFF: IconName = IconName(\"ufo-off\")\n UMBRELLA: IconName = IconName(\"umbrella\")\n UMBRELLA_FILLED: IconName = IconName(\"umbrella-filled\")\n UMBRELLA_OFF: IconName = IconName(\"umbrella-off\")\n UNDERLINE: IconName = IconName(\"underline\")\n UNLINK: IconName = IconName(\"unlink\")\n UPLOAD: IconName = IconName(\"upload\")\n URGENT: IconName = IconName(\"urgent\")\n USB: IconName = IconName(\"usb\")\n USER: IconName = IconName(\"user\")\n USER_BOLT: IconName = IconName(\"user-bolt\")\n USER_CANCEL: IconName = IconName(\"user-cancel\")\n USER_CHECK: IconName = IconName(\"user-check\")\n USER_CIRCLE: IconName = IconName(\"user-circle\")\n USER_CODE: IconName = IconName(\"user-code\")\n USER_COG: IconName = IconName(\"user-cog\")\n USER_DOLLAR: IconName = IconName(\"user-dollar\")\n USER_DOWN: IconName = IconName(\"user-down\")\n USER_EDIT: IconName = IconName(\"user-edit\")\n USER_EXCLAMATION: IconName = IconName(\"user-exclamation\")\n USER_HEART: IconName = IconName(\"user-heart\")\n USER_MINUS: IconName = IconName(\"user-minus\")\n USER_OFF: IconName = IconName(\"user-off\")\n USER_PAUSE: IconName = IconName(\"user-pause\")\n USER_PIN: IconName = IconName(\"user-pin\")\n USER_PLUS: IconName = IconName(\"user-plus\")\n USER_QUESTION: IconName = IconName(\"user-question\")\n USER_SEARCH: IconName = IconName(\"user-search\")\n USER_SHARE: IconName = IconName(\"user-share\")\n USER_SHIELD: IconName = IconName(\"user-shield\")\n USER_STAR: IconName = IconName(\"user-star\")\n USER_UP: IconName = IconName(\"user-up\")\n USER_X: IconName = IconName(\"user-x\")\n USERS: IconName = IconName(\"users\")\n USERS_GROUP: IconName = IconName(\"users-group\")\n USERS_MINUS: IconName = IconName(\"users-minus\")\n USERS_PLUS: IconName = IconName(\"users-plus\")\n UV_INDEX: IconName = IconName(\"uv-index\")\n UX_CIRCLE: IconName = IconName(\"ux-circle\")\n VACCINE: IconName = IconName(\"vaccine\")\n VACCINE_BOTTLE: IconName = IconName(\"vaccine-bottle\")\n VACCINE_BOTTLE_OFF: IconName = IconName(\"vaccine-bottle-off\")\n VACCINE_OFF: IconName = IconName(\"vaccine-off\")\n VACUUM_CLEANER: IconName = IconName(\"vacuum-cleaner\")\n VARIABLE: IconName = IconName(\"variable\")\n VARIABLE_MINUS: IconName = IconName(\"variable-minus\")\n VARIABLE_OFF: IconName = IconName(\"variable-off\")\n VARIABLE_PLUS: IconName = IconName(\"variable-plus\")\n VECTOR: IconName = IconName(\"vector\")\n VECTOR_BEZIER: IconName = IconName(\"vector-bezier\")\n VECTOR_BEZIER_2: IconName = IconName(\"vector-bezier-2\")\n VECTOR_BEZIER_ARC: IconName = IconName(\"vector-bezier-arc\")\n VECTOR_BEZIER_CIRCLE: IconName = IconName(\"vector-bezier-circle\")\n VECTOR_OFF: IconName = IconName(\"vector-off\")\n VECTOR_SPLINE: IconName = IconName(\"vector-spline\")\n VECTOR_TRIANGLE: IconName = IconName(\"vector-triangle\")\n VECTOR_TRIANGLE_OFF: IconName = IconName(\"vector-triangle-off\")\n VENUS: IconName = IconName(\"venus\")\n VERSIONS: IconName = IconName(\"versions\")\n VERSIONS_FILLED: IconName = IconName(\"versions-filled\")\n VERSIONS_OFF: IconName = IconName(\"versions-off\")\n VIDEO: IconName = IconName(\"video\")\n VIDEO_MINUS: IconName = IconName(\"video-minus\")\n VIDEO_OFF: IconName = IconName(\"video-off\")\n VIDEO_PLUS: IconName = IconName(\"video-plus\")\n VIEW_360: IconName = IconName(\"view-360\")\n VIEW_360_OFF: IconName = IconName(\"view-360-off\")\n VIEWFINDER: IconName = IconName(\"viewfinder\")\n VIEWFINDER_OFF: IconName = IconName(\"viewfinder-off\")\n VIEWPORT_NARROW: IconName = IconName(\"viewport-narrow\")\n VIEWPORT_WIDE: IconName = IconName(\"viewport-wide\")\n VINYL: IconName = IconName(\"vinyl\")\n VIP: IconName = IconName(\"vip\")\n VIP_OFF: IconName = IconName(\"vip-off\")\n VIRUS: IconName = IconName(\"virus\")\n VIRUS_OFF: IconName = IconName(\"virus-off\")\n VIRUS_SEARCH: IconName = IconName(\"virus-search\")\n VOCABULARY: IconName = IconName(\"vocabulary\")\n VOCABULARY_OFF: IconName = IconName(\"vocabulary-off\")\n VOLCANO: IconName = IconName(\"volcano\")\n VOLUME: IconName = IconName(\"volume\")\n VOLUME_2: IconName = IconName(\"volume-2\")\n VOLUME_3: IconName = IconName(\"volume-3\")\n VOLUME_OFF: IconName = IconName(\"volume-off\")\n WALK: IconName = IconName(\"walk\")\n WALL: IconName = IconName(\"wall\")\n WALL_OFF: IconName = IconName(\"wall-off\")\n WALLET: IconName = IconName(\"wallet\")\n WALLET_OFF: IconName = IconName(\"wallet-off\")\n WALLPAPER: IconName = IconName(\"wallpaper\")\n WALLPAPER_OFF: IconName = IconName(\"wallpaper-off\")\n WAND: IconName = IconName(\"wand\")\n WAND_OFF: IconName = IconName(\"wand-off\")\n WASH: IconName = IconName(\"wash\")\n WASH_DRY: IconName = IconName(\"wash-dry\")\n WASH_DRY_1: IconName = IconName(\"wash-dry-1\")\n WASH_DRY_2: IconName = IconName(\"wash-dry-2\")\n WASH_DRY_3: IconName = IconName(\"wash-dry-3\")\n WASH_DRY_A: IconName = IconName(\"wash-dry-a\")\n WASH_DRY_DIP: IconName = IconName(\"wash-dry-dip\")\n WASH_DRY_F: IconName = IconName(\"wash-dry-f\")\n WASH_DRY_FLAT: IconName = IconName(\"wash-dry-flat\")\n WASH_DRY_HANG: IconName = IconName(\"wash-dry-hang\")\n WASH_DRY_OFF: IconName = IconName(\"wash-dry-off\")\n WASH_DRY_P: IconName = IconName(\"wash-dry-p\")\n WASH_DRY_SHADE: IconName = IconName(\"wash-dry-shade\")\n WASH_DRY_W: IconName = IconName(\"wash-dry-w\")\n WASH_DRYCLEAN: IconName = IconName(\"wash-dryclean\")\n WASH_DRYCLEAN_OFF: IconName = IconName(\"wash-dryclean-off\")\n WASH_ECO: IconName = IconName(\"wash-eco\")\n WASH_GENTLE: IconName = IconName(\"wash-gentle\")\n WASH_HAND: IconName = IconName(\"wash-hand\")\n WASH_MACHINE: IconName = IconName(\"wash-machine\")\n WASH_OFF: IconName = IconName(\"wash-off\")\n WASH_PRESS: IconName = IconName(\"wash-press\")\n WASH_TEMPERATURE_1: IconName = IconName(\"wash-temperature-1\")\n WASH_TEMPERATURE_2: IconName = IconName(\"wash-temperature-2\")\n WASH_TEMPERATURE_3: IconName = IconName(\"wash-temperature-3\")\n WASH_TEMPERATURE_4: IconName = IconName(\"wash-temperature-4\")\n WASH_TEMPERATURE_5: IconName = IconName(\"wash-temperature-5\")\n WASH_TEMPERATURE_6: IconName = IconName(\"wash-temperature-6\")\n WASH_TUMBLE_DRY: IconName = IconName(\"wash-tumble-dry\")\n WASH_TUMBLE_OFF: IconName = IconName(\"wash-tumble-off\")\n WATERPOLO: IconName = IconName(\"waterpolo\")\n WAVE_SAW_TOOL: IconName = IconName(\"wave-saw-tool\")\n WAVE_SINE: IconName = IconName(\"wave-sine\")\n WAVE_SQUARE: IconName = IconName(\"wave-square\")\n WEBHOOK: IconName = IconName(\"webhook\")\n WEBHOOK_OFF: IconName = IconName(\"webhook-off\")\n WEIGHT: IconName = IconName(\"weight\")\n WHEELCHAIR: IconName = IconName(\"wheelchair\")\n WHEELCHAIR_OFF: IconName = IconName(\"wheelchair-off\")\n WHIRL: IconName = IconName(\"whirl\")\n WIFI: IconName = IconName(\"wifi\")\n WIFI_0: IconName = IconName(\"wifi-0\")\n WIFI_1: IconName = IconName(\"wifi-1\")\n WIFI_2: IconName = IconName(\"wifi-2\")\n WIFI_OFF: IconName = IconName(\"wifi-off\")\n WIND: IconName = IconName(\"wind\")\n WIND_OFF: IconName = IconName(\"wind-off\")\n WINDMILL: IconName = IconName(\"windmill\")\n WINDMILL_FILLED: IconName = IconName(\"windmill-filled\")\n WINDMILL_OFF: IconName = IconName(\"windmill-off\")\n WINDOW: IconName = IconName(\"window\")\n WINDOW_MAXIMIZE: IconName = IconName(\"window-maximize\")\n WINDOW_MINIMIZE: IconName = IconName(\"window-minimize\")\n WINDOW_OFF: IconName = IconName(\"window-off\")\n WINDSOCK: IconName = IconName(\"windsock\")\n WIPER: IconName = IconName(\"wiper\")\n WIPER_WASH: IconName = IconName(\"wiper-wash\")\n WOMAN: IconName = IconName(\"woman\")\n WOOD: IconName = IconName(\"wood\")\n WORLD: IconName = IconName(\"world\")\n WORLD_BOLT: IconName = IconName(\"world-bolt\")\n WORLD_CANCEL: IconName = IconName(\"world-cancel\")\n WORLD_CHECK: IconName = IconName(\"world-check\")\n WORLD_CODE: IconName = IconName(\"world-code\")\n WORLD_COG: IconName = IconName(\"world-cog\")\n WORLD_DOLLAR: IconName = IconName(\"world-dollar\")\n WORLD_DOWN: IconName = IconName(\"world-down\")\n WORLD_DOWNLOAD: IconName = IconName(\"world-download\")\n WORLD_EXCLAMATION: IconName = IconName(\"world-exclamation\")\n WORLD_HEART: IconName = IconName(\"world-heart\")\n WORLD_LATITUDE: IconName = IconName(\"world-latitude\")\n WORLD_LONGITUDE: IconName = IconName(\"world-longitude\")\n WORLD_MINUS: IconName = IconName(\"world-minus\")\n WORLD_OFF: IconName = IconName(\"world-off\")\n WORLD_PAUSE: IconName = IconName(\"world-pause\")\n WORLD_PIN: IconName = IconName(\"world-pin\")\n WORLD_PLUS: IconName = IconName(\"world-plus\")\n WORLD_QUESTION: IconName = IconName(\"world-question\")\n WORLD_SEARCH: IconName = IconName(\"world-search\")\n WORLD_SHARE: IconName = IconName(\"world-share\")\n WORLD_STAR: IconName = IconName(\"world-star\")\n WORLD_UP: IconName = IconName(\"world-up\")\n WORLD_UPLOAD: IconName = IconName(\"world-upload\")\n WORLD_WWW: IconName = IconName(\"world-www\")\n WORLD_X: IconName = IconName(\"world-x\")\n WRECKING_BALL: IconName = IconName(\"wrecking-ball\")\n WRITING: IconName = IconName(\"writing\")\n WRITING_OFF: IconName = IconName(\"writing-off\")\n WRITING_SIGN: IconName = IconName(\"writing-sign\")\n WRITING_SIGN_OFF: IconName = IconName(\"writing-sign-off\")\n X: IconName = IconName(\"x\")\n XBOX_A: IconName = IconName(\"xbox-a\")\n XBOX_B: IconName = IconName(\"xbox-b\")\n XBOX_X: IconName = IconName(\"xbox-x\")\n XBOX_Y: IconName = IconName(\"xbox-y\")\n XD: IconName = IconName(\"xd\")\n YIN_YANG: IconName = IconName(\"yin-yang\")\n YIN_YANG_FILLED: IconName = IconName(\"yin-yang-filled\")\n YOGA: IconName = IconName(\"yoga\")\n ZEPPELIN: IconName = IconName(\"zeppelin\")\n ZEPPELIN_OFF: IconName = IconName(\"zeppelin-off\")\n ZIP: IconName = IconName(\"zip\")\n ZODIAC_AQUARIUS: IconName = IconName(\"zodiac-aquarius\")\n ZODIAC_ARIES: IconName = IconName(\"zodiac-aries\")\n ZODIAC_CANCER: IconName = IconName(\"zodiac-cancer\")\n ZODIAC_CAPRICORN: IconName = IconName(\"zodiac-capricorn\")\n ZODIAC_GEMINI: IconName = IconName(\"zodiac-gemini\")\n ZODIAC_LEO: IconName = IconName(\"zodiac-leo\")\n ZODIAC_LIBRA: IconName = IconName(\"zodiac-libra\")\n ZODIAC_PISCES: IconName = IconName(\"zodiac-pisces\")\n ZODIAC_SAGITTARIUS: IconName = IconName(\"zodiac-sagittarius\")\n ZODIAC_SCORPIO: IconName = IconName(\"zodiac-scorpio\")\n ZODIAC_TAURUS: IconName = IconName(\"zodiac-taurus\")\n ZODIAC_VIRGO: IconName = IconName(\"zodiac-virgo\")\n ZOOM_CANCEL: IconName = IconName(\"zoom-cancel\")\n ZOOM_CHECK: IconName = IconName(\"zoom-check\")\n ZOOM_CHECK_FILLED: IconName = IconName(\"zoom-check-filled\")\n ZOOM_CODE: IconName = IconName(\"zoom-code\")\n ZOOM_EXCLAMATION: IconName = IconName(\"zoom-exclamation\")\n ZOOM_FILLED: IconName = IconName(\"zoom-filled\")\n ZOOM_IN: IconName = IconName(\"zoom-in\")\n ZOOM_IN_AREA: IconName = IconName(\"zoom-in-area\")\n ZOOM_IN_AREA_FILLED: IconName = IconName(\"zoom-in-area-filled\")\n ZOOM_IN_FILLED: IconName = IconName(\"zoom-in-filled\")\n ZOOM_MONEY: IconName = IconName(\"zoom-money\")\n ZOOM_OUT: IconName = IconName(\"zoom-out\")\n ZOOM_OUT_AREA: IconName = IconName(\"zoom-out-area\")\n ZOOM_OUT_FILLED: IconName = IconName(\"zoom-out-filled\")\n ZOOM_PAN: IconName = IconName(\"zoom-pan\")\n ZOOM_QUESTION: IconName = IconName(\"zoom-question\")\n ZOOM_REPLACE: IconName = IconName(\"zoom-replace\")\n ZOOM_RESET: IconName = IconName(\"zoom-reset\")\n ZZZ: IconName = IconName(\"zzz\")\n ZZZ_OFF: IconName = IconName(\"zzz-off\")\n" }, { "path": "viser/src/viser/_icons_generate_enum.py", "content": "\"\"\"Helper script for dumping Tabler icon names into a Literal type.\"\"\"\n\nimport tarfile\nfrom pathlib import Path\n\nHERE_DIR = Path(__file__).absolute().parent\nICON_DIR = HERE_DIR / \"_icons\"\n\n\ndef enum_name_from_icon(name: str) -> str:\n \"\"\"Capitalize an icon name for use as an enum name.\"\"\"\n name = name.upper()\n name = name.replace(\"-\", \"_\")\n if name[0].isdigit():\n name = \"ICON_\" + name\n return name\n\n\nif __name__ == \"__main__\":\n with tarfile.open(ICON_DIR / \"tabler-icons.tar\") as tar:\n icon_names = sorted([name.partition(\".svg\")[0] for name in tar.getnames()])\n\n # Generate stub file. This is used by type checkers.\n (HERE_DIR / \"_icons_enum.pyi\").write_text(\n \"\\n\".join(\n [\n \"# Automatically generated by `_icons_generate_enum.py`\",\n \"# See https://tabler-icons.io/\",\n \"import enum\",\n \"from typing import NewType\",\n \"\",\n \"IconName = NewType('IconName', str)\",\n '\"\"\"Name of an icon. Should be generated via `viser.Icon.*`.\"\"\"',\n \"\",\n \"class Icon:\",\n ' \"\"\"\\'Enum\\' class for referencing Tabler icons.',\n \"\",\n \" We don't subclass enum.Enum for performance reasons -- importing an enum with\",\n \" thousands of names can result in import times in the hundreds of milliseconds.\",\n ' \"\"\"',\n \"\",\n ]\n + [\n # Prefix all icon names with ICON_, since some of them start with\n # numbers and can't directly be used as Python names.\n f\" {enum_name_from_icon(icon)}: IconName = IconName('{icon}')\"\n for icon in icon_names\n ]\n )\n )\n\n # Generate source. This is used at runtime + by Sphinx for documentation.\n (HERE_DIR / \"_icons_enum.py\").write_text(\n \"\\n\".join(\n [\n \"# Automatically generated by `_icons_generate_enum.py`\",\n \"# See https://tabler-icons.io/\",\n \"from typing import NewType\",\n \"\",\n \"IconName = NewType('IconName', str)\",\n '\"\"\"Name of an icon. Should be generated via `viser.Icon.*`.\"\"\"',\n \"\",\n \"\",\n \"class _IconStringConverter(type):\",\n \" def __getattr__(self, __name: str) -> IconName:\",\n ' if not __name.startswith(\"_\"):',\n ' return IconName(__name.lower().replace(\"_\", \"-\"))',\n \" else:\",\n \" raise AttributeError()\",\n \"\",\n \"\",\n \"class Icon(metaclass=_IconStringConverter):\",\n ' \"\"\"\\'Enum\\' class for referencing Tabler icons.',\n \"\",\n \" We don't subclass enum.Enum for performance reasons -- importing an enum with\",\n \" thousands of names can result in import times in the hundreds of milliseconds.\",\n \"\",\n \" Attributes:\",\n ]\n + [\n # Prefix all icon names with ICON_, since some of them start with\n # numbers and can't directly be used as Python names.\n f\" {enum_name_from_icon(icon)} (IconName): The :code:`{icon}` icon.\"\n for icon in icon_names\n ]\n + [' \"\"\"']\n )\n )\n" }, { "path": "viser/src/viser/_messages.py", "content": "\"\"\"Message type definitions. For synchronization with the TypeScript definitions, see\n`_typescript_interface_gen.py.`\"\"\"\n\nfrom __future__ import annotations\n\nimport dataclasses\nimport uuid\nfrom typing import (\n Any,\n Callable,\n ClassVar,\n Dict,\n Optional,\n Tuple,\n Type,\n TypeVar,\n Union,\n)\n\nimport numpy as onp\nimport numpy.typing as onpt\nfrom typing_extensions import Annotated, Literal, NotRequired, TypedDict, override\n\nfrom . import infra, theme\n\nGuiSliderMark = TypedDict(\"GuiSliderMark\", {\"value\": float, \"label\": NotRequired[str]})\nColor = Literal[\n \"dark\",\n \"gray\",\n \"red\",\n \"pink\",\n \"grape\",\n \"violet\",\n \"indigo\",\n \"blue\",\n \"cyan\",\n \"green\",\n \"lime\",\n \"yellow\",\n \"orange\",\n \"teal\",\n]\n\n\nclass Message(infra.Message):\n _tags: ClassVar[Tuple[str, ...]] = tuple()\n\n @override\n def redundancy_key(self) -> str:\n \"\"\"Returns a unique key for this message, used for detecting redundant\n messages.\n\n For example: if we send 1000 GuiSetValue messages for the same GUI element, we\n should only keep the latest messages.\n \"\"\"\n parts = [type(self).__name__]\n\n # Scene node manipulation messages all have a \"name\" field.\n node_name = getattr(self, \"name\", None)\n if node_name is not None:\n parts.append(node_name)\n\n # GUI and notification messages all have an \"id\" field.\n node_name = getattr(self, \"id\", None)\n if node_name is not None:\n parts.append(node_name)\n\n return \"_\".join(parts)\n\n\nT = TypeVar(\"T\", bound=Type[Message])\n\n\ndef tag_class(tag: str) -> Callable[[T], T]:\n \"\"\"Decorator for tagging a class with a `type` field.\"\"\"\n\n def wrapper(cls: T) -> T:\n cls._tags = (cls._tags or ()) + (tag,)\n return cls\n\n return wrapper\n\n\n@dataclasses.dataclass\nclass RunJavascriptMessage(Message):\n \"\"\"Message for running some arbitrary Javascript on the client.\n We use this to set up the Plotly.js package, via the plotly.min.js source\n code.\"\"\"\n\n source: str\n\n @override\n def redundancy_key(self) -> str:\n # Never cull these messages.\n return str(uuid.uuid4())\n\n\n@dataclasses.dataclass\nclass NotificationMessage(Message):\n \"\"\"Notification message.\"\"\"\n\n mode: Literal[\"show\", \"update\"]\n id: str\n title: str\n body: str\n loading: bool\n with_close_button: bool\n auto_close: Union[int, Literal[False]]\n color: Optional[Color]\n\n\n@dataclasses.dataclass\nclass RemoveNotificationMessage(Message):\n \"\"\"Remove a specific notification.\"\"\"\n\n id: str\n\n\n@dataclasses.dataclass\nclass ViewerCameraMessage(Message):\n \"\"\"Message for a posed viewer camera.\n Pose is in the form T_world_camera, OpenCV convention, +Z forward.\"\"\"\n\n wxyz: Tuple[float, float, float, float]\n position: Tuple[float, float, float]\n fov: float\n aspect: float\n look_at: Tuple[float, float, float]\n up_direction: Tuple[float, float, float]\n\n\n# The list of scene pointer events supported by the viser frontend.\nScenePointerEventType = Literal[\"click\", \"rect-select\"]\n\n\n@dataclasses.dataclass\nclass ScenePointerMessage(Message):\n \"\"\"Message for a raycast-like pointer in the scene.\n origin is the viewing camera position, in world coordinates.\n direction is the vector if a ray is projected from the camera through the clicked pixel,\n \"\"\"\n\n # Later we can add `double_click`, `move`, `down`, `up`, etc.\n event_type: ScenePointerEventType\n ray_origin: Optional[Tuple[float, float, float]]\n ray_direction: Optional[Tuple[float, float, float]]\n screen_pos: Tuple[Tuple[float, float], ...]\n\n\n@dataclasses.dataclass\nclass ScenePointerEnableMessage(Message):\n \"\"\"Message to enable/disable scene click events.\"\"\"\n\n enable: bool\n event_type: ScenePointerEventType\n\n @override\n def redundancy_key(self) -> str:\n return (\n type(self).__name__ + \"-\" + self.event_type + \"-\" + str(self.enable).lower()\n )\n\n\n@dataclasses.dataclass\nclass CameraFrustumMessage(Message):\n \"\"\"Variant of CameraMessage used for visualizing camera frustums.\n\n OpenCV convention, +Z forward.\"\"\"\n\n name: str\n fov: float\n aspect: float\n scale: float\n color: int\n thickness: float\n image_media_type: Optional[Literal[\"image/jpeg\", \"image/png\"]]\n image_binary: Optional[bytes]\n\n\n@dataclasses.dataclass\nclass GlbMessage(Message):\n \"\"\"GlTF message.\"\"\"\n\n name: str\n glb_data: bytes\n scale: float\n\n\n@dataclasses.dataclass\nclass FrameMessage(Message):\n \"\"\"Coordinate frame message.\"\"\"\n\n name: str\n show_axes: bool\n axes_length: float\n axes_radius: float\n origin_radius: float\n\n\n@dataclasses.dataclass\nclass BatchedAxesMessage(Message):\n \"\"\"Batched axes message.\n\n Positions and orientations should follow a `T_parent_local` convention, which\n corresponds to the R matrix and t vector in `p_parent = [R | t] p_local`.\"\"\"\n\n name: str\n wxyzs_batched: onpt.NDArray[onp.float32]\n positions_batched: onpt.NDArray[onp.float32]\n axes_length: float\n axes_radius: float\n\n\n@dataclasses.dataclass\nclass GridMessage(Message):\n \"\"\"Grid message. Helpful for visualizing things like ground planes.\"\"\"\n\n name: str\n\n width: float\n height: float\n width_segments: int\n height_segments: int\n\n plane: Literal[\"xz\", \"xy\", \"yx\", \"yz\", \"zx\", \"zy\"]\n\n cell_color: int\n cell_thickness: float\n cell_size: float\n\n section_color: int\n section_thickness: float\n section_size: float\n\n\n@dataclasses.dataclass\nclass LabelMessage(Message):\n \"\"\"Add a 2D label to the scene.\"\"\"\n\n name: str\n text: str\n\n\n@dataclasses.dataclass\nclass Gui3DMessage(Message):\n \"\"\"Add a 3D gui element to the scene.\"\"\"\n\n order: float\n name: str\n container_id: str\n\n\n@dataclasses.dataclass\nclass PointCloudMessage(Message):\n \"\"\"Point cloud message.\n\n Positions are internally canonicalized to float32, colors to uint8.\n\n Float color inputs should be in the range [0,1], int color inputs should be in the\n range [0,255].\"\"\"\n\n name: str\n points: onpt.NDArray[onp.float32]\n colors: onpt.NDArray[onp.uint8]\n point_size: float\n point_ball_norm: float\n\n def __post_init__(self):\n # Check shapes.\n assert self.points.shape == self.colors.shape\n assert self.points.shape[-1] == 3\n\n # Check dtypes.\n assert self.points.dtype == onp.float32\n assert self.colors.dtype == onp.uint8\n\n\n@dataclasses.dataclass\nclass MeshBoneMessage(Message):\n \"\"\"Message for a bone of a skinned mesh.\"\"\"\n\n name: str\n\n\n@dataclasses.dataclass\nclass MeshMessage(Message):\n \"\"\"Mesh message.\n\n Vertices are internally canonicalized to float32, faces to uint32.\"\"\"\n\n name: str\n vertices: onpt.NDArray[onp.float32]\n faces: onpt.NDArray[onp.uint32]\n\n color: Optional[int]\n vertex_colors: Optional[onpt.NDArray[onp.uint8]]\n\n wireframe: bool\n opacity: Optional[float]\n flat_shading: bool\n side: Literal[\"front\", \"back\", \"double\"]\n material: Literal[\"standard\", \"toon3\", \"toon5\"]\n\n def __post_init__(self):\n # Check shapes.\n assert self.vertices.shape[-1] == 3\n assert self.faces.shape[-1] == 3\n\n\n@dataclasses.dataclass\nclass SkinnedMeshMessage(MeshMessage):\n \"\"\"Mesh message.\n\n Vertices are internally canonicalized to float32, faces to uint32.\"\"\"\n\n bone_wxyzs: Tuple[Tuple[float, float, float, float], ...]\n bone_positions: Tuple[Tuple[float, float, float], ...]\n skin_indices: onpt.NDArray[onp.uint16]\n skin_weights: onpt.NDArray[onp.float32]\n\n def __post_init__(self):\n # Check shapes.\n assert self.vertices.shape[-1] == 3\n assert self.faces.shape[-1] == 3\n assert self.skin_weights is not None\n assert (\n self.skin_indices.shape\n == self.skin_weights.shape\n == (self.vertices.shape[0], 4)\n )\n\n\n@dataclasses.dataclass\nclass SetBoneOrientationMessage(Message):\n \"\"\"Server -> client message to set a skinned mesh bone's orientation.\n\n As with all other messages, transforms take the `T_parent_local` convention.\"\"\"\n\n name: str\n bone_index: int\n wxyz: Tuple[float, float, float, float]\n\n @override\n def redundancy_key(self) -> str:\n return type(self).__name__ + \"-\" + self.name + \"-\" + str(self.bone_index)\n\n\n@dataclasses.dataclass\nclass SetBonePositionMessage(Message):\n \"\"\"Server -> client message to set a skinned mesh bone's position.\n\n As with all other messages, transforms take the `T_parent_local` convention.\"\"\"\n\n name: str\n bone_index: int\n position: Tuple[float, float, float]\n\n @override\n def redundancy_key(self) -> str:\n return type(self).__name__ + \"-\" + self.name + \"-\" + str(self.bone_index)\n\n\n@dataclasses.dataclass\nclass TransformControlsMessage(Message):\n \"\"\"Message for transform gizmos.\"\"\"\n\n name: str\n scale: float\n line_width: float\n fixed: bool\n auto_transform: bool\n active_axes: Tuple[bool, bool, bool]\n disable_axes: bool\n disable_sliders: bool\n disable_rotations: bool\n translation_limits: Tuple[\n Tuple[float, float], Tuple[float, float], Tuple[float, float]\n ]\n rotation_limits: Tuple[\n Tuple[float, float], Tuple[float, float], Tuple[float, float]\n ]\n depth_test: bool\n opacity: float\n\n\n@dataclasses.dataclass\nclass SetCameraPositionMessage(Message):\n \"\"\"Server -> client message to set the camera's position.\"\"\"\n\n position: Tuple[float, float, float]\n\n\n@dataclasses.dataclass\nclass SetCameraUpDirectionMessage(Message):\n \"\"\"Server -> client message to set the camera's up direction.\"\"\"\n\n position: Tuple[float, float, float]\n\n\n@dataclasses.dataclass\nclass SetCameraLookAtMessage(Message):\n \"\"\"Server -> client message to set the camera's look-at point.\"\"\"\n\n look_at: Tuple[float, float, float]\n\n\n@dataclasses.dataclass\nclass SetCameraFovMessage(Message):\n \"\"\"Server -> client message to set the camera's field of view.\"\"\"\n\n fov: float\n\n\n@dataclasses.dataclass\nclass SetOrientationMessage(Message):\n \"\"\"Server -> client message to set a scene node's orientation.\n\n As with all other messages, transforms take the `T_parent_local` convention.\"\"\"\n\n name: str\n wxyz: Tuple[float, float, float, float]\n\n\n@dataclasses.dataclass\nclass SetPositionMessage(Message):\n \"\"\"Server -> client message to set a scene node's position.\n\n As with all other messages, transforms take the `T_parent_local` convention.\"\"\"\n\n name: str\n position: Tuple[float, float, float]\n\n\n@dataclasses.dataclass\nclass TransformControlsUpdateMessage(Message):\n \"\"\"Client -> server message when a transform control is updated.\n\n As with all other messages, transforms take the `T_parent_local` convention.\"\"\"\n\n name: str\n wxyz: Tuple[float, float, float, float]\n position: Tuple[float, float, float]\n\n\n@dataclasses.dataclass\nclass BackgroundImageMessage(Message):\n \"\"\"Message for rendering a background image.\"\"\"\n\n media_type: Literal[\"image/jpeg\", \"image/png\"]\n rgb_bytes: bytes\n depth_bytes: Optional[bytes]\n\n\n@dataclasses.dataclass\nclass ImageMessage(Message):\n \"\"\"Message for rendering 2D images.\"\"\"\n\n name: str\n media_type: Literal[\"image/jpeg\", \"image/png\"]\n data: bytes\n render_width: float\n render_height: float\n\n\n@dataclasses.dataclass\nclass RemoveSceneNodeMessage(Message):\n \"\"\"Remove a particular node from the scene.\"\"\"\n\n name: str\n\n\n@dataclasses.dataclass\nclass SetSceneNodeVisibilityMessage(Message):\n \"\"\"Set the visibility of a particular node in the scene.\"\"\"\n\n name: str\n visible: bool\n\n\n@dataclasses.dataclass\nclass SetSceneNodeClickableMessage(Message):\n \"\"\"Set the clickability of a particular node in the scene.\"\"\"\n\n name: str\n clickable: bool\n\n\n@dataclasses.dataclass\nclass SceneNodeClickMessage(Message):\n \"\"\"Message for clicked objects.\"\"\"\n\n name: str\n instance_index: Optional[int]\n \"\"\"Instance index. Currently only used for batched axes.\"\"\"\n ray_origin: Tuple[float, float, float]\n ray_direction: Tuple[float, float, float]\n screen_pos: Tuple[float, float]\n\n\n@dataclasses.dataclass\nclass ResetSceneMessage(Message):\n \"\"\"Reset scene.\"\"\"\n\n\n@dataclasses.dataclass\nclass ResetGuiMessage(Message):\n \"\"\"Reset GUI.\"\"\"\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddFolderMessage(Message):\n order: float\n id: str\n label: str\n container_id: str\n expand_by_default: bool\n visible: bool\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddMarkdownMessage(Message):\n order: float\n id: str\n markdown: str\n container_id: str\n visible: bool\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddProgressBarMessage(Message):\n order: float\n id: str\n value: float\n animated: bool\n color: Optional[Color]\n container_id: str\n visible: bool\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddPlotlyMessage(Message):\n order: float\n id: str\n plotly_json_str: str\n aspect: float\n container_id: str\n visible: bool\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddTabGroupMessage(Message):\n order: float\n id: str\n container_id: str\n tab_labels: Tuple[str, ...]\n tab_icons_html: Tuple[Union[str, None], ...]\n tab_container_ids: Tuple[str, ...]\n visible: bool\n\n\n@dataclasses.dataclass\nclass _GuiAddInputBase(Message):\n \"\"\"Base message type containing fields commonly used by GUI inputs.\"\"\"\n\n order: float\n id: str\n label: str\n container_id: str\n hint: Optional[str]\n value: Any\n visible: bool\n disabled: bool\n\n\n@dataclasses.dataclass\nclass GuiModalMessage(Message):\n order: float\n id: str\n title: str\n\n\n@dataclasses.dataclass\nclass GuiCloseModalMessage(Message):\n id: str\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddButtonMessage(_GuiAddInputBase):\n # All GUI elements currently need an `value` field.\n # This makes our job on the frontend easier.\n value: bool\n color: Optional[Color]\n icon_html: Optional[str]\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddUploadButtonMessage(_GuiAddInputBase):\n color: Optional[Color]\n icon_html: Optional[str]\n mime_type: str\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddSliderMessage(_GuiAddInputBase):\n min: float\n max: float\n step: Optional[float]\n value: float\n precision: int\n marks: Optional[Tuple[GuiSliderMark, ...]] = None\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddMultiSliderMessage(_GuiAddInputBase):\n min: float\n max: float\n step: Optional[float]\n min_range: Optional[float]\n precision: int\n fixed_endpoints: bool = False\n marks: Optional[Tuple[GuiSliderMark, ...]] = None\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddNumberMessage(_GuiAddInputBase):\n value: float\n precision: int\n step: float\n min: Optional[float]\n max: Optional[float]\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddRgbMessage(_GuiAddInputBase):\n value: Tuple[int, int, int]\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddRgbaMessage(_GuiAddInputBase):\n value: Tuple[int, int, int, int]\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddCheckboxMessage(_GuiAddInputBase):\n value: bool\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddVector2Message(_GuiAddInputBase):\n value: Tuple[float, float]\n min: Optional[Tuple[float, float]]\n max: Optional[Tuple[float, float]]\n step: float\n precision: int\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddVector3Message(_GuiAddInputBase):\n value: Tuple[float, float, float]\n min: Optional[Tuple[float, float, float]]\n max: Optional[Tuple[float, float, float]]\n step: float\n precision: int\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddTextMessage(_GuiAddInputBase):\n value: str\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddDropdownMessage(_GuiAddInputBase):\n value: str\n options: Tuple[str, ...]\n\n\n@tag_class(\"GuiAddComponentMessage\")\n@dataclasses.dataclass\nclass GuiAddButtonGroupMessage(_GuiAddInputBase):\n value: str\n options: Tuple[str, ...]\n\n\n@dataclasses.dataclass\nclass GuiRemoveMessage(Message):\n \"\"\"Sent server->client to remove a GUI element.\"\"\"\n\n id: str\n\n\n@dataclasses.dataclass\nclass GuiUpdateMessage(Message):\n \"\"\"Sent client<->server when any property of a GUI component is changed.\"\"\"\n\n id: str\n updates: Annotated[\n Dict[str, Any],\n infra.TypeScriptAnnotationOverride(\"Partial\"),\n ]\n \"\"\"Mapping from property name to new value.\"\"\"\n\n @override\n def redundancy_key(self) -> str:\n return (\n type(self).__name__\n + \"-\"\n + self.id\n + \"-\"\n + \",\".join(list(self.updates.keys()))\n )\n\n\n@dataclasses.dataclass\nclass ThemeConfigurationMessage(Message):\n \"\"\"Message from server->client to configure parts of the GUI.\"\"\"\n\n titlebar_content: Optional[theme.TitlebarConfig]\n control_layout: Literal[\"floating\", \"collapsible\", \"fixed\"]\n control_width: Literal[\"small\", \"medium\", \"large\"]\n show_logo: bool\n show_share_button: bool\n dark_mode: bool\n colors: Optional[Tuple[str, str, str, str, str, str, str, str, str, str]]\n\n\n@dataclasses.dataclass\nclass CatmullRomSplineMessage(Message):\n \"\"\"Message from server->client carrying Catmull-Rom spline information.\"\"\"\n\n name: str\n positions: Tuple[Tuple[float, float, float], ...]\n curve_type: Literal[\"centripetal\", \"chordal\", \"catmullrom\"]\n tension: float\n closed: bool\n line_width: float\n color: int\n segments: Optional[int]\n\n\n@dataclasses.dataclass\nclass CubicBezierSplineMessage(Message):\n \"\"\"Message from server->client carrying Cubic Bezier spline information.\"\"\"\n\n name: str\n positions: Tuple[Tuple[float, float, float], ...]\n control_points: Tuple[Tuple[float, float, float], ...]\n line_width: float\n color: int\n segments: Optional[int]\n\n\n@dataclasses.dataclass\nclass GaussianSplatsMessage(Message):\n \"\"\"Message from server->client carrying splattable Gaussians.\"\"\"\n\n name: str\n\n # Memory layout is borrowed from:\n # https://github.com/antimatter15/splat\n buffer: onpt.NDArray[onp.uint32]\n \"\"\"Our buffer will contain:\n - x as f32\n - y as f32\n - z as f32\n - (unused)\n - cov1 (f16), cov2 (f16)\n - cov3 (f16), cov4 (f16)\n - cov5 (f16), cov6 (f16)\n - rgba (int32)\n Where cov1-6 are the upper triangular elements of the covariance matrix.\"\"\"\n\n\n@dataclasses.dataclass\nclass GetRenderRequestMessage(Message):\n \"\"\"Message from server->client requesting a render of the current viewport.\"\"\"\n\n format: Literal[\"image/jpeg\", \"image/png\"]\n height: int\n width: int\n quality: int\n\n\n@dataclasses.dataclass\nclass GetRenderResponseMessage(Message):\n \"\"\"Message from client->server carrying a render.\"\"\"\n\n payload: bytes\n\n\n@dataclasses.dataclass\nclass FileTransferStart(Message):\n \"\"\"Signal that a file is about to be sent.\"\"\"\n\n source_component_id: Optional[str]\n \"\"\"Origin GUI component, used for client->server file uploads.\"\"\"\n transfer_uuid: str\n filename: str\n mime_type: str\n part_count: int\n size_bytes: int\n\n @override\n def redundancy_key(self) -> str:\n return type(self).__name__ + \"-\" + self.transfer_uuid\n\n\n@dataclasses.dataclass\nclass FileTransferPart(Message):\n \"\"\"Send a file for clients to download or upload files from client.\"\"\"\n\n # TODO: it would make sense to rename all \"id\" instances to \"uuid\" for GUI component ids.\n source_component_id: Optional[str]\n transfer_uuid: str\n part: int\n content: bytes\n\n @override\n def redundancy_key(self) -> str:\n return type(self).__name__ + \"-\" + self.transfer_uuid + \"-\" + str(self.part)\n\n\n@dataclasses.dataclass\nclass FileTransferPartAck(Message):\n \"\"\"Send a file for clients to download or upload files from client.\"\"\"\n\n source_component_id: Optional[str]\n transfer_uuid: str\n transferred_bytes: int\n total_bytes: int\n\n @override\n def redundancy_key(self) -> str:\n return (\n type(self).__name__\n + \"-\"\n + self.transfer_uuid\n + \"-\"\n + str(self.transferred_bytes)\n )\n\n\n@dataclasses.dataclass\nclass ShareUrlRequest(Message):\n \"\"\"Message from client->server to connect to the share URL server.\"\"\"\n\n\n@dataclasses.dataclass\nclass ShareUrlUpdated(Message):\n \"\"\"Message from server->client to indicate that the share URL has been updated.\"\"\"\n\n share_url: Optional[str]\n\n\n@dataclasses.dataclass\nclass ShareUrlDisconnect(Message):\n \"\"\"Message from client->server to disconnect from the share URL server.\"\"\"\n\n\n@dataclasses.dataclass\nclass SetGuiPanelLabelMessage(Message):\n \"\"\"Message from server->client to set the label of the GUI panel.\"\"\"\n\n label: Optional[str]\n" }, { "path": "viser/src/viser/_notification_handle.py", "content": "from __future__ import annotations\n\nimport dataclasses\nfrom typing import Literal\n\nfrom ._gui_api import Color\nfrom ._messages import NotificationMessage, RemoveNotificationMessage\nfrom .infra._infra import WebsockClientConnection\n\n\n@dataclasses.dataclass\nclass _NotificationHandleState:\n websock_interface: WebsockClientConnection\n id: str\n title: str\n body: str\n loading: bool\n with_close_button: bool\n auto_close: int | Literal[False]\n color: Color | None\n\n\n@dataclasses.dataclass\nclass NotificationHandle:\n \"\"\"Handle for a notification in our visualizer.\"\"\"\n\n _impl: _NotificationHandleState\n\n def _sync_with_client(self, first: bool = False) -> None:\n m = NotificationMessage(\n \"show\" if first else \"update\",\n self._impl.id,\n self._impl.title,\n self._impl.body,\n self._impl.loading,\n self._impl.with_close_button,\n self._impl.auto_close,\n self._impl.color,\n )\n self._impl.websock_interface.queue_message(m)\n\n @property\n def title(self) -> str:\n \"\"\"Title to display on the notification.\"\"\"\n return self._impl.title\n\n @title.setter\n def title(self, title: str) -> None:\n if title == self._impl.title:\n return\n\n self._impl.title = title\n self._sync_with_client()\n\n @property\n def body(self) -> str:\n \"\"\"Message to display on the notification body.\"\"\"\n return self._impl.body\n\n @body.setter\n def body(self, body: str) -> None:\n if body == self._impl.body:\n return\n\n self._impl.body = body\n self._sync_with_client()\n\n @property\n def loading(self) -> bool:\n \"\"\"Whether the notification shows loading icon.\"\"\"\n return self._impl.loading\n\n @loading.setter\n def loading(self, loading: bool) -> None:\n if loading == self._impl.loading:\n return\n\n self._impl.loading = loading\n self._sync_with_client()\n\n @property\n def with_close_button(self) -> bool:\n \"\"\"Whether the notification can be manually closed.\"\"\"\n return self._impl.with_close_button\n\n @with_close_button.setter\n def with_close_button(self, with_close_button: bool) -> None:\n if with_close_button == self._impl.with_close_button:\n return\n\n self._impl.with_close_button = with_close_button\n self._sync_with_client()\n\n @property\n def auto_close(self) -> int | Literal[False]:\n \"\"\"Time in ms before the notification automatically closes;\n otherwise False such that the notification never closes on its own.\"\"\"\n return self._impl.auto_close\n\n @auto_close.setter\n def auto_close(self, auto_close: int | Literal[False]) -> None:\n if auto_close == self._impl.auto_close:\n return\n\n self._impl.auto_close = auto_close\n self._sync_with_client()\n\n @property\n def color(self) -> Color | None:\n \"\"\"Color of the notification.\"\"\"\n return self._impl.color\n\n @color.setter\n def color(self, color: Color | None) -> None:\n if color == self._impl.color:\n return\n\n self._impl.color = color\n self._sync_with_client()\n\n def remove(self) -> None:\n self._impl.websock_interface.queue_message(\n RemoveNotificationMessage(self._impl.id)\n )\n" }, { "path": "viser/src/viser/_scene_api.py", "content": "from __future__ import annotations\n\nimport io\nimport time\nimport warnings\nfrom concurrent.futures import ThreadPoolExecutor\nfrom typing import TYPE_CHECKING, Callable, Tuple, TypeVar, Union, cast, get_args\n\nimport imageio.v3 as iio\nimport numpy as onp\nimport numpy.typing as onpt\nfrom typing_extensions import Literal, ParamSpec, TypeAlias, assert_never\n\nfrom . import _messages\nfrom . import transforms as tf\nfrom ._scene_handles import (\n BatchedAxesHandle,\n BoneState,\n CameraFrustumHandle,\n FrameHandle,\n GaussianSplatHandle,\n GlbHandle,\n Gui3dContainerHandle,\n ImageHandle,\n LabelHandle,\n MeshHandle,\n MeshSkinnedBoneHandle,\n MeshSkinnedHandle,\n PointCloudHandle,\n SceneNodeHandle,\n SceneNodePointerEvent,\n ScenePointerEvent,\n TransformControlsHandle,\n _SceneNodeHandleState,\n _TransformControlsState,\n)\n\nif TYPE_CHECKING:\n import trimesh\n\n from ._viser import ClientHandle, ViserServer\n from .infra import ClientId\n\n\nP = ParamSpec(\"P\")\n\n\ndef _colors_to_uint8(colors: onp.ndarray) -> onpt.NDArray[onp.uint8]:\n \"\"\"Convert intensity values to uint8. We assume the range [0,1] for floats, and\n [0,255] for integers. Accepts any shape.\"\"\"\n if colors.dtype != onp.uint8:\n if onp.issubdtype(colors.dtype, onp.floating):\n colors = onp.clip(colors * 255.0, 0, 255).astype(onp.uint8)\n if onp.issubdtype(colors.dtype, onp.integer):\n colors = onp.clip(colors, 0, 255).astype(onp.uint8)\n return colors\n\n\nRgbTupleOrArray: TypeAlias = Union[\n Tuple[int, int, int], Tuple[float, float, float], onp.ndarray\n]\n\n\ndef _encode_rgb(rgb: RgbTupleOrArray) -> int:\n if isinstance(rgb, onp.ndarray):\n assert rgb.shape == (3,)\n rgb_fixed = tuple(\n value if onp.issubdtype(type(value), onp.integer) else int(value * 255)\n for value in rgb\n )\n assert len(rgb_fixed) == 3\n return int(rgb_fixed[0] * (256**2) + rgb_fixed[1] * 256 + rgb_fixed[2])\n\n\ndef _encode_image_binary(\n image: onp.ndarray,\n format: Literal[\"png\", \"jpeg\"],\n jpeg_quality: int | None = None,\n) -> tuple[Literal[\"image/png\", \"image/jpeg\"], bytes]:\n media_type: Literal[\"image/png\", \"image/jpeg\"]\n image = _colors_to_uint8(image)\n with io.BytesIO() as data_buffer:\n if format == \"png\":\n media_type = \"image/png\"\n iio.imwrite(data_buffer, image, extension=\".png\")\n elif format == \"jpeg\":\n media_type = \"image/jpeg\"\n iio.imwrite(\n data_buffer,\n image[..., :3], # Strip alpha.\n extension=\".jpeg\",\n quality=75 if jpeg_quality is None else jpeg_quality,\n )\n else:\n assert_never(format)\n binary = data_buffer.getvalue()\n return media_type, binary\n\n\nTVector = TypeVar(\"TVector\", bound=tuple)\n\n\ndef cast_vector(vector: TVector | onp.ndarray, length: int) -> TVector:\n if not isinstance(vector, tuple):\n assert cast(onp.ndarray, vector).shape == (\n length,\n ), f\"Expected vector of shape {(length,)}, but got {vector.shape} instead\"\n return cast(TVector, tuple(map(float, vector)))\n\n\nclass SceneApi:\n \"\"\"Interface for adding 3D primitives to the scene.\n\n Used by both our global server object, for sharing the same GUI elements\n with all clients, and by individual client handles.\"\"\"\n\n def __init__(\n self,\n owner: ViserServer | ClientHandle, # Who do I belong to?\n thread_executor: ThreadPoolExecutor,\n ) -> None:\n from ._viser import ViserServer\n\n self._owner = owner\n \"\"\"Entity that owns this API.\"\"\"\n\n self._websock_interface = (\n owner._websock_server\n if isinstance(owner, ViserServer)\n else owner._websock_connection\n )\n \"\"\"Interface for sending and listening to messages.\"\"\"\n\n self.world_axes: FrameHandle = FrameHandle(\n _SceneNodeHandleState(\n \"/WorldAxes\",\n self,\n wxyz=onp.array([1.0, 0.0, 0.0, 0.0]),\n position=onp.zeros(3),\n )\n )\n \"\"\"Handle for the world axes, which are created by default.\"\"\"\n\n # Hide world axes on initialization.\n if isinstance(owner, ViserServer):\n self.world_axes.visible = False\n\n self._handle_from_transform_controls_name: dict[\n str, TransformControlsHandle\n ] = {}\n self._handle_from_node_name: dict[str, SceneNodeHandle] = {}\n\n self._scene_pointer_cb: Callable[[ScenePointerEvent], None] | None = None\n self._scene_pointer_done_cb: Callable[[], None] = lambda: None\n self._scene_pointer_event_type: _messages.ScenePointerEventType | None = None\n\n self._websock_interface.register_handler(\n _messages.TransformControlsUpdateMessage,\n self._handle_transform_controls_updates,\n )\n self._websock_interface.register_handler(\n _messages.SceneNodeClickMessage,\n self._handle_node_click_updates,\n )\n self._websock_interface.register_handler(\n _messages.ScenePointerMessage,\n self._handle_scene_pointer_updates,\n )\n\n self._thread_executor = thread_executor\n\n def set_up_direction(\n self,\n direction: Literal[\"+x\", \"+y\", \"+z\", \"-x\", \"-y\", \"-z\"]\n | tuple[float, float, float]\n | onp.ndarray,\n ) -> None:\n \"\"\"Set the global up direction of the scene. By default we follow +Z-up\n (similar to Blender, 3DS Max, ROS, etc), the most common alternative is\n +Y (OpenGL, Maya, etc).\n\n Args:\n direction: New up direction. Can either be a string (one of +x, +y,\n +z, -x, -y, -z) or a length-3 direction vector.\n \"\"\"\n if isinstance(direction, str):\n direction = {\n \"+x\": (1, 0, 0),\n \"+y\": (0, 1, 0),\n \"+z\": (0, 0, 1),\n \"-x\": (-1, 0, 0),\n \"-y\": (0, -1, 0),\n \"-z\": (0, 0, -1),\n }[direction]\n assert not isinstance(direction, str)\n\n default_three_up = onp.array([0.0, 1.0, 0.0])\n direction = onp.asarray(direction)\n\n def rotate_between(before: onp.ndarray, after: onp.ndarray) -> tf.SO3:\n assert before.shape == after.shape == (3,)\n before = before / onp.linalg.norm(before)\n after = after / onp.linalg.norm(after)\n\n angle = onp.arccos(onp.clip(onp.dot(before, after), -1, 1))\n axis = onp.cross(before, after)\n if onp.allclose(axis, onp.zeros(3), rtol=1e-3, atol=1e-5):\n unit_vector = onp.arange(3) == onp.argmin(onp.abs(before))\n axis = onp.cross(before, unit_vector)\n axis = axis / onp.linalg.norm(axis)\n return tf.SO3.exp(angle * axis)\n\n R_threeworld_world = rotate_between(direction, default_three_up)\n\n # Rotate the world frame such that:\n # If we set +Y to up, +X and +Z should face the camera.\n # If we set +Z to up, +X and +Y should face the camera.\n # In App.tsx, the camera is initialized at [-3, 3, -3] in the threejs\n # coordinate frame.\n desired_fwd = onp.array([-1.0, 0.0, -1.0]) / onp.sqrt(2.0)\n current_fwd = R_threeworld_world @ (onp.ones(3) / onp.sqrt(3.0))\n current_fwd = current_fwd * onp.array([1.0, 0.0, 1.0])\n current_fwd = current_fwd / onp.linalg.norm(current_fwd)\n R_threeworld_world = (\n tf.SO3.from_y_radians( # Rotate around the null space / up direction.\n onp.arctan2(\n onp.cross(current_fwd, desired_fwd)[1],\n onp.dot(current_fwd, desired_fwd),\n ),\n )\n @ R_threeworld_world\n )\n\n if not onp.any(onp.isnan(R_threeworld_world.wxyz)):\n # Set the orientation of the root node.\n self._websock_interface.queue_message(\n _messages.SetOrientationMessage(\n \"\", cast_vector(R_threeworld_world.wxyz, 4)\n )\n )\n\n def set_global_visibility(self, visible: bool) -> None:\n \"\"\"Set visibility for all scene nodes. If set to False, all scene nodes\n will be hidden.\n\n This can be useful when we've called\n :meth:`SceneApi.set_background_image()`, and want to hide everything\n except for the background.\n\n Args:\n visible: Whether or not all scene nodes should be visible.\n \"\"\"\n self._websock_interface.queue_message(\n _messages.SetSceneNodeVisibilityMessage(\"\", visible)\n )\n\n def add_glb(\n self,\n name: str,\n glb_data: bytes,\n scale=1.0,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> GlbHandle:\n \"\"\"Add a general 3D asset via binary glTF (GLB).\n\n For glTF files, it's often simpler to use `trimesh.load()` with\n `.add_mesh_trimesh()`. This will call `.add_glb()` under the hood.\n\n For glTF features not supported by trimesh, glTF to GLB conversion can\n also be done programatically with libraries like `pygltflib`.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n glb_data: A binary payload.\n scale: A scale for resizing the GLB asset.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n self._websock_interface.queue_message(\n _messages.GlbMessage(name, glb_data, scale)\n )\n return GlbHandle._make(self, name, wxyz, position, visible)\n\n def add_spline_catmull_rom(\n self,\n name: str,\n positions: tuple[tuple[float, float, float], ...] | onp.ndarray,\n curve_type: Literal[\"centripetal\", \"chordal\", \"catmullrom\"] = \"centripetal\",\n tension: float = 0.5,\n closed: bool = False,\n line_width: float = 1,\n color: RgbTupleOrArray = (20, 20, 20),\n segments: int | None = None,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> SceneNodeHandle:\n \"\"\"Add a spline to the scene using Catmull-Rom interpolation.\n\n This method creates a spline based on a set of positions and interpolates\n them using the Catmull-Rom algorithm. This can be used to create smooth curves.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n positions: A tuple of 3D positions (x, y, z) defining the spline's path.\n curve_type: Type of the curve ('centripetal', 'chordal', 'catmullrom').\n tension: Tension of the curve. Affects the tightness of the curve.\n closed: Boolean indicating if the spline is closed (forms a loop).\n line_width: Width of the spline line.\n color: Color of the spline as an RGB tuple.\n segments: Number of segments to divide the spline into.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n if isinstance(positions, onp.ndarray):\n assert len(positions.shape) == 2 and positions.shape[1] == 3\n positions = tuple(map(tuple, positions)) # type: ignore\n assert len(positions[0]) == 3\n assert isinstance(positions, tuple)\n self._websock_interface.queue_message(\n _messages.CatmullRomSplineMessage(\n name,\n positions,\n curve_type,\n tension,\n closed,\n line_width,\n _encode_rgb(color),\n segments=segments,\n )\n )\n return SceneNodeHandle._make(self, name, wxyz, position, visible)\n\n def add_spline_cubic_bezier(\n self,\n name: str,\n positions: tuple[tuple[float, float, float], ...] | onp.ndarray,\n control_points: tuple[tuple[float, float, float], ...] | onp.ndarray,\n line_width: float = 1,\n color: RgbTupleOrArray = (20, 20, 20),\n segments: int | None = None,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> SceneNodeHandle:\n \"\"\"Add a spline to the scene using Cubic Bezier interpolation.\n\n This method allows for the creation of a cubic Bezier spline based on given\n positions and control points. It is useful for creating complex, smooth,\n curving shapes.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n positions: A tuple of 3D positions (x, y, z) defining the spline's key points.\n control_points: A tuple of control points for Bezier curve shaping.\n line_width: Width of the spline line.\n color: Color of the spline as an RGB tuple.\n segments: Number of segments to divide the spline into.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n\n if isinstance(positions, onp.ndarray):\n assert len(positions.shape) == 2 and positions.shape[1] == 3\n positions = tuple(map(tuple, positions)) # type: ignore\n if isinstance(control_points, onp.ndarray):\n assert len(control_points.shape) == 2 and control_points.shape[1] == 3\n control_points = tuple(map(tuple, control_points)) # type: ignore\n\n assert isinstance(positions, tuple)\n assert isinstance(control_points, tuple)\n assert len(control_points) == (2 * len(positions) - 2)\n self._websock_interface.queue_message(\n _messages.CubicBezierSplineMessage(\n name,\n positions,\n control_points,\n line_width,\n _encode_rgb(color),\n segments=segments,\n )\n )\n return SceneNodeHandle._make(self, name, wxyz, position, visible)\n\n def add_camera_frustum(\n self,\n name: str,\n fov: float,\n aspect: float,\n scale: float = 0.3,\n color: RgbTupleOrArray = (20, 20, 20),\n image: onp.ndarray | None = None,\n format: Literal[\"png\", \"jpeg\"] = \"jpeg\",\n jpeg_quality: int | None = None,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n thickness: float = 1.0,\n ) -> CameraFrustumHandle:\n \"\"\"Add a camera frustum to the scene for visualization.\n\n This method adds a frustum representation, typically used to visualize the\n field of view of a camera. It's helpful for understanding the perspective\n and coverage of a camera in the 3D space.\n\n Like all cameras in the viser Python API, frustums follow the OpenCV [+Z forward,\n +X right, +Y down] convention. fov is vertical in radians; aspect is width over height\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n fov: Field of view of the camera (in radians).\n aspect: Aspect ratio of the camera (width over height).\n scale: Scale factor for the size of the frustum.\n color: Color of the frustum as an RGB tuple.\n image: Optional image to be displayed on the frustum.\n format: Format of the provided image ('png' or 'jpeg').\n jpeg_quality: Quality of the jpeg image (if jpeg format is used).\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n\n if image is not None:\n media_type, binary = _encode_image_binary(\n image, format, jpeg_quality=jpeg_quality\n )\n else:\n media_type = None\n binary = None\n\n self._websock_interface.queue_message(\n _messages.CameraFrustumMessage(\n name=name,\n fov=fov,\n aspect=aspect,\n scale=scale,\n thickness=thickness,\n # (255, 255, 255) => 0xffffff, etc\n color=_encode_rgb(color),\n image_media_type=media_type,\n image_binary=binary,\n )\n )\n return CameraFrustumHandle._make(self, name, wxyz, position, visible)\n\n def add_frame(\n self,\n name: str,\n show_axes: bool = True,\n axes_length: float = 0.5,\n axes_radius: float = 0.025,\n origin_radius: float | None = None,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> FrameHandle:\n \"\"\"Add a coordinate frame to the scene.\n\n This method is used for adding a visual representation of a coordinate\n frame, which can help in understanding the orientation and position of\n objects in 3D space.\n\n For cases where we want to visualize many coordinate frames, like\n trajectories containing thousands or tens of thousands of frames,\n batching and calling :meth:`add_batched_axes()` may be a better choice\n than calling :meth:`add_frame()` in a loop.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n show_axes: Boolean to indicate whether to show the frame as a set of axes + origin sphere.\n axes_length: Length of each axis.\n axes_radius: Radius of each axis.\n origin_radius: Radius of the origin sphere. If not set, defaults to `2 * axes_radius`.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n if origin_radius is None:\n origin_radius = axes_radius * 2\n self._websock_interface.queue_message(\n _messages.FrameMessage(\n name=name,\n show_axes=show_axes,\n axes_length=axes_length,\n axes_radius=axes_radius,\n origin_radius=origin_radius,\n )\n )\n return FrameHandle._make(self, name, wxyz, position, visible)\n\n def add_batched_axes(\n self,\n name: str,\n batched_wxyzs: tuple[tuple[float, float, float, float], ...] | onp.ndarray,\n batched_positions: tuple[tuple[float, float, float], ...] | onp.ndarray,\n axes_length: float = 0.5,\n axes_radius: float = 0.025,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> BatchedAxesHandle:\n \"\"\"Visualize batched sets of coordinate frame axes.\n\n The functionality of :meth:`add_batched_axes()` overlaps significantly\n with :meth:`add_frame()` when `show_axes=True`. The primary difference\n is that :meth:`add_batched_axes()` supports multiple axes via the\n `wxyzs_batched` (shape Nx4) and `positions_batched` (shape Nx3)\n arguments.\n\n Axes that are batched and rendered via a single call to\n `add_batched_axes()` are instanced on the client; this will be much\n faster to render than `add_frame()` called in a loop.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n batched_wxyzs: Float array of shape (N,4).\n batched_positions: Float array of shape (N,3).\n axes_length: Length of each axis.\n axes_radius: Radius of each axis.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n This will be applied to all axes.\n position: Translation to parent frame from local frame (t_pl).\n This will be applied to all axes.\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n batched_wxyzs = onp.asarray(batched_wxyzs)\n batched_positions = onp.asarray(batched_positions)\n\n num_axes = batched_wxyzs.shape[0]\n assert batched_wxyzs.shape == (num_axes, 4)\n assert batched_positions.shape == (num_axes, 3)\n self._websock_interface.queue_message(\n _messages.BatchedAxesMessage(\n name=name,\n wxyzs_batched=batched_wxyzs.astype(onp.float32),\n positions_batched=batched_positions.astype(onp.float32),\n axes_length=axes_length,\n axes_radius=axes_radius,\n )\n )\n return BatchedAxesHandle._make(self, name, wxyz, position, visible)\n\n def add_grid(\n self,\n name: str,\n width: float = 10.0,\n height: float = 10.0,\n width_segments: int = 10,\n height_segments: int = 10,\n plane: Literal[\"xz\", \"xy\", \"yx\", \"yz\", \"zx\", \"zy\"] = \"xy\",\n cell_color: RgbTupleOrArray = (200, 200, 200),\n cell_thickness: float = 1.0,\n cell_size: float = 0.5,\n section_color: RgbTupleOrArray = (140, 140, 140),\n section_thickness: float = 1.0,\n section_size: float = 1.0,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> SceneNodeHandle:\n \"\"\"Add a 2D grid to the scene.\n\n This can be useful as a size, orientation, or ground plane reference.\n\n Args:\n name: Name of the grid.\n width: Width of the grid.\n height: Height of the grid.\n width_segments: Number of segments along the width.\n height_segments: Number of segments along the height.\n plane: The plane in which the grid is oriented (e.g., 'xy', 'yz').\n cell_color: Color of the grid cells as an RGB tuple.\n cell_thickness: Thickness of the grid lines.\n cell_size: Size of each cell in the grid.\n section_color: Color of the grid sections as an RGB tuple.\n section_thickness: Thickness of the section lines.\n section_size: Size of each section in the grid.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n self._websock_interface.queue_message(\n _messages.GridMessage(\n name=name,\n width=width,\n height=height,\n width_segments=width_segments,\n height_segments=height_segments,\n plane=plane,\n cell_color=_encode_rgb(cell_color),\n cell_thickness=cell_thickness,\n cell_size=cell_size,\n section_color=_encode_rgb(section_color),\n section_thickness=section_thickness,\n section_size=section_size,\n )\n )\n return SceneNodeHandle._make(self, name, wxyz, position, visible)\n\n def add_label(\n self,\n name: str,\n text: str,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> LabelHandle:\n \"\"\"Add a 2D label to the scene.\n\n This method creates a text label in the 3D scene, which can be used to annotate\n or provide information about specific points or objects.\n\n Args:\n name: Name of the label.\n text: Text content of the label.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n self._websock_interface.queue_message(_messages.LabelMessage(name, text))\n return LabelHandle._make(self, name, wxyz, position, visible=visible)\n\n def add_point_cloud(\n self,\n name: str,\n points: onp.ndarray,\n colors: onp.ndarray | tuple[float, float, float],\n point_size: float = 0.1,\n point_shape: Literal[\n \"square\", \"diamond\", \"circle\", \"rounded\", \"sparkle\"\n ] = \"square\",\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> PointCloudHandle:\n \"\"\"Add a point cloud to the scene.\n\n Args:\n name: Name of scene node. Determines location in kinematic tree.\n points: Location of points. Should have shape (N, 3).\n colors: Colors of points. Should have shape (N, 3) or (3,).\n point_size: Size of each point.\n point_shape: Shape to draw each point.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n colors_cast = _colors_to_uint8(onp.asarray(colors))\n assert (\n len(points.shape) == 2 and points.shape[-1] == 3\n ), \"Shape of points should be (N, 3).\"\n assert colors_cast.shape in {\n points.shape,\n (3,),\n }, \"Shape of colors should be (N, 3) or (3,).\"\n\n if colors_cast.shape == (3,):\n colors_cast = onp.tile(colors_cast[None, :], reps=(points.shape[0], 1))\n\n self._websock_interface.queue_message(\n _messages.PointCloudMessage(\n name=name,\n points=points.astype(onp.float32),\n colors=colors_cast,\n point_size=point_size,\n point_ball_norm={\n \"square\": float(\"inf\"),\n \"diamond\": 1.0,\n \"circle\": 2.0,\n \"rounded\": 3.0,\n \"sparkle\": 0.6,\n }[point_shape],\n )\n )\n return PointCloudHandle._make(self, name, wxyz, position, visible)\n\n def add_mesh_skinned(\n self,\n name: str,\n vertices: onp.ndarray,\n faces: onp.ndarray,\n bone_wxyzs: tuple[tuple[float, float, float, float], ...] | onp.ndarray,\n bone_positions: tuple[tuple[float, float, float], ...] | onp.ndarray,\n skin_weights: onp.ndarray,\n color: RgbTupleOrArray = (90, 200, 255),\n wireframe: bool = False,\n opacity: float | None = None,\n material: Literal[\"standard\", \"toon3\", \"toon5\"] = \"standard\",\n flat_shading: bool = False,\n side: Literal[\"front\", \"back\", \"double\"] = \"front\",\n wxyz: Tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> MeshSkinnedHandle:\n \"\"\"Add a skinned mesh to the scene, which we can deform using a set of\n bone transformations.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n vertices: A numpy array of vertex positions. Should have shape (V, 3).\n faces: A numpy array of faces, where each face is represented by indices of\n vertices. Should have shape (F,)\n bone_wxyzs: Nested tuple or array of initial bone orientations.\n bone_positions: Nested tuple or array of initial bone positions.\n skin_weights: A numpy array of skin weights. Should have shape (V, B) where B\n is the number of bones. Only the top 4 bone weights for each\n vertex will be used.\n color: Color of the mesh as an RGB tuple.\n wireframe: Boolean indicating if the mesh should be rendered as a wireframe.\n opacity: Opacity of the mesh. None means opaque.\n material: Material type of the mesh ('standard', 'toon3', 'toon5').\n This argument is ignored when wireframe=True.\n flat_shading: Whether to do flat shading. This argument is ignored\n when wireframe=True.\n side: Side of the surface to render ('front', 'back', 'double').\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation from parent frame to local frame (t_pl).\n visible: Whether or not this mesh is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n if wireframe and material != \"standard\":\n warnings.warn(\n f\"Invalid combination of {wireframe=} and {material=}. Material argument will be ignored.\",\n stacklevel=2,\n )\n if wireframe and flat_shading:\n warnings.warn(\n f\"Invalid combination of {wireframe=} and {flat_shading=}. Flat shading argument will be ignored.\",\n stacklevel=2,\n )\n\n num_bones = len(bone_wxyzs)\n assert skin_weights.shape == (vertices.shape[0], num_bones)\n\n # Take the four biggest indices.\n top4_skin_indices = onp.argsort(skin_weights, axis=-1)[:, -4:]\n top4_skin_weights = skin_weights[\n onp.arange(vertices.shape[0])[:, None], top4_skin_indices\n ]\n assert (\n top4_skin_weights.shape == top4_skin_indices.shape == (vertices.shape[0], 4)\n )\n\n bone_wxyzs = onp.asarray(bone_wxyzs)\n bone_positions = onp.asarray(bone_positions)\n assert bone_wxyzs.shape == (num_bones, 4)\n assert bone_positions.shape == (num_bones, 3)\n self._websock_interface.queue_message(\n _messages.SkinnedMeshMessage(\n name,\n vertices.astype(onp.float32),\n faces.astype(onp.uint32),\n # (255, 255, 255) => 0xffffff, etc\n color=_encode_rgb(color),\n vertex_colors=None,\n wireframe=wireframe,\n opacity=opacity,\n flat_shading=flat_shading,\n side=side,\n material=material,\n bone_wxyzs=tuple(\n (\n float(wxyz[0]),\n float(wxyz[1]),\n float(wxyz[2]),\n float(wxyz[3]),\n )\n for wxyz in bone_wxyzs.astype(onp.float32)\n ),\n bone_positions=tuple(\n (float(xyz[0]), float(xyz[1]), float(xyz[2]))\n for xyz in bone_positions.astype(onp.float32)\n ),\n skin_indices=top4_skin_indices.astype(onp.uint16),\n skin_weights=top4_skin_weights.astype(onp.float32),\n )\n )\n handle = MeshHandle._make(self, name, wxyz, position, visible)\n return MeshSkinnedHandle(\n handle._impl,\n bones=tuple(\n MeshSkinnedBoneHandle(\n _impl=BoneState(\n name=name,\n websock_interface=self._websock_interface,\n bone_index=i,\n wxyz=bone_wxyzs[i],\n position=bone_positions[i],\n )\n )\n for i in range(num_bones)\n ),\n )\n\n def add_mesh_simple(\n self,\n name: str,\n vertices: onp.ndarray,\n faces: onp.ndarray,\n color: RgbTupleOrArray = (90, 200, 255),\n wireframe: bool = False,\n opacity: float | None = None,\n material: Literal[\"standard\", \"toon3\", \"toon5\"] = \"standard\",\n flat_shading: bool = False,\n side: Literal[\"front\", \"back\", \"double\"] = \"front\",\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> MeshHandle:\n \"\"\"Add a mesh to the scene.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n vertices: A numpy array of vertex positions. Should have shape (V, 3).\n faces: A numpy array of faces, where each face is represented by indices of\n vertices. Should have shape (F,)\n color: Color of the mesh as an RGB tuple.\n wireframe: Boolean indicating if the mesh should be rendered as a wireframe.\n opacity: Opacity of the mesh. None means opaque.\n material: Material type of the mesh ('standard', 'toon3', 'toon5').\n This argument is ignored when wireframe=True.\n flat_shading: Whether to do flat shading. This argument is ignored\n when wireframe=True.\n side: Side of the surface to render ('front', 'back', 'double').\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation from parent frame to local frame (t_pl).\n visible: Whether or not this mesh is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n if wireframe and material != \"standard\":\n warnings.warn(\n f\"Invalid combination of {wireframe=} and {material=}. Material argument will be ignored.\",\n stacklevel=2,\n )\n if wireframe and flat_shading:\n warnings.warn(\n f\"Invalid combination of {wireframe=} and {flat_shading=}. Flat shading argument will be ignored.\",\n stacklevel=2,\n )\n\n self._websock_interface.queue_message(\n _messages.MeshMessage(\n name,\n vertices.astype(onp.float32),\n faces.astype(onp.uint32),\n # (255, 255, 255) => 0xffffff, etc\n color=_encode_rgb(color),\n vertex_colors=None,\n wireframe=wireframe,\n opacity=opacity,\n flat_shading=flat_shading,\n side=side,\n material=material,\n )\n )\n return MeshHandle._make(self, name, wxyz, position, visible)\n\n def add_mesh_trimesh(\n self,\n name: str,\n mesh: trimesh.Trimesh,\n scale: float = 1.0,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> GlbHandle:\n \"\"\"Add a trimesh mesh to the scene. Internally calls `self.add_glb()`.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n mesh: A trimesh mesh object.\n scale: A scale for resizing the mesh.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n\n with io.BytesIO() as data_buffer:\n mesh.export(data_buffer, file_type=\"glb\")\n glb_data = data_buffer.getvalue()\n return self.add_glb(\n name,\n glb_data=glb_data,\n scale=scale,\n wxyz=wxyz,\n position=position,\n visible=visible,\n )\n\n def _add_gaussian_splats(\n self,\n name: str,\n centers: onp.ndarray,\n covariances: onp.ndarray,\n rgbs: onp.ndarray,\n opacities: onp.ndarray,\n wxyz: Tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> GaussianSplatHandle:\n \"\"\"Add a model to render using Gaussian Splatting.\n\n **Work-in-progress.** This feature is experimental and still under\n development. It may be changed or removed.\n\n Arguments:\n name: Scene node name.\n centers: Centers of Gaussians. (N, 3).\n covariances: Second moment for each Gaussian. (N, 3, 3).\n rgbs: Color for each Gaussian. (N, 3).\n opacities: Opacity for each Gaussian. (N, 1).\n wxyz: R_parent_local transformation.\n position: t_parent_local transformation.\n visibile: Initial visibility of scene node.\n\n Returns:\n Scene node handle.\n \"\"\"\n num_gaussians = centers.shape[0]\n assert centers.shape == (num_gaussians, 3)\n assert rgbs.shape == (num_gaussians, 3)\n assert opacities.shape == (num_gaussians, 1)\n assert covariances.shape == (num_gaussians, 3, 3)\n\n # Get cholesky factor of covariance. This helps retain precision when\n # we convert to float16.\n cov_cholesky_triu = (\n onp.linalg.cholesky(covariances.astype(onp.float64) + onp.ones(3) * 1e-7)\n .swapaxes(-1, -2) # tril => triu\n .reshape((-1, 9))[:, onp.array([0, 1, 2, 4, 5, 8])]\n )\n buffer = onp.concatenate(\n [\n # First texelFetch.\n # - xyz (96 bits): centers.\n centers.astype(onp.float32).view(onp.uint8),\n # - w (32 bits): this is reserved for use by the renderer.\n onp.zeros((num_gaussians, 4), dtype=onp.uint8),\n # Second texelFetch.\n # - xyz (96 bits): upper-triangular Cholesky factor of covariance.\n cov_cholesky_triu.astype(onp.float16).copy().view(onp.uint8),\n # - w (32 bits): rgba.\n _colors_to_uint8(rgbs),\n _colors_to_uint8(opacities),\n ],\n axis=-1,\n ).view(onp.uint32)\n assert buffer.shape == (num_gaussians, 8)\n\n self._websock_interface.queue_message(\n _messages.GaussianSplatsMessage(\n name=name,\n buffer=buffer,\n )\n )\n node_handle = GaussianSplatHandle._make(self, name, wxyz, position, visible)\n return node_handle\n\n def add_box(\n self,\n name: str,\n color: RgbTupleOrArray,\n dimensions: tuple[float, float, float] | onp.ndarray = (1.0, 1.0, 1.0),\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> MeshHandle:\n \"\"\"Add a box to the scene.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n color: Color of the box as an RGB tuple.\n dimensions: Dimensions of the box (x, y, z).\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation from parent frame to local frame (t_pl).\n visible: Whether or not this box is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n import trimesh.creation\n\n mesh = trimesh.creation.box(dimensions)\n\n return self.add_mesh_simple(\n name=name,\n vertices=mesh.vertices,\n faces=mesh.faces,\n color=color,\n flat_shading=True,\n position=position,\n wxyz=wxyz,\n visible=visible,\n )\n\n def add_icosphere(\n self,\n name: str,\n radius: float,\n color: RgbTupleOrArray,\n subdivisions: int = 3,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> MeshHandle:\n \"\"\"Add an icosphere to the scene.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n radius: Radius of the icosphere.\n color: Color of the icosphere as an RGB tuple.\n subdivisions: Number of subdivisions to use when creating the icosphere.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation from parent frame to local frame (t_pl).\n visible: Whether or not this icosphere is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n import trimesh.creation\n\n mesh = trimesh.creation.icosphere(subdivisions=subdivisions, radius=radius)\n\n # We use add_mesh_simple() because it lets us do smooth shading;\n # add_mesh_trimesh() currently does not.\n return self.add_mesh_simple(\n name=name,\n vertices=mesh.vertices,\n faces=mesh.faces,\n color=color,\n flat_shading=False,\n position=position,\n wxyz=wxyz,\n visible=visible,\n )\n\n def set_background_image(\n self,\n image: onp.ndarray,\n format: Literal[\"png\", \"jpeg\"] = \"jpeg\",\n jpeg_quality: int | None = None,\n depth: onp.ndarray | None = None,\n ) -> None:\n \"\"\"Set a background image for the scene, optionally with depth compositing.\n\n Args:\n image: The image to set as the background. Should have shape (H, W, 3).\n format: Format to transport and display the image using ('png' or 'jpeg').\n jpeg_quality: Quality of the jpeg image (if jpeg format is used).\n depth: Optional depth image to use to composite background with scene elements.\n \"\"\"\n media_type, rgb_bytes = _encode_image_binary(\n image, format, jpeg_quality=jpeg_quality\n )\n\n # Encode depth if provided. We use a 3-channel PNG to represent a fixed point\n # depth at each pixel.\n depth_bytes = None\n if depth is not None:\n # Convert to fixed-point.\n # We'll support from 0 -> (2^24 - 1) / 100_000.\n #\n # This translates to a range of [0, 167.77215], with a precision of 1e-5.\n assert len(depth.shape) == 2 or (\n len(depth.shape) == 3 and depth.shape[2] == 1\n ), \"Depth should have shape (H,W) or (H,W,1).\"\n depth = onp.clip(depth * 100_000, 0, 2**24 - 1).astype(onp.uint32)\n assert depth is not None # Appease mypy.\n intdepth: onp.ndarray = depth.reshape((*depth.shape[:2], 1)).view(onp.uint8)\n assert intdepth.shape == (*depth.shape[:2], 4)\n with io.BytesIO() as data_buffer:\n iio.imwrite(data_buffer, intdepth[:, :, :3], extension=\".png\")\n depth_bytes = data_buffer.getvalue()\n\n self._websock_interface.queue_message(\n _messages.BackgroundImageMessage(\n media_type=media_type,\n rgb_bytes=rgb_bytes,\n depth_bytes=depth_bytes,\n )\n )\n\n def add_image(\n self,\n name: str,\n image: onp.ndarray,\n render_width: float,\n render_height: float,\n format: Literal[\"png\", \"jpeg\"] = \"jpeg\",\n jpeg_quality: int | None = None,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> ImageHandle:\n \"\"\"Add a 2D image to the scene.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n image: A numpy array representing the image.\n render_width: Width at which the image should be rendered in the scene.\n render_height: Height at which the image should be rendered in the scene.\n format: Format to transport and display the image using ('png' or 'jpeg').\n jpeg_quality: Quality of the jpeg image (if jpeg format is used).\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation from parent frame to local frame (t_pl).\n visible: Whether or not this image is initially visible.\n\n Returns:\n Handle for manipulating scene node.\n \"\"\"\n\n media_type, binary = _encode_image_binary(\n image, format, jpeg_quality=jpeg_quality\n )\n self._websock_interface.queue_message(\n _messages.ImageMessage(\n name=name,\n media_type=media_type,\n data=binary,\n render_width=render_width,\n render_height=render_height,\n )\n )\n return ImageHandle._make(self, name, wxyz, position, visible)\n\n def add_transform_controls(\n self,\n name: str,\n scale: float = 1.0,\n line_width: float = 2.5,\n fixed: bool = False,\n auto_transform: bool = True,\n active_axes: tuple[bool, bool, bool] = (True, True, True),\n disable_axes: bool = False,\n disable_sliders: bool = False,\n disable_rotations: bool = False,\n translation_limits: tuple[\n tuple[float, float], tuple[float, float], tuple[float, float]\n ] = ((-1000.0, 1000.0), (-1000.0, 1000.0), (-1000.0, 1000.0)),\n rotation_limits: tuple[\n tuple[float, float], tuple[float, float], tuple[float, float]\n ] = ((-1000.0, 1000.0), (-1000.0, 1000.0), (-1000.0, 1000.0)),\n depth_test: bool = True,\n opacity: float = 1.0,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> TransformControlsHandle:\n \"\"\"Add a transform gizmo for interacting with the scene.\n\n This method adds a transform control (gizmo) to the scene, allowing for interactive\n manipulation of objects in terms of their position, rotation, and scale.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n scale: Scale of the transform controls.\n line_width: Width of the lines used in the gizmo.\n fixed: Boolean indicating if the gizmo should be fixed in position.\n auto_transform: Whether the transform should be applied automatically.\n active_axes: tuple of booleans indicating active axes.\n disable_axes: Boolean to disable axes interaction.\n disable_sliders: Boolean to disable slider interaction.\n disable_rotations: Boolean to disable rotation interaction.\n translation_limits: Limits for translation.\n rotation_limits: Limits for rotation.\n depth_test: Boolean indicating if depth testing should be used when rendering.\n opacity: Opacity of the gizmo.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation from parent frame to local frame (t_pl).\n visible: Whether or not this gizmo is initially visible.\n\n Returns:\n Handle for manipulating (and reading state of) scene node.\n \"\"\"\n self._websock_interface.queue_message(\n _messages.TransformControlsMessage(\n name=name,\n scale=scale,\n line_width=line_width,\n fixed=fixed,\n auto_transform=auto_transform,\n active_axes=active_axes,\n disable_axes=disable_axes,\n disable_sliders=disable_sliders,\n disable_rotations=disable_rotations,\n translation_limits=translation_limits,\n rotation_limits=rotation_limits,\n depth_test=depth_test,\n opacity=opacity,\n )\n )\n\n def sync_cb(client_id: ClientId, state: TransformControlsHandle) -> None:\n message_orientation = _messages.SetOrientationMessage(\n name=name,\n wxyz=tuple(map(float, state._impl.wxyz)), # type: ignore\n )\n message_orientation.excluded_self_client = client_id\n self._websock_interface.queue_message(message_orientation)\n\n message_position = _messages.SetPositionMessage(\n name=name,\n position=tuple(map(float, state._impl.position)), # type: ignore\n )\n message_position.excluded_self_client = client_id\n self._websock_interface.queue_message(message_position)\n\n node_handle = SceneNodeHandle._make(self, name, wxyz, position, visible)\n state_aux = _TransformControlsState(\n last_updated=time.time(),\n update_cb=[],\n sync_cb=sync_cb,\n )\n handle = TransformControlsHandle(node_handle._impl, state_aux)\n self._handle_from_transform_controls_name[name] = handle\n return handle\n\n def reset(self) -> None:\n \"\"\"Reset the scene.\"\"\"\n self._websock_interface.queue_message(_messages.ResetSceneMessage())\n\n def _get_client_handle(self, client_id: ClientId) -> ClientHandle:\n \"\"\"Private helper for getting a client handle from its ID.\"\"\"\n # Avoid circular imports.\n from ._viser import ViserServer\n\n # Implementation-wise, note that MessageApi is never directly instantiated.\n # Instead, it serves as a mixin/base class for either ViserServer, which\n # maintains a registry of connected clients, or ClientHandle, which should\n # only ever be dealing with its own client_id.\n if isinstance(self._owner, ViserServer):\n # TODO: there's a potential race condition here when the client disconnects.\n # This probably applies to multiple other parts of the code, we should\n # revisit all of the cases where we index into connected_clients.\n return self._owner._connected_clients[client_id]\n else:\n assert client_id == self._owner.client_id\n return self._owner\n\n def _handle_transform_controls_updates(\n self, client_id: ClientId, message: _messages.TransformControlsUpdateMessage\n ) -> None:\n \"\"\"Callback for handling transform gizmo messages.\"\"\"\n handle = self._handle_from_transform_controls_name.get(message.name, None)\n if handle is None:\n return\n\n # Update state.\n wxyz = onp.array(message.wxyz)\n position = onp.array(message.position)\n with self._owner.atomic():\n handle._impl.wxyz = wxyz\n handle._impl.position = position\n handle._impl_aux.last_updated = time.time()\n\n # Trigger callbacks.\n for cb in handle._impl_aux.update_cb:\n cb(handle)\n if handle._impl_aux.sync_cb is not None:\n handle._impl_aux.sync_cb(client_id, handle)\n\n def _handle_node_click_updates(\n self, client_id: ClientId, message: _messages.SceneNodeClickMessage\n ) -> None:\n \"\"\"Callback for handling click messages.\"\"\"\n handle = self._handle_from_node_name.get(message.name, None)\n if handle is None or handle._impl.click_cb is None:\n return\n for cb in handle._impl.click_cb:\n event = SceneNodePointerEvent(\n client=self._get_client_handle(client_id),\n client_id=client_id,\n event=\"click\",\n target=handle,\n ray_origin=message.ray_origin,\n ray_direction=message.ray_direction,\n screen_pos=message.screen_pos,\n instance_index=message.instance_index,\n )\n cb(event) # type: ignore\n\n def _handle_scene_pointer_updates(\n self, client_id: ClientId, message: _messages.ScenePointerMessage\n ):\n \"\"\"Callback for handling click messages.\"\"\"\n event = ScenePointerEvent(\n client=self._get_client_handle(client_id),\n client_id=client_id,\n event_type=message.event_type,\n ray_origin=message.ray_origin,\n ray_direction=message.ray_direction,\n screen_pos=message.screen_pos,\n )\n # Call the callback if it exists, and the after-run callback.\n if self._scene_pointer_cb is None:\n return\n self._scene_pointer_cb(event)\n\n def on_pointer_event(\n self, event_type: Literal[\"click\", \"rect-select\"]\n ) -> Callable[\n [Callable[[ScenePointerEvent], None]], Callable[[ScenePointerEvent], None]\n ]:\n \"\"\"Add a callback for scene pointer events.\n\n Args:\n event_type: event to listen to.\n \"\"\"\n # Ensure the event type is valid.\n assert event_type in get_args(_messages.ScenePointerEventType)\n\n from ._viser import ClientHandle, ViserServer\n\n def cleanup_previous_event(target: ViserServer | ClientHandle):\n # If the server or client does not have a scene pointer callback, return.\n if target.scene._scene_pointer_cb is None:\n return\n\n # Remove callback.\n target.scene.remove_pointer_callback()\n\n def decorator(\n func: Callable[[ScenePointerEvent], None],\n ) -> Callable[[ScenePointerEvent], None]:\n # Check if another scene pointer event was previously registered.\n # If so, we need to clear the previous event and register the new one.\n cleanup_previous_event(self._owner)\n\n # If called on the server handle, remove all clients' callbacks.\n if isinstance(self._owner, ViserServer):\n for client in self._owner.get_clients().values():\n cleanup_previous_event(client)\n\n # If called on the client handle, and server handle has a callback, remove the server's callback.\n # (If the server has a callback, none of the clients should have callbacks.)\n elif isinstance(self._owner, ClientHandle):\n server = self._owner._viser_server\n cleanup_previous_event(server)\n\n self._scene_pointer_cb = func\n self._scene_pointer_event_type = event_type\n\n self._websock_interface.queue_message(\n _messages.ScenePointerEnableMessage(enable=True, event_type=event_type)\n )\n return func\n\n return decorator\n\n def on_pointer_callback_removed(\n self,\n func: Callable[[], None],\n ) -> Callable[[], None]:\n \"\"\"Add a callback to run automatically when the callback for a scene\n pointer event is removed. This will be triggered exactly once, either\n manually (via :meth:`remove_pointer_callback()`) or automatically (if\n the scene pointer event is overridden with another call to\n :meth:`on_pointer_event()`).\n\n Args:\n func: Callback for when scene pointer events are removed.\n \"\"\"\n self._scene_pointer_done_cb = func\n return func\n\n def remove_pointer_callback(\n self,\n ) -> None:\n \"\"\"Remove the currently attached scene pointer event. This will trigger\n any callback attached to `.on_scene_pointer_removed()`.\"\"\"\n\n if self._scene_pointer_cb is None:\n warnings.warn(\n \"No scene pointer callback exists for this server/client, ignoring.\",\n stacklevel=2,\n )\n return\n\n # Notify client that the listener has been removed.\n event_type = self._scene_pointer_event_type\n assert event_type is not None\n self._websock_interface.queue_message(\n _messages.ScenePointerEnableMessage(enable=False, event_type=event_type)\n )\n self._owner.flush()\n\n # Run cleanup callback.\n self._scene_pointer_done_cb()\n\n # Reset the callback and event type, on the python side.\n self._scene_pointer_cb = None\n self._scene_pointer_done_cb = lambda: None\n self._scene_pointer_event_type = None\n\n def add_3d_gui_container(\n self,\n name: str,\n wxyz: tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),\n position: tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),\n visible: bool = True,\n ) -> Gui3dContainerHandle:\n \"\"\"Add a 3D gui container to the scene. The returned container handle can be\n used as a context to place GUI elements into the 3D scene.\n\n Args:\n name: A scene tree name. Names in the format of /parent/child can be used to\n define a kinematic tree.\n wxyz: Quaternion rotation to parent frame from local frame (R_pl).\n position: Translation to parent frame from local frame (t_pl).\n visible: Whether or not this scene node is initially visible.\n\n Returns:\n Handle for manipulating scene node. Can be used as a context to place GUI\n elements inside of the container.\n \"\"\"\n\n # Avoids circular import.\n from ._gui_api import _make_unique_id\n\n # New name to make the type checker happy; ViserServer and ClientHandle inherit\n # from both GuiApi and MessageApi. The pattern below is unideal.\n gui_api = self._owner.gui\n\n # Remove the 3D GUI container if it already exists. This will make sure\n # contained GUI elements are removed, preventing potential memory leaks.\n if name in self._handle_from_node_name:\n self._handle_from_node_name[name].remove()\n\n container_id = _make_unique_id()\n self._websock_interface.queue_message(\n _messages.Gui3DMessage(\n order=time.time(),\n name=name,\n container_id=container_id,\n )\n )\n node_handle = SceneNodeHandle._make(self, name, wxyz, position, visible=visible)\n return Gui3dContainerHandle(node_handle._impl, gui_api, container_id)\n" }, { "path": "viser/src/viser/_scene_handles.py", "content": "from __future__ import annotations\n\nimport dataclasses\nfrom typing import TYPE_CHECKING, Callable, Generic, Literal, TypeVar\n\nimport numpy as onp\n\nfrom . import _messages\nfrom .infra._infra import WebsockClientConnection, WebsockServer\n\nif TYPE_CHECKING:\n from ._gui_api import GuiApi\n from ._gui_handles import SupportsRemoveProtocol\n from ._scene_api import SceneApi\n from ._viser import ClientHandle\n from .infra import ClientId\n\n\n@dataclasses.dataclass(frozen=True)\nclass ScenePointerEvent:\n \"\"\"Event passed to pointer callbacks for the scene (currently only clicks).\"\"\"\n\n client: ClientHandle\n \"\"\"Client that triggered this event.\"\"\"\n client_id: int\n \"\"\"ID of client that triggered this event.\"\"\"\n event_type: _messages.ScenePointerEventType\n \"\"\"Type of event that was triggered. Currently we only support clicks and box selections.\"\"\"\n ray_origin: tuple[float, float, float] | None\n \"\"\"Origin of 3D ray corresponding to this click, in world coordinates.\"\"\"\n ray_direction: tuple[float, float, float] | None\n \"\"\"Direction of 3D ray corresponding to this click, in world coordinates.\"\"\"\n screen_pos: tuple[tuple[float, float], ...]\n \"\"\"Screen position of the click on the screen (OpenCV image coordinates, 0 to 1).\n (0, 0) is the upper-left corner, (1, 1) is the bottom-right corner.\n For a box selection, this includes the min- and max- corners of the box.\"\"\"\n\n @property\n def event(self):\n \"\"\"Deprecated. Use `event_type` instead.\"\"\"\n return self.event_type\n\n\nTSceneNodeHandle = TypeVar(\"TSceneNodeHandle\", bound=\"SceneNodeHandle\")\n\n\n@dataclasses.dataclass\nclass _SceneNodeHandleState:\n name: str\n api: SceneApi\n wxyz: onp.ndarray = dataclasses.field(\n default_factory=lambda: onp.array([1.0, 0.0, 0.0, 0.0])\n )\n position: onp.ndarray = dataclasses.field(\n default_factory=lambda: onp.array([0.0, 0.0, 0.0])\n )\n visible: bool = True\n # TODO: we should remove SceneNodeHandle as an argument here.\n click_cb: list[Callable[[SceneNodePointerEvent[SceneNodeHandle]], None]] | None = (\n None\n )\n\n\n@dataclasses.dataclass\nclass SceneNodeHandle:\n \"\"\"Handle base class for interacting with scene nodes.\"\"\"\n\n _impl: _SceneNodeHandleState\n\n @classmethod\n def _make(\n cls: type[TSceneNodeHandle],\n api: SceneApi,\n name: str,\n wxyz: tuple[float, float, float, float] | onp.ndarray,\n position: tuple[float, float, float] | onp.ndarray,\n visible: bool,\n ) -> TSceneNodeHandle:\n out = cls(_SceneNodeHandleState(name, api))\n api._handle_from_node_name[name] = out\n\n out.wxyz = wxyz\n out.position = position\n\n # Toggle visibility to make sure we send a\n # SetSceneNodeVisibilityMessage to the client.\n out._impl.visible = not visible\n out.visible = visible\n return out\n\n @property\n def wxyz(self) -> onp.ndarray:\n \"\"\"Orientation of the scene node. This is the quaternion representation of the R\n in `p_parent = [R | t] p_local`. Synchronized to clients automatically when assigned.\n \"\"\"\n return self._impl.wxyz\n\n @wxyz.setter\n def wxyz(self, wxyz: tuple[float, float, float, float] | onp.ndarray) -> None:\n from ._scene_api import cast_vector\n\n wxyz_cast = cast_vector(wxyz, 4)\n self._impl.wxyz = onp.asarray(wxyz)\n self._impl.api._websock_interface.queue_message(\n _messages.SetOrientationMessage(self._impl.name, wxyz_cast)\n )\n\n @property\n def position(self) -> onp.ndarray:\n \"\"\"Position of the scene node. This is equivalent to the t in\n `p_parent = [R | t] p_local`. Synchronized to clients automatically when assigned.\n \"\"\"\n return self._impl.position\n\n @position.setter\n def position(self, position: tuple[float, float, float] | onp.ndarray) -> None:\n from ._scene_api import cast_vector\n\n position_cast = cast_vector(position, 3)\n self._impl.position = onp.asarray(position)\n self._impl.api._websock_interface.queue_message(\n _messages.SetPositionMessage(self._impl.name, position_cast)\n )\n\n @property\n def visible(self) -> bool:\n \"\"\"Whether the scene node is visible or not. Synchronized to clients automatically when assigned.\"\"\"\n return self._impl.visible\n\n @visible.setter\n def visible(self, visible: bool) -> None:\n if visible == self._impl.visible:\n return\n self._impl.api._websock_interface.queue_message(\n _messages.SetSceneNodeVisibilityMessage(self._impl.name, visible)\n )\n self._impl.visible = visible\n\n def remove(self) -> None:\n \"\"\"Remove the node from the scene.\"\"\"\n self._impl.api._websock_interface.queue_message(\n _messages.RemoveSceneNodeMessage(self._impl.name)\n )\n\n\n@dataclasses.dataclass(frozen=True)\nclass SceneNodePointerEvent(Generic[TSceneNodeHandle]):\n \"\"\"Event passed to pointer callbacks for scene nodes (currently only clicks).\"\"\"\n\n client: ClientHandle\n \"\"\"Client that triggered this event.\"\"\"\n client_id: int\n \"\"\"ID of client that triggered this event.\"\"\"\n event: Literal[\"click\"]\n \"\"\"Type of event that was triggered. Currently we only support clicks.\"\"\"\n target: TSceneNodeHandle\n \"\"\"Scene node that was clicked.\"\"\"\n ray_origin: tuple[float, float, float]\n \"\"\"Origin of 3D ray corresponding to this click, in world coordinates.\"\"\"\n ray_direction: tuple[float, float, float]\n \"\"\"Direction of 3D ray corresponding to this click, in world coordinates.\"\"\"\n screen_pos: tuple[float, float]\n \"\"\"Screen position of the click on the screen (OpenCV image coordinates, 0 to 1).\n (0, 0) is the upper-left corner, (1, 1) is the bottom-right corner.\"\"\"\n instance_index: int | None\n \"\"\"Instance ID of the clicked object, if applicable. Currently this is `None` for all objects except for the output of :meth:`SceneApi.add_batched_axes()`.\"\"\"\n\n\n@dataclasses.dataclass\nclass _ClickableSceneNodeHandle(SceneNodeHandle):\n def on_click(\n self: TSceneNodeHandle,\n func: Callable[[SceneNodePointerEvent[TSceneNodeHandle]], None],\n ) -> Callable[[SceneNodePointerEvent[TSceneNodeHandle]], None]:\n \"\"\"Attach a callback for when a scene node is clicked.\"\"\"\n self._impl.api._websock_interface.queue_message(\n _messages.SetSceneNodeClickableMessage(self._impl.name, True)\n )\n if self._impl.click_cb is None:\n self._impl.click_cb = []\n self._impl.click_cb.append(func) # type: ignore\n return func\n\n\n@dataclasses.dataclass\nclass CameraFrustumHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for camera frustums.\"\"\"\n\n\n@dataclasses.dataclass\nclass PointCloudHandle(SceneNodeHandle):\n \"\"\"Handle for point clouds. Does not support click events.\"\"\"\n\n\n@dataclasses.dataclass\nclass BatchedAxesHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for batched coordinate frames.\"\"\"\n\n\n@dataclasses.dataclass\nclass FrameHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for coordinate frames.\"\"\"\n\n\n@dataclasses.dataclass\nclass MeshHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for mesh objects.\"\"\"\n\n\n@dataclasses.dataclass\nclass GaussianSplatHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for Gaussian splatting objects.\n\n **Work-in-progress.** Gaussian rendering is still under development.\n \"\"\"\n\n\n@dataclasses.dataclass\nclass MeshSkinnedHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for skinned mesh objects.\"\"\"\n\n bones: tuple[MeshSkinnedBoneHandle, ...]\n \"\"\"Bones of the skinned mesh. These handles can be used for reading and\n writing poses, which are defined relative to the mesh root.\"\"\"\n\n\n@dataclasses.dataclass\nclass BoneState:\n name: str\n websock_interface: WebsockServer | WebsockClientConnection\n bone_index: int\n wxyz: onp.ndarray\n position: onp.ndarray\n\n\n@dataclasses.dataclass\nclass MeshSkinnedBoneHandle:\n \"\"\"Handle for reading and writing the poses of bones in a skinned mesh.\"\"\"\n\n _impl: BoneState\n\n @property\n def wxyz(self) -> onp.ndarray:\n \"\"\"Orientation of the bone. This is the quaternion representation of the R\n in `p_parent = [R | t] p_local`. Synchronized to clients automatically when assigned.\n \"\"\"\n return self._impl.wxyz\n\n @wxyz.setter\n def wxyz(self, wxyz: tuple[float, float, float, float] | onp.ndarray) -> None:\n from ._scene_api import cast_vector\n\n wxyz_cast = cast_vector(wxyz, 4)\n self._impl.wxyz = onp.asarray(wxyz)\n self._impl.websock_interface.queue_message(\n _messages.SetBoneOrientationMessage(\n self._impl.name, self._impl.bone_index, wxyz_cast\n )\n )\n\n @property\n def position(self) -> onp.ndarray:\n \"\"\"Position of the bone. This is equivalent to the t in\n `p_parent = [R | t] p_local`. Synchronized to clients automatically when assigned.\n \"\"\"\n return self._impl.position\n\n @position.setter\n def position(self, position: tuple[float, float, float] | onp.ndarray) -> None:\n from ._scene_api import cast_vector\n\n position_cast = cast_vector(position, 3)\n self._impl.position = onp.asarray(position)\n self._impl.websock_interface.queue_message(\n _messages.SetBonePositionMessage(\n self._impl.name, self._impl.bone_index, position_cast\n )\n )\n\n\n@dataclasses.dataclass\nclass GlbHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for GLB objects.\"\"\"\n\n\n@dataclasses.dataclass\nclass ImageHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for 2D images, rendered in 3D.\"\"\"\n\n\n@dataclasses.dataclass\nclass LabelHandle(SceneNodeHandle):\n \"\"\"Handle for 2D label objects. Does not support click events.\"\"\"\n\n\n@dataclasses.dataclass\nclass _TransformControlsState:\n last_updated: float\n update_cb: list[Callable[[TransformControlsHandle], None]]\n sync_cb: None | Callable[[ClientId, TransformControlsHandle], None] = None\n\n\n@dataclasses.dataclass\nclass TransformControlsHandle(_ClickableSceneNodeHandle):\n \"\"\"Handle for interacting with transform control gizmos.\"\"\"\n\n _impl_aux: _TransformControlsState\n\n @property\n def update_timestamp(self) -> float:\n return self._impl_aux.last_updated\n\n def on_update(\n self, func: Callable[[TransformControlsHandle], None]\n ) -> Callable[[TransformControlsHandle], None]:\n \"\"\"Attach a callback for when the gizmo is moved.\"\"\"\n self._impl_aux.update_cb.append(func)\n return func\n\n\n@dataclasses.dataclass\nclass Gui3dContainerHandle(SceneNodeHandle):\n \"\"\"Use as a context to place GUI elements into a 3D GUI container.\"\"\"\n\n _gui_api: GuiApi\n _container_id: str\n _container_id_restore: str | None = None\n _children: dict[str, SupportsRemoveProtocol] = dataclasses.field(\n default_factory=dict\n )\n\n def __enter__(self) -> Gui3dContainerHandle:\n self._container_id_restore = self._gui_api._get_container_id()\n self._gui_api._set_container_id(self._container_id)\n return self\n\n def __exit__(self, *args) -> None:\n del args\n assert self._container_id_restore is not None\n self._gui_api._set_container_id(self._container_id_restore)\n self._container_id_restore = None\n\n def __post_init__(self) -> None:\n self._gui_api._container_handle_from_id[self._container_id] = self\n\n def remove(self) -> None:\n \"\"\"Permanently remove this GUI container from the visualizer.\"\"\"\n\n # Call scene node remove.\n super().remove()\n\n # Clean up contained GUI elements.\n for child in tuple(self._children.values()):\n child.remove()\n self._gui_api._container_handle_from_id.pop(self._container_id)\n" }, { "path": "viser/src/viser/_tunnel.py", "content": "from __future__ import annotations\n\nimport asyncio\nimport multiprocessing as mp\nimport threading\nfrom functools import lru_cache\nfrom multiprocessing.managers import DictProxy\nfrom pathlib import Path\nfrom typing import Callable, Literal\n\nimport rich\n\n\n@lru_cache\ndef _is_multiprocess_ok() -> bool:\n import __main__\n\n if hasattr(__main__, \"__file__\"):\n src = Path(__main__.__file__).read_text()\n return \"\\nif __name__\" in src and \"__main__\" in src\n else:\n return True\n\n\nclass ViserTunnel:\n \"\"\"Tunneling utility for internal use.\n\n This is chaotic academic software, and we'd appreciate if you refrained\n from red-teaming it. :)\n \"\"\"\n\n def __init__(self, share_domain: str, local_port: int) -> None:\n self._share_domain = share_domain\n self._local_port = local_port\n\n # Heuristic for `if __name__ == \"__main__\"` check.\n self._multiprocess_ok = _is_multiprocess_ok()\n if not self._multiprocess_ok:\n rich.print(\n \"[bold](viser)[/bold] No `if __name__ == '__main__'` check found; creating share URL tunnel in a thread\"\n )\n\n self._process: mp.Process | None = None\n self._thread: threading.Thread | None = None\n self._event_loop: asyncio.AbstractEventLoop | None = None\n\n self._shared_state: DictProxy | dict\n if self._multiprocess_ok:\n manager = mp.Manager()\n self._connect_event = manager.Event()\n self._disconnect_event = manager.Event()\n self._close_event = None # Only used for threads. For processes, we just kill the tunnel process.\n self._shared_state = manager.dict()\n else:\n self._connect_event = threading.Event()\n self._disconnect_event = threading.Event()\n self._close_event = asyncio.Event()\n self._shared_state = {}\n\n self._shared_state[\"status\"] = \"ready\"\n self._shared_state[\"url\"] = None\n\n def on_disconnect(self, callback: Callable[[], None]) -> None:\n def call_on_disconnect() -> None:\n try:\n self._disconnect_event.wait()\n except EOFError:\n return\n callback()\n\n threading.Thread(target=call_on_disconnect, daemon=True).start()\n\n def on_connect(self, callback: Callable[[int], None]) -> None:\n \"\"\"Establish the tunnel connection.\n\n Returns URL if tunnel succeeds, otherwise None.\"\"\"\n assert self._process is None\n\n self._shared_state[\"status\"] = \"connecting\"\n\n def wait_job() -> None:\n try:\n self._connect_event.wait()\n except EOFError:\n return\n callback(self._shared_state[\"max_conn_count\"])\n\n threading.Thread(target=wait_job, daemon=True).start()\n\n # Note that this will generally require an __name__ == \"__main__\" check\n # on the origin script.\n if self._multiprocess_ok:\n self._process = mp.Process(\n target=_connect_job,\n daemon=True,\n args=(\n self._connect_event,\n self._disconnect_event,\n self._close_event,\n self._share_domain,\n self._local_port,\n self._shared_state,\n None,\n ),\n )\n self._process.start()\n else:\n self._thread = threading.Thread(\n target=_connect_job,\n daemon=True,\n args=(\n self._connect_event,\n self._disconnect_event,\n self._close_event,\n self._share_domain,\n self._local_port,\n self._shared_state,\n self,\n ),\n )\n self._thread.start()\n\n def get_url(self) -> str | None:\n \"\"\"Get tunnel URL. None if not connected (or connection failed).\"\"\"\n return self._shared_state[\"url\"]\n\n def get_status(\n self,\n ) -> Literal[\"ready\", \"connecting\", \"failed\", \"connected\", \"closed\"]:\n return self._shared_state[\"status\"]\n\n def close(self) -> None:\n \"\"\"Close the tunnel.\"\"\"\n if self._process is not None:\n self._process.kill()\n self._process.join()\n self._disconnect_event.set()\n if self._thread is not None:\n assert self._event_loop is not None\n\n @self._event_loop.call_soon_threadsafe\n def _() -> None:\n assert self._close_event is not None\n self._close_event.set()\n\n self._thread.join()\n self._disconnect_event.set()\n\n\ndef _connect_job(\n connect_event: threading.Event,\n disconnect_event: threading.Event,\n close_event: asyncio.Event | None, # Only for threads.\n share_domain: str,\n local_port: int,\n shared_state: DictProxy | dict,\n event_loop_target: ViserTunnel | None, # Only for threads.\n) -> None:\n event_loop = asyncio.new_event_loop()\n asyncio.set_event_loop(event_loop)\n if event_loop_target is not None:\n event_loop_target._event_loop = event_loop\n if close_event is None:\n close_event = asyncio.Event()\n\n try:\n event_loop.run_until_complete(\n _make_tunnel(\n connect_event,\n disconnect_event,\n close_event,\n share_domain,\n local_port,\n shared_state,\n )\n )\n event_loop.close()\n except KeyboardInterrupt:\n event_loop.call_soon_threadsafe(close_event.set)\n tasks = asyncio.all_tasks(event_loop)\n event_loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))\n event_loop.close()\n\n\nasync def _make_tunnel(\n connect_event: threading.Event,\n disconnect_event: threading.Event,\n close_event: asyncio.Event | None,\n share_domain: str,\n local_port: int,\n shared_state: DictProxy | dict,\n) -> None:\n share_domain = \"share.viser.studio\"\n\n import requests\n\n try:\n response = requests.request(\n \"GET\",\n url=f\"https://{share_domain}/?request_forward\",\n headers={\"Content-Type\": \"application/json\"},\n )\n if response.status_code != 200:\n shared_state[\"status\"] = \"failed\"\n return\n except requests.exceptions.ConnectionError:\n shared_state[\"status\"] = \"failed\"\n return\n except Exception as e:\n shared_state[\"status\"] = \"failed\"\n raise e\n\n res = response.json()\n shared_state[\"url\"] = res[\"url\"]\n shared_state[\"max_conn_count\"] = res[\"max_conn_count\"]\n shared_state[\"status\"] = \"connected\"\n connect_event.set()\n\n await asyncio.gather(\n *[\n asyncio.create_task(\n _simple_proxy(\n \"127.0.0.1\",\n local_port,\n share_domain,\n res[\"port\"],\n close_event if close_event is not None else asyncio.Event(),\n )\n )\n for _ in range(res[\"max_conn_count\"])\n ]\n )\n\n shared_state[\"url\"] = None\n shared_state[\"status\"] = \"closed\"\n disconnect_event.set()\n\n\nasync def _simple_proxy(\n local_host: str,\n local_port: int,\n remote_host: str,\n remote_port: int,\n close_event: asyncio.Event,\n) -> None:\n \"\"\"Establish a connection to the tunnel server.\"\"\"\n\n async def close_writer(writer: asyncio.StreamWriter) -> None:\n \"\"\"Utility for closing a writer and waiting until done, while suppressing errors\n from broken connections.\"\"\"\n try:\n if not writer.is_closing():\n writer.close()\n await writer.wait_closed()\n except ConnectionError:\n pass\n\n async def relay(r: asyncio.StreamReader, w: asyncio.StreamWriter) -> None:\n \"\"\"Simple data passthrough from one stream to another.\"\"\"\n try:\n while True:\n data = await r.read(4096)\n if len(data) == 0:\n # Done!\n break\n w.write(data)\n await w.drain()\n except Exception:\n pass\n finally:\n await close_writer(w)\n\n while True:\n local_w = None\n remote_w = None\n try:\n local_r, local_w = await asyncio.open_connection(local_host, local_port)\n remote_r, remote_w = await asyncio.open_connection(remote_host, remote_port)\n await asyncio.wait(\n [\n asyncio.gather(\n asyncio.create_task(relay(local_r, remote_w)),\n asyncio.create_task(relay(remote_r, local_w)),\n ),\n asyncio.create_task(close_event.wait()),\n ],\n return_when=asyncio.FIRST_COMPLETED,\n )\n except Exception:\n pass\n finally:\n # Be extra sure that connections are closed.\n if local_w is not None:\n await close_writer(local_w)\n if remote_w is not None:\n await close_writer(remote_w)\n\n if close_event.is_set():\n break\n\n # Throttle connection attempts.\n await asyncio.sleep(0.1)\n" }, { "path": "viser/src/viser/_viser.py", "content": "from __future__ import annotations\n\nimport dataclasses\nimport io\nimport mimetypes\nimport threading\nimport time\nimport warnings\nfrom pathlib import Path\nfrom typing import TYPE_CHECKING, Any, Callable, ContextManager\n\nimport imageio.v3 as iio\nimport numpy as onp\nimport numpy.typing as npt\nimport rich\nfrom rich import box, style\nfrom rich.panel import Panel\nfrom rich.table import Table\nfrom typing_extensions import Literal\n\nfrom . import _client_autobuild, _messages, infra\nfrom . import transforms as tf\nfrom ._gui_api import Color, GuiApi, _make_unique_id\nfrom ._notification_handle import NotificationHandle, _NotificationHandleState\nfrom ._scene_api import SceneApi, cast_vector\nfrom ._tunnel import ViserTunnel\nfrom .infra._infra import RecordHandle\n\n\nclass _BackwardsCompatibilityShim:\n \"\"\"Shims for backward compatibility with viser API from version\n `<=0.1.30`.\"\"\"\n\n def __getattr__(self, name: str) -> Any:\n fixed_name = {\n # Map from old method names (viser v0.1.*) to new methods names.\n \"reset_scene\": \"reset\",\n \"set_global_scene_node_visibility\": \"set_global_visibility\",\n \"on_scene_pointer\": \"on_pointer_event\",\n \"on_scene_pointer_removed\": \"on_pointer_callback_removed\",\n \"remove_scene_pointer_callback\": \"remove_pointer_callback\",\n \"add_mesh\": \"add_mesh_simple\",\n }.get(name, name)\n if hasattr(self.scene, fixed_name):\n warnings.warn(\n f\"{type(self).__name__}.{name} has been deprecated, use {type(self).__name__}.scene.{fixed_name} instead. Alternatively, pin to `viser<0.2.0`.\",\n category=DeprecationWarning,\n stacklevel=2,\n )\n return object.__getattribute__(self.scene, fixed_name)\n\n fixed_name = name.replace(\"add_gui_\", \"add_\").replace(\"set_gui_\", \"set_\")\n if hasattr(self.gui, fixed_name):\n warnings.warn(\n f\"{type(self).__name__}.{name} has been deprecated, use {type(self).__name__}.gui.{fixed_name} instead. Alternatively, pin to `viser<0.2.0`.\",\n category=DeprecationWarning,\n stacklevel=2,\n )\n return object.__getattribute__(self.gui, fixed_name)\n\n raise AttributeError(\n f\"'{type(self).__name__}' object has no attribute '{name}'\"\n )\n\n\n@dataclasses.dataclass\nclass _CameraHandleState:\n \"\"\"Information about a client's camera state.\"\"\"\n\n client: ClientHandle\n wxyz: npt.NDArray[onp.float64]\n position: npt.NDArray[onp.float64]\n fov: float\n aspect: float\n look_at: npt.NDArray[onp.float64]\n up_direction: npt.NDArray[onp.float64]\n update_timestamp: float\n camera_cb: list[Callable[[CameraHandle], None]]\n\n\nclass CameraHandle:\n \"\"\"A handle for reading and writing the camera state of a particular\n client. Typically accessed via :attr:`ClientHandle.camera`.\"\"\"\n\n def __init__(self, client: ClientHandle) -> None:\n self._state = _CameraHandleState(\n client,\n wxyz=onp.zeros(4),\n position=onp.zeros(3),\n fov=0.0,\n aspect=0.0,\n look_at=onp.zeros(3),\n up_direction=onp.zeros(3),\n update_timestamp=0.0,\n camera_cb=[],\n )\n\n @property\n def client(self) -> ClientHandle:\n \"\"\"Client that this camera corresponds to.\"\"\"\n return self._state.client\n\n @property\n def wxyz(self) -> npt.NDArray[onp.float64]:\n \"\"\"Corresponds to the R in `P_world = [R | t] p_camera`. Synchronized\n automatically when assigned.\"\"\"\n assert self._state.update_timestamp != 0.0\n return self._state.wxyz\n\n # Note: asymmetric properties are supported in Pyright, but not yet in mypy.\n # - https://github.com/python/mypy/issues/3004\n # - https://github.com/python/mypy/pull/11643\n @wxyz.setter\n def wxyz(self, wxyz: tuple[float, float, float, float] | onp.ndarray) -> None:\n R_world_camera = tf.SO3(onp.asarray(wxyz)).as_matrix()\n look_distance = onp.linalg.norm(self.look_at - self.position)\n\n # We're following OpenCV conventions: look_direction is +Z, up_direction is -Y,\n # right_direction is +X.\n look_direction = R_world_camera[:, 2]\n up_direction = -R_world_camera[:, 1]\n right_direction = R_world_camera[:, 0]\n\n # Minimize our impact on the orbit controls by keeping the new up direction as\n # close to the old one as possible.\n projected_up_direction = (\n self.up_direction\n - float(self.up_direction @ right_direction) * right_direction\n )\n up_cosine = float(up_direction @ projected_up_direction)\n if abs(up_cosine) < 0.05:\n projected_up_direction = up_direction\n elif up_cosine < 0.0:\n projected_up_direction = up_direction\n\n new_look_at = look_direction * look_distance + self.position\n\n # Update lookat and up direction.\n self.look_at = new_look_at\n self.up_direction = projected_up_direction\n\n # The internal camera orientation should be set in the look_at /\n # up_direction setters. We can uncomment this assert to check this.\n # assert onp.allclose(self._state.wxyz, wxyz) or onp.allclose(\n # self._state.wxyz, -wxyz\n # )\n\n @property\n def position(self) -> npt.NDArray[onp.float64]:\n \"\"\"Corresponds to the t in `P_world = [R | t] p_camera`. Synchronized\n automatically when assigned.\n\n The `look_at` point and `up_direction` vectors are maintained when updating\n `position`, which means that updates to `position` will often also affect `wxyz`.\n \"\"\"\n assert self._state.update_timestamp != 0.0\n return self._state.position\n\n @position.setter\n def position(self, position: tuple[float, float, float] | onp.ndarray) -> None:\n offset = onp.asarray(position) - onp.array(self.position) # type: ignore\n self._state.position = onp.asarray(position)\n self.look_at = onp.array(self.look_at) + offset\n self._state.update_timestamp = time.time()\n self._state.client._websock_connection.queue_message(\n _messages.SetCameraPositionMessage(cast_vector(position, 3))\n )\n\n def _update_wxyz(self) -> None:\n \"\"\"Compute and update the camera orientation from the internal look_at, position, and up vectors.\"\"\"\n z = self._state.look_at - self._state.position\n z /= onp.linalg.norm(z)\n y = tf.SO3.exp(z * onp.pi) @ self._state.up_direction\n y = y - onp.dot(z, y) * z\n y /= onp.linalg.norm(y)\n x = onp.cross(y, z)\n self._state.wxyz = tf.SO3.from_matrix(onp.stack([x, y, z], axis=1)).wxyz\n\n @property\n def fov(self) -> float:\n \"\"\"Vertical field of view of the camera, in radians. Synchronized automatically\n when assigned.\"\"\"\n assert self._state.update_timestamp != 0.0\n return self._state.fov\n\n @fov.setter\n def fov(self, fov: float) -> None:\n self._state.fov = fov\n self._state.update_timestamp = time.time()\n self._state.client._websock_connection.queue_message(\n _messages.SetCameraFovMessage(fov)\n )\n\n @property\n def aspect(self) -> float:\n \"\"\"Canvas width divided by height. Not assignable.\"\"\"\n assert self._state.update_timestamp != 0.0\n return self._state.aspect\n\n @property\n def update_timestamp(self) -> float:\n assert self._state.update_timestamp != 0.0\n return self._state.update_timestamp\n\n @property\n def look_at(self) -> npt.NDArray[onp.float64]:\n \"\"\"Look at point for the camera. Synchronized automatically when set.\"\"\"\n assert self._state.update_timestamp != 0.0\n return self._state.look_at\n\n @look_at.setter\n def look_at(self, look_at: tuple[float, float, float] | onp.ndarray) -> None:\n self._state.look_at = onp.asarray(look_at)\n self._state.update_timestamp = time.time()\n self._update_wxyz()\n self._state.client._websock_connection.queue_message(\n _messages.SetCameraLookAtMessage(cast_vector(look_at, 3))\n )\n\n @property\n def up_direction(self) -> npt.NDArray[onp.float64]:\n \"\"\"Up direction for the camera. Synchronized automatically when set.\"\"\"\n assert self._state.update_timestamp != 0.0\n return self._state.up_direction\n\n @up_direction.setter\n def up_direction(\n self, up_direction: tuple[float, float, float] | onp.ndarray\n ) -> None:\n self._state.up_direction = onp.asarray(up_direction)\n self._update_wxyz()\n self._state.update_timestamp = time.time()\n self._state.client._websock_connection.queue_message(\n _messages.SetCameraUpDirectionMessage(cast_vector(up_direction, 3))\n )\n\n def on_update(\n self, callback: Callable[[CameraHandle], None]\n ) -> Callable[[CameraHandle], None]:\n \"\"\"Attach a callback to run when a new camera message is received.\"\"\"\n self._state.camera_cb.append(callback)\n return callback\n\n def get_render(\n self, height: int, width: int, transport_format: Literal[\"png\", \"jpeg\"] = \"jpeg\"\n ) -> onp.ndarray:\n \"\"\"Request a render from a client, block until it's done and received, then\n return it as a numpy array.\n\n Args:\n height: Height of rendered image. Should be <= the browser height.\n width: Width of rendered image. Should be <= the browser width.\n transport_format: Image transport format. JPEG will return a lossy (H, W, 3) RGB array. PNG will\n return a lossless (H, W, 4) RGBA array, but can cause memory issues on the frontend if called\n too quickly for higher-resolution images.\n \"\"\"\n\n # Listen for a render reseponse message, which should contain the rendered\n # image.\n render_ready_event = threading.Event()\n out: onp.ndarray | None = None\n\n connection = self.client._websock_connection\n\n def got_render_cb(\n client_id: int, message: _messages.GetRenderResponseMessage\n ) -> None:\n del client_id\n connection.unregister_handler(\n _messages.GetRenderResponseMessage, got_render_cb\n )\n nonlocal out\n out = iio.imread(\n io.BytesIO(message.payload),\n extension=f\".{transport_format}\",\n )\n render_ready_event.set()\n\n connection.register_handler(_messages.GetRenderResponseMessage, got_render_cb)\n self.client._websock_connection.queue_message(\n _messages.GetRenderRequestMessage(\n \"image/jpeg\" if transport_format == \"jpeg\" else \"image/png\",\n height=height,\n width=width,\n # Only used for JPEG. The main reason to use a lower quality version\n # value is (unfortunately) to make life easier for the Javascript\n # garbage collector.\n quality=80,\n )\n )\n render_ready_event.wait()\n assert out is not None\n return out\n\n\n# Don't inherit from _BackwardsCompatibilityShim during type checking, because\n# this will unnecessarily suppress type errors. (from the overriding of\n# __getattr__).\nclass ClientHandle(_BackwardsCompatibilityShim if not TYPE_CHECKING else object):\n \"\"\"A handle is created for each client that connects to a server. Handles can be\n used to communicate with just one client, as well as for reading and writing of\n camera state.\n\n Similar to :class:`ViserServer`, client handles also expose scene and GUI\n interfaces at :attr:`ClientHandle.scene` and :attr:`ClientHandle.gui`. If\n these are used, for example via a client's\n :meth:`SceneApi.add_point_cloud()` method, created elements are local to\n only one specific client.\n \"\"\"\n\n def __init__(\n self, conn: infra.WebsockClientConnection, server: ViserServer\n ) -> None:\n # Private attributes.\n self._websock_connection = conn\n self._viser_server = server\n\n # Public attributes.\n self.scene: SceneApi = SceneApi(\n self, thread_executor=server._websock_server._thread_executor\n )\n \"\"\"Handle for interacting with the 3D scene.\"\"\"\n self.gui: GuiApi = GuiApi(\n self, thread_executor=server._websock_server._thread_executor\n )\n \"\"\"Handle for interacting with the GUI.\"\"\"\n self.client_id: int = conn.client_id\n \"\"\"Unique ID for this client.\"\"\"\n self.camera: CameraHandle = CameraHandle(self)\n \"\"\"Handle for reading from and manipulating the client's viewport camera.\"\"\"\n\n def flush(self) -> None:\n \"\"\"Flush the outgoing message buffer. Any buffered messages will immediately be\n sent. (by default they are windowed)\"\"\"\n self._viser_server._websock_server.flush_client(self.client_id)\n\n def atomic(self) -> ContextManager[None]:\n \"\"\"Returns a context where: all outgoing messages are grouped and applied by\n clients atomically.\n\n This should be treated as a soft constraint that's helpful for things\n like animations, or when we want position and orientation updates to\n happen synchronously.\n\n Returns:\n Context manager.\n \"\"\"\n return self._websock_connection.atomic()\n\n def send_file_download(\n self, filename: str, content: bytes, chunk_size: int = 1024 * 1024\n ) -> None:\n \"\"\"Send a file for a client or clients to download.\n\n Args:\n filename: Name of the file to send. Used to infer MIME type.\n content: Content of the file.\n chunk_size: Number of bytes to send at a time.\n \"\"\"\n mime_type = mimetypes.guess_type(filename, strict=False)[0]\n if mime_type is None:\n mime_type = \"application/octet-stream\"\n\n parts = [\n content[i * chunk_size : (i + 1) * chunk_size]\n for i in range(int(onp.ceil(len(content) / chunk_size)))\n ]\n\n uuid = _make_unique_id()\n self._websock_connection.queue_message(\n _messages.FileTransferStart(\n source_component_id=None,\n transfer_uuid=uuid,\n filename=filename,\n mime_type=mime_type,\n part_count=len(parts),\n size_bytes=len(content),\n )\n )\n for i, part in enumerate(parts):\n self._websock_connection.queue_message(\n _messages.FileTransferPart(\n None,\n transfer_uuid=uuid,\n part=i,\n content=part,\n )\n )\n self.flush()\n\n def add_notification(\n self,\n title: str,\n body: str,\n loading: bool = False,\n with_close_button: bool = True,\n auto_close: int | Literal[False] = False,\n color: Color | None = None,\n ) -> NotificationHandle:\n \"\"\"Add a notification to the client's interface.\n\n This method creates a new notification that will be displayed at the\n top left corner of the client's viewer. Notifications are useful for\n providing alerts or status updates to users.\n\n Args:\n title: Title to display on the notification.\n body: Message to display on the notification body.\n loading: Whether the notification shows loading icon.\n with_close_button: Whether the notification can be manually closed.\n auto_close: Time in ms before the notification automatically closes;\n otherwise False such that the notification never closes on its own.\n\n Returns:\n A handle that can be used to interact with the GUI element.\n \"\"\"\n handle = NotificationHandle(\n _NotificationHandleState(\n websock_interface=self._websock_connection,\n id=_make_unique_id(),\n title=title,\n body=body,\n loading=loading,\n with_close_button=with_close_button,\n auto_close=auto_close,\n color=color,\n )\n )\n handle._sync_with_client(first=True)\n return handle\n\n\nclass ViserServer(_BackwardsCompatibilityShim if not TYPE_CHECKING else object):\n \"\"\":class:`ViserServer` is the main class for working with viser. On\n instantiation, it (a) launches a thread with a web server and (b) provides\n a high-level API for interactive 3D visualization.\n\n **Core API.** Clients can connect via a web browser, and will be shown two\n components: a 3D scene and a 2D GUI panel. Methods belonging to\n :attr:`ViserServer.scene` can be used to add 3D primitives to the scene.\n Methods belonging to :attr:`ViserServer.gui` can be used to add 2D GUI\n elements.\n\n **Shared state.** Elements added to the server object, for example via a\n server's :meth:`SceneApi.add_point_cloud` or :meth:`GuiApi.add_button`,\n will have state that's shared and synchronized automatically between all\n connected clients. To show elements that are local to a single client, see\n :attr:`ClientHandle.scene` and :attr:`ClientHandle.gui`.\n\n Args:\n host: Host to bind server to.\n port: Port to bind server to.\n label: Label shown at the top of the GUI panel.\n \"\"\"\n\n # Hide deprecated arguments from docstring and type checkers.\n def __init__(\n self,\n host: str = \"0.0.0.0\",\n port: int = 8080,\n label: str | None = None,\n verbose: bool = True,\n **_deprecated_kwargs,\n ):\n # Create server.\n server = infra.WebsockServer(\n host=host,\n port=port,\n message_class=_messages.Message,\n http_server_root=Path(__file__).absolute().parent / \"client\" / \"build\",\n verbose=verbose,\n client_api_version=1,\n )\n self._websock_server = server\n\n _client_autobuild.ensure_client_is_built()\n\n self._connection = server\n self._connected_clients: dict[int, ClientHandle] = {}\n self._client_lock = threading.Lock()\n self._client_connect_cb: list[Callable[[ClientHandle], None]] = []\n self._client_disconnect_cb: list[Callable[[ClientHandle], None]] = []\n\n # For new clients, register and add a handler for camera messages.\n @server.on_client_connect\n def _(conn: infra.WebsockClientConnection) -> None:\n client = ClientHandle(conn, server=self)\n first = True\n\n def handle_camera_message(\n client_id: infra.ClientId, message: _messages.ViewerCameraMessage\n ) -> None:\n nonlocal first\n\n assert client_id == client.client_id\n\n # Update the client's camera.\n with client.atomic():\n client.camera._state = _CameraHandleState(\n client,\n onp.array(message.wxyz),\n onp.array(message.position),\n message.fov,\n message.aspect,\n onp.array(message.look_at),\n onp.array(message.up_direction),\n time.time(),\n camera_cb=client.camera._state.camera_cb,\n )\n\n # We consider a client to be connected after the first camera message is\n # received.\n if first:\n first = False\n with self._client_lock:\n self._connected_clients[conn.client_id] = client\n for cb in self._client_connect_cb:\n cb(client)\n\n for camera_cb in client.camera._state.camera_cb:\n camera_cb(client.camera)\n\n conn.register_handler(_messages.ViewerCameraMessage, handle_camera_message)\n\n # Remove clients when they disconnect.\n @server.on_client_disconnect\n def _(conn: infra.WebsockClientConnection) -> None:\n with self._client_lock:\n if conn.client_id not in self._connected_clients:\n return\n\n handle = self._connected_clients.pop(conn.client_id)\n for cb in self._client_disconnect_cb:\n cb(handle)\n\n # Start the server.\n server.start()\n\n self.scene: SceneApi = SceneApi(self, thread_executor=server._thread_executor)\n \"\"\"Handle for interacting with the 3D scene.\"\"\"\n\n self.gui: GuiApi = GuiApi(self, thread_executor=server._thread_executor)\n \"\"\"Handle for interacting with the GUI.\"\"\"\n\n server.register_handler(\n _messages.ShareUrlDisconnect,\n lambda client_id, msg: self.disconnect_share_url(),\n )\n server.register_handler(\n _messages.ShareUrlRequest, lambda client_id, msg: self.request_share_url()\n )\n\n # Form status print.\n port = server._port # Port may have changed.\n http_url = f\"http://{host}:{port}\"\n ws_url = f\"ws://{host}:{port}\"\n table = Table(\n title=None,\n show_header=False,\n box=box.MINIMAL,\n title_style=style.Style(bold=True),\n )\n table.add_row(\"HTTP\", http_url)\n table.add_row(\"Websocket\", ws_url)\n rich.print(Panel(table, title=\"[bold]viser[/bold]\", expand=False))\n\n self._share_tunnel: ViserTunnel | None = None\n\n # Create share tunnel if requested.\n # This is deprecated: we should use get_share_url() instead.\n share = _deprecated_kwargs.get(\"share\", False)\n if share:\n self.request_share_url()\n\n self.scene.reset()\n self.gui.reset()\n self.gui.set_panel_label(label)\n\n def get_host(self) -> str:\n \"\"\"Returns the host address of the Viser server.\n\n Returns:\n Host address as string.\n \"\"\"\n return self._websock_server._host\n\n def get_port(self) -> int:\n \"\"\"Returns the port of the Viser server. This could be different from the\n originally requested one.\n\n Returns:\n Port as integer.\n \"\"\"\n return self._websock_server._port\n\n def request_share_url(self, verbose: bool = True) -> str | None:\n \"\"\"Request a share URL for the Viser server, which allows for public access.\n On the first call, will block until a connecting with the share URL server is\n established. Afterwards, the URL will be returned directly.\n\n This is an experimental feature that relies on an external server; it shouldn't\n be relied on for critical applications.\n\n Returns:\n Share URL as string, or None if connection fails or is closed.\n \"\"\"\n if self._share_tunnel is not None:\n # Tunnel already exists.\n while self._share_tunnel.get_status() in (\"ready\", \"connecting\"):\n time.sleep(0.05)\n return self._share_tunnel.get_url()\n else:\n # Create a new tunnel!.\n if verbose:\n rich.print(\"[bold](viser)[/bold] Share URL requested!\")\n\n connect_event = threading.Event()\n\n self._share_tunnel = ViserTunnel(\n \"share.viser.studio\", self._websock_server._port\n )\n\n @self._share_tunnel.on_disconnect\n def _() -> None:\n rich.print(\"[bold](viser)[/bold] Disconnected from share URL\")\n self._share_tunnel = None\n self._websock_server.unsafe_send_message(\n _messages.ShareUrlUpdated(None)\n )\n\n @self._share_tunnel.on_connect\n def _(max_clients: int) -> None:\n assert self._share_tunnel is not None\n share_url = self._share_tunnel.get_url()\n if verbose:\n if share_url is None:\n rich.print(\"[bold](viser)[/bold] Could not generate share URL\")\n else:\n rich.print(\n f\"[bold](viser)[/bold] Generated share URL (expires in 24 hours, max {max_clients} clients): {share_url}\"\n )\n self._websock_server.unsafe_send_message(\n _messages.ShareUrlUpdated(share_url)\n )\n connect_event.set()\n\n connect_event.wait()\n\n url = self._share_tunnel.get_url()\n return url\n\n def disconnect_share_url(self) -> None:\n \"\"\"Disconnect from the share URL server.\"\"\"\n if self._share_tunnel is not None:\n self._share_tunnel.close()\n else:\n rich.print(\n \"[bold](viser)[/bold] Tried to disconnect from share URL, but already disconnected\"\n )\n\n def stop(self) -> None:\n \"\"\"Stop the Viser server and associated threads and tunnels.\"\"\"\n self._websock_server.stop()\n if self._share_tunnel is not None:\n self._share_tunnel.close()\n\n def get_clients(self) -> dict[int, ClientHandle]:\n \"\"\"Creates and returns a copy of the mapping from connected client IDs to\n handles.\n\n Returns:\n Dictionary of clients.\n \"\"\"\n with self._client_lock:\n return self._connected_clients.copy()\n\n def on_client_connect(\n self, cb: Callable[[ClientHandle], None]\n ) -> Callable[[ClientHandle], None]:\n \"\"\"Attach a callback to run for newly connected clients.\"\"\"\n with self._client_lock:\n clients = self._connected_clients.copy().values()\n self._client_connect_cb.append(cb)\n\n # Trigger callback on any already-connected clients.\n # If we have:\n #\n # server = viser.ViserServer()\n # server.on_client_connect(...)\n #\n # This makes sure that the the callback is applied to any clients that\n # connect between the two lines.\n for client in clients:\n cb(client)\n return cb\n\n def on_client_disconnect(\n self, cb: Callable[[ClientHandle], None]\n ) -> Callable[[ClientHandle], None]:\n \"\"\"Attach a callback to run when clients disconnect.\"\"\"\n self._client_disconnect_cb.append(cb)\n return cb\n\n def flush(self) -> None:\n \"\"\"Flush the outgoing message buffer. Any buffered messages will immediately be\n sent. (by default they are windowed)\"\"\"\n self._websock_server.flush()\n\n def atomic(self) -> ContextManager[None]:\n \"\"\"Returns a context where: all outgoing messages are grouped and applied by\n clients atomically.\n\n This should be treated as a soft constraint that's helpful for things\n like animations, or when we want position and orientation updates to\n happen synchronously.\n\n Returns:\n Context manager.\n \"\"\"\n return self._websock_server.atomic()\n\n def send_file_download(\n self, filename: str, content: bytes, chunk_size: int = 1024 * 1024\n ) -> None:\n \"\"\"Send a file for a client or clients to download.\n\n Args:\n filename: Name of the file to send. Used to infer MIME type.\n content: Content of the file.\n chunk_size: Number of bytes to send at a time.\n \"\"\"\n for client in self.get_clients().values():\n client.send_file_download(filename, content, chunk_size)\n\n def _start_scene_recording(self) -> RecordHandle:\n \"\"\"Start recording outgoing messages for playback or\n embedding. Includes only the scene.\n\n **Work-in-progress.** This API may be changed or removed.\n \"\"\"\n recorder = self._websock_server.start_recording(\n # Don't record GUI messages. This feels brittle.\n filter=lambda message: \"Gui\" not in type(message).__name__\n )\n # Insert current scene state.\n for message in self._websock_server._broadcast_buffer.message_from_id.values():\n recorder._insert_message(message)\n return recorder\n" }, { "path": "viser/src/viser/client/.eslintrc.js", "content": "module.exports = {\n settings: {\n react: {\n version: \"detect\", // React version. \"detect\" automatically picks the version you have installed.\n // You can also use `16.0`, `16.3`, etc, if you want to override the detected value.\n // It will default to \"latest\" and warn if missing, and to \"detect\" in the future\n },\n },\n env: {\n browser: true,\n es2021: true,\n },\n extends: [\n \"eslint:recommended\",\n \"plugin:react/recommended\",\n \"plugin:@typescript-eslint/recommended\",\n ],\n overrides: [],\n parser: \"@typescript-eslint/parser\",\n parserOptions: {\n ecmaVersion: \"latest\",\n sourceType: \"module\",\n },\n plugins: [\"react\", \"@typescript-eslint\", \"react-refresh\"],\n ignorePatterns: [\"build/\", \".eslintrc.js\"],\n rules: {\n // https://github.com/jsx-eslint/eslint-plugin-react/issues/3423\n \"react/no-unknown-property\": \"off\",\n \"no-constant-condition\": \"off\",\n // Suppress errors for missing 'import React' in files.\n \"react/react-in-jsx-scope\": \"off\",\n \"@typescript-eslint/ban-ts-comment\": \"off\",\n \"@typescript-eslint/no-explicit-any\": \"off\",\n \"@typescript-eslint/no-non-null-assertion\": \"off\",\n \"react/prop-types\": [\n \"error\",\n {\n skipUndeclared: true,\n },\n ],\n \"react-refresh/only-export-components\": \"warn\",\n },\n};\n" }, { "path": "viser/src/viser/client/.gitignore", "content": "# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.\n\n# dependencies\n/node_modules\n/.pnp\n.pnp.js\n\n# testing\n/coverage\n\n# misc\n.DS_Store\n.env.local\n.env.development.local\n.env.test.local\n.env.production.local\n\nnpm-debug.log*\nyarn-debug.log*\nyarn-error.log*\n" }, { "path": "viser/src/viser/client/index.html", "content": "\n\n \n \n \n \n \n \n \n \n \n \n \n Viser\n \n \n \n
\n \n \n\n" }, { "path": "viser/src/viser/client/package.json", "content": "{\n \"name\": \"viser\",\n \"version\": \"0.1.0\",\n \"private\": true,\n \"dependencies\": {\n \"@mantine/core\": \"^7.6.2\",\n \"@mantine/dates\": \"^7.6.2\",\n \"@mantine/hooks\": \"^7.6.2\",\n \"@mantine/notifications\": \"^7.6.2\",\n \"@mantine/vanilla-extract\": \"^7.6.2\",\n \"@mdx-js/mdx\": \"^3.0.1\",\n \"@mdx-js/react\": \"^3.0.1\",\n \"@msgpack/msgpack\": \"^3.0.0-beta2\",\n \"@react-three/drei\": \"^9.64.0\",\n \"@react-three/fiber\": \"^8.12.0\",\n \"@tabler/icons-react\": \"^3.1.0\",\n \"@types/node\": \"^20.11.30\",\n \"@types/react\": \"^18.0.33\",\n \"@types/react-dom\": \"^18.0.11\",\n \"@types/three\": \"^0.162.0\",\n \"@vanilla-extract/css\": \"^1.14.1\",\n \"@vitejs/plugin-react\": \"^4.0.1\",\n \"await-lock\": \"^2.2.2\",\n \"clsx\": \"^2.1.0\",\n \"colortranslator\": \"^4.1.0\",\n \"dayjs\": \"^1.11.10\",\n \"detect-browser\": \"^5.3.0\",\n \"fflate\": \"^0.8.2\",\n \"hold-event\": \"^1.1.0\",\n \"immer\": \"^10.0.4\",\n \"its-fine\": \"^1.2.5\",\n \"mantine-react-table\": \"^2.0.0-beta.0\",\n \"postcss\": \"^8.4.38\",\n \"prettier\": \"^3.0.3\",\n \"react\": \"^18.2.0\",\n \"react-dom\": \"^18.2.0\",\n \"react-error-boundary\": \"^4.0.10\",\n \"react-qr-code\": \"^2.0.12\",\n \"react-router-dom\": \"^6.10.0\",\n \"rehype-color-chips\": \"^0.1.3\",\n \"remark-gfm\": \"^4.0.0\",\n \"three\": \"^0.162.0\",\n \"vite\": \"^5.2.6\",\n \"vite-plugin-svgr\": \"^4.2.0\",\n \"vite-tsconfig-paths\": \"^4.2.0\",\n \"web-vitals\": \"^3.3.1\",\n \"ws\": \"^8.13.0\",\n \"zustand\": \"^4.3.7\"\n },\n \"scripts\": {\n \"start\": \"vite --host\",\n \"build\": \"tsc && vite build\",\n \"serve\": \"vite preview\"\n },\n \"eslintConfig\": {\n \"extends\": [\n \"react-app\"\n ]\n },\n \"browserslist\": {\n \"production\": [\n \"last 2 chrome versions\",\n \"last 2 firefox versions\",\n \"last 2 safari versions\"\n ],\n \"development\": [\n \"last 1 chrome version\",\n \"last 1 firefox version\",\n \"last 1 safari version\"\n ]\n },\n \"devDependencies\": {\n \"@types/msgpack\": \"0.0.34\",\n \"@types/uuid\": \"^9.0.8\",\n \"@types/wicg-file-system-access\": \"^2023.10.5\",\n \"@typescript-eslint/eslint-plugin\": \"^7.4.0\",\n \"@typescript-eslint/parser\": \"^7.4.0\",\n \"@typescript-eslint/typescript-estree\": \"^7.4.0\",\n \"@vanilla-extract/vite-plugin\": \"^4.0.6\",\n \"browserslist-to-esbuild\": \"^2.1.1\",\n \"eslint\": \"^8.43.0\",\n \"eslint-plugin-react\": \"^7.32.2\",\n \"eslint-plugin-react-refresh\": \"^0.4.1\",\n \"postcss-preset-mantine\": \"^1.13.0\",\n \"typescript\": \"^5.0.4\",\n \"vite-plugin-eslint\": \"^1.8.1\"\n }\n}\n" }, { "path": "viser/src/viser/client/postcss.config.cjs", "content": "module.exports = {\n plugins: {\n \"postcss-preset-mantine\": {},\n \"postcss-simple-vars\": {\n variables: {\n \"mantine-breakpoint-xs\": \"36em\",\n \"mantine-breakpoint-sm\": \"48em\",\n \"mantine-breakpoint-md\": \"62em\",\n \"mantine-breakpoint-lg\": \"75em\",\n \"mantine-breakpoint-xl\": \"88em\",\n },\n },\n },\n};\n" }, { "path": "viser/src/viser/client/public/manifest.json", "content": "{\n \"short_name\": \"Viser\",\n \"name\": \"Viser\",\n \"icons\": [\n {\n \"src\": \"favicon.svg\",\n \"sizes\": \"any\",\n \"type\": \"image/x-icon\"\n }\n ],\n \"start_url\": \".\",\n \"display\": \"standalone\",\n \"theme_color\": \"#000000\",\n \"background_color\": \"#ffffff\"\n}\n" }, { "path": "viser/src/viser/client/src/App.css.ts", "content": "import { globalStyle } from \"@vanilla-extract/css\";\n\nglobalStyle(\".mantine-ScrollArea-scrollbar\", {\n zIndex: 100,\n});\n" }, { "path": "viser/src/viser/client/src/App.tsx", "content": "// @refresh reset\nimport \"@mantine/core/styles.css\";\nimport \"@mantine/notifications/styles.css\";\nimport \"./App.css\";\n\nimport { Notifications } from \"@mantine/notifications\";\n\nimport {\n CameraControls,\n Environment,\n PerformanceMonitor,\n Stats,\n} from \"@react-three/drei\";\nimport * as THREE from \"three\";\nimport { Canvas, useThree, useFrame } from \"@react-three/fiber\";\n\nimport { SynchronizedCameraControls } from \"./CameraControls\";\nimport {\n Anchor,\n Box,\n ColorSchemeScript,\n Image,\n MantineProvider,\n Modal,\n Tooltip,\n createTheme,\n useMantineTheme,\n} from \"@mantine/core\";\nimport React, { useEffect } from \"react\";\nimport { SceneNodeThreeObject, UseSceneTree } from \"./SceneTree\";\n\nimport \"./index.css\";\n\nimport ControlPanel from \"./ControlPanel/ControlPanel\";\nimport { UseGui, useGuiState } from \"./ControlPanel/GuiState\";\nimport { searchParamKey } from \"./SearchParamsUtils\";\nimport { WebsocketMessageProducer } from \"./WebsocketInterface\";\n\nimport { Titlebar } from \"./Titlebar\";\nimport { ViserModal } from \"./Modal\";\nimport { useSceneTreeState } from \"./SceneTreeState\";\nimport { GetRenderRequestMessage, Message } from \"./WebsocketMessages\";\nimport { useThrottledMessageSender } from \"./WebsocketFunctions\";\nimport { useDisclosure } from \"@mantine/hooks\";\nimport { rayToViserCoords } from \"./WorldTransformUtils\";\nimport { ndcFromPointerXy, opencvXyFromPointerXy } from \"./ClickUtils\";\nimport { theme } from \"./AppTheme\";\nimport { FrameSynchronizedMessageHandler } from \"./MessageHandler\";\nimport { PlaybackFromFile } from \"./FilePlayback\";\nimport { SplatRenderContext } from \"./Splatting/GaussianSplats\";\nimport { BrowserWarning } from \"./BrowserWarning\";\n\nexport type ViewerContextContents = {\n messageSource: \"websocket\" | \"file_playback\";\n // Zustand hooks.\n useSceneTree: UseSceneTree;\n useGui: UseGui;\n // Useful references.\n // TODO: there's really no reason these all need to be their own ref objects.\n // We could have just one ref to a global mutable struct.\n sendMessageRef: React.MutableRefObject<(message: Message) => void>;\n canvasRef: React.MutableRefObject;\n sceneRef: React.MutableRefObject;\n cameraRef: React.MutableRefObject;\n backgroundMaterialRef: React.MutableRefObject;\n cameraControlRef: React.MutableRefObject;\n sendCameraRef: React.MutableRefObject<(() => void) | null>;\n resetCameraViewRef: React.MutableRefObject<(() => void) | null>;\n // Scene node attributes.\n // This is intentionally placed outside of the Zustand state to reduce overhead.\n nodeAttributesFromName: React.MutableRefObject<{\n [name: string]:\n | undefined\n | {\n poseUpdateState?: \"updated\" | \"needsUpdate\" | \"waitForMakeObject\";\n wxyz?: [number, number, number, number];\n position?: [number, number, number];\n visibility?: boolean; // Visibility state from the server.\n overrideVisibility?: boolean; // Override from the GUI.\n };\n }>;\n nodeRefFromName: React.MutableRefObject<{\n [name: string]: undefined | THREE.Object3D;\n }>;\n messageQueueRef: React.MutableRefObject;\n // Requested a render.\n getRenderRequestState: React.MutableRefObject<\n \"ready\" | \"triggered\" | \"pause\" | \"in_progress\"\n >;\n getRenderRequest: React.MutableRefObject;\n // Track click drag events.\n scenePointerInfo: React.MutableRefObject<{\n enabled: false | \"click\" | \"rect-select\"; // Enable box events.\n dragStart: [number, number]; // First mouse position.\n dragEnd: [number, number]; // Final mouse position.\n isDragging: boolean;\n }>;\n // 2D canvas for drawing -- can be used to give feedback on cursor movement, or more.\n canvas2dRef: React.MutableRefObject;\n // Poses for bones in skinned meshes.\n skinnedMeshState: React.MutableRefObject<{\n [name: string]: {\n initialized: boolean;\n poses: {\n wxyz: [number, number, number, number];\n position: [number, number, number];\n }[];\n };\n }>;\n};\nexport const ViewerContext = React.createContext(\n null,\n);\n\nTHREE.ColorManagement.enabled = true;\n\nfunction ViewerRoot() {\n // What websocket server should we connect to?\n function getDefaultServerFromUrl() {\n // https://localhost:8080/ => ws://localhost:8080\n // https://localhost:8080/?server=some_url => ws://localhost:8080\n let server = window.location.href;\n server = server.replace(\"http://\", \"ws://\");\n server = server.replace(\"https://\", \"wss://\");\n server = server.split(\"?\")[0];\n if (server.endsWith(\"/\")) server = server.slice(0, -1);\n return server;\n }\n const servers = new URLSearchParams(window.location.search).getAll(\n searchParamKey,\n );\n const initialServer =\n servers.length >= 1 ? servers[0] : getDefaultServerFromUrl();\n\n // Playback mode for embedding viser.\n const searchParams = new URLSearchParams(window.location.search);\n const playbackPath = searchParams.get(\"playbackPath\");\n const darkMode = searchParams.get(\"darkMode\") !== null;\n const showStats = searchParams.get(\"showStats\") !== null;\n\n // Values that can be globally accessed by components in a viewer.\n const nodeRefFromName = React.useRef<{\n [name: string]: undefined | THREE.Object3D;\n }>({});\n const viewer: ViewerContextContents = {\n messageSource: playbackPath === null ? \"websocket\" : \"file_playback\",\n useSceneTree: useSceneTreeState(nodeRefFromName),\n useGui: useGuiState(initialServer),\n sendMessageRef: React.useRef(\n playbackPath == null\n ? (message) =>\n console.log(\n `Tried to send ${message.type} but websocket is not connected!`,\n )\n : () => null,\n ),\n canvasRef: React.useRef(null),\n sceneRef: React.useRef(null),\n cameraRef: React.useRef(null),\n backgroundMaterialRef: React.useRef(null),\n cameraControlRef: React.useRef(null),\n sendCameraRef: React.useRef(null),\n resetCameraViewRef: React.useRef(null),\n // Scene node attributes that aren't placed in the zustand state for performance reasons.\n nodeAttributesFromName: React.useRef({\n \"\": {\n wxyz: (() => {\n const quat = new THREE.Quaternion().setFromEuler(\n new THREE.Euler(Math.PI / 2, Math.PI, -Math.PI / 2),\n );\n return [quat.w, quat.x, quat.y, quat.z];\n })(),\n },\n }),\n nodeRefFromName: nodeRefFromName,\n messageQueueRef: React.useRef([]),\n getRenderRequestState: React.useRef(\"ready\"),\n getRenderRequest: React.useRef(null),\n scenePointerInfo: React.useRef({\n enabled: false,\n dragStart: [0, 0],\n dragEnd: [0, 0],\n isDragging: false,\n }),\n canvas2dRef: React.useRef(null),\n skinnedMeshState: React.useRef({}),\n };\n\n // Set dark default if specified in URL.\n if (darkMode) viewer.useGui.getState().theme.dark_mode = darkMode;\n\n return (\n \n \n {viewer.messageSource === \"websocket\" ? (\n \n ) : null}\n {viewer.messageSource === \"file_playback\" ? (\n \n ) : null}\n {showStats ? : null}\n \n \n );\n}\n\nfunction ViewerContents({ children }: { children: React.ReactNode }) {\n const viewer = React.useContext(ViewerContext)!;\n const darkMode = viewer.useGui((state) => state.theme.dark_mode);\n const colors = viewer.useGui((state) => state.theme.colors);\n const controlLayout = viewer.useGui((state) => state.theme.control_layout);\n return (\n <>\n \n \n {children}\n \n \n \n \n \n \n ({\n backgroundColor: darkMode ? theme.colors.dark[9] : \"#fff\",\n flexGrow: 1,\n overflow: \"hidden\",\n height: \"100%\",\n })}\n >\n \n \n \n \n {viewer.useGui((state) => state.theme.show_logo) &&\n viewer.messageSource == \"websocket\" ? (\n \n ) : null}\n \n {viewer.messageSource == \"websocket\" ? (\n \n ) : null}\n \n \n \n \n );\n}\n\nfunction ViewerCanvas({ children }: { children: React.ReactNode }) {\n const viewer = React.useContext(ViewerContext)!;\n const sendClickThrottled = useThrottledMessageSender(20);\n const theme = useMantineTheme();\n\n const initDistanceScale = parseFloat(\n new URLSearchParams(window.location.search).get(\"initDistanceScale\") ??\n \"1.0\",\n );\n\n return (\n {\n const pointerInfo = viewer.scenePointerInfo.current!;\n\n // Only handle pointer events if enabled.\n if (pointerInfo.enabled === false) return;\n\n // Keep track of the first click position.\n const canvasBbox = viewer.canvasRef.current!.getBoundingClientRect();\n pointerInfo.dragStart = [\n e.clientX - canvasBbox.left,\n e.clientY - canvasBbox.top,\n ];\n pointerInfo.dragEnd = pointerInfo.dragStart;\n\n // Check if pointer position is in bounds.\n if (ndcFromPointerXy(viewer, pointerInfo.dragEnd) === null) return;\n\n // Only allow one drag event at a time.\n if (pointerInfo.isDragging) return;\n pointerInfo.isDragging = true;\n\n // Disable camera controls -- we don't want the camera to move while we're dragging.\n viewer.cameraControlRef.current!.enabled = false;\n\n const ctx = viewer.canvas2dRef.current!.getContext(\"2d\")!;\n ctx.clearRect(0, 0, ctx.canvas.width, ctx.canvas.height);\n }}\n onPointerMove={(e) => {\n const pointerInfo = viewer.scenePointerInfo.current!;\n\n // Only handle if click events are enabled, and if pointer is down (i.e., dragging).\n if (pointerInfo.enabled === false || !pointerInfo.isDragging) return;\n\n // Check if pointer position is in boudns.\n const canvasBbox = viewer.canvasRef.current!.getBoundingClientRect();\n const pointerXy: [number, number] = [\n e.clientX - canvasBbox.left,\n e.clientY - canvasBbox.top,\n ];\n if (ndcFromPointerXy(viewer, pointerXy) === null) return;\n\n // Check if mouse position has changed sufficiently from last position.\n // Uses 3px as a threshood, similar to drag detection in\n // `SceneNodeClickMessage` from `SceneTree.tsx`.\n pointerInfo.dragEnd = pointerXy;\n if (\n Math.abs(pointerInfo.dragEnd[0] - pointerInfo.dragStart[0]) <= 3 &&\n Math.abs(pointerInfo.dragEnd[1] - pointerInfo.dragStart[1]) <= 3\n )\n return;\n\n // If we're listening for scene box events, draw the box on the 2D canvas for user feedback.\n if (pointerInfo.enabled === \"rect-select\") {\n const ctx = viewer.canvas2dRef.current!.getContext(\"2d\")!;\n ctx.clearRect(0, 0, ctx.canvas.width, ctx.canvas.height);\n ctx.beginPath();\n ctx.fillStyle = theme.primaryColor;\n ctx.strokeStyle = \"blue\";\n ctx.globalAlpha = 0.2;\n ctx.fillRect(\n pointerInfo.dragStart[0],\n pointerInfo.dragStart[1],\n pointerInfo.dragEnd[0] - pointerInfo.dragStart[0],\n pointerInfo.dragEnd[1] - pointerInfo.dragStart[1],\n );\n ctx.globalAlpha = 1.0;\n ctx.stroke();\n }\n }}\n onPointerUp={() => {\n const pointerInfo = viewer.scenePointerInfo.current!;\n\n // Re-enable camera controls! Was disabled in `onPointerDown`, to allow\n // for mouse drag w/o camera movement.\n viewer.cameraControlRef.current!.enabled = true;\n\n // Only handle if click events are enabled, and if pointer was down (i.e., dragging).\n if (pointerInfo.enabled === false || !pointerInfo.isDragging) return;\n\n const ctx = viewer.canvas2dRef.current!.getContext(\"2d\")!;\n ctx.clearRect(0, 0, ctx.canvas.width, ctx.canvas.height);\n\n // If there's only one pointer, send a click message.\n // The message will return origin/direction lists of length 1.\n if (pointerInfo.enabled === \"click\") {\n const raycaster = new THREE.Raycaster();\n\n // Raycaster expects NDC coordinates, so we convert the click event to NDC.\n const mouseVector = ndcFromPointerXy(viewer, pointerInfo.dragEnd);\n if (mouseVector === null) return;\n raycaster.setFromCamera(mouseVector, viewer.cameraRef.current!);\n const ray = rayToViserCoords(viewer, raycaster.ray);\n\n // Send OpenCV image coordinates to the server (normalized).\n const mouseVectorOpenCV = opencvXyFromPointerXy(\n viewer,\n pointerInfo.dragEnd,\n );\n\n sendClickThrottled({\n type: \"ScenePointerMessage\",\n event_type: \"click\",\n ray_origin: [ray.origin.x, ray.origin.y, ray.origin.z],\n ray_direction: [ray.direction.x, ray.direction.y, ray.direction.z],\n screen_pos: [[mouseVectorOpenCV.x, mouseVectorOpenCV.y]],\n });\n } else if (pointerInfo.enabled === \"rect-select\") {\n // If the ScenePointerEvent had mouse drag movement, we will send a \"box\" message:\n // Use the first and last mouse positions to create a box.\n // Again, click should be in openCV image coordinates (normalized).\n const firstMouseVector = opencvXyFromPointerXy(\n viewer,\n pointerInfo.dragStart,\n );\n const lastMouseVector = opencvXyFromPointerXy(\n viewer,\n pointerInfo.dragEnd,\n );\n\n const x_min = Math.min(firstMouseVector.x, lastMouseVector.x);\n const x_max = Math.max(firstMouseVector.x, lastMouseVector.x);\n const y_min = Math.min(firstMouseVector.y, lastMouseVector.y);\n const y_max = Math.max(firstMouseVector.y, lastMouseVector.y);\n\n // Send the upper-left and lower-right corners of the box.\n const screenBoxList: [number, number][] = [\n [x_min, y_min],\n [x_max, y_max],\n ];\n\n sendClickThrottled({\n type: \"ScenePointerMessage\",\n event_type: \"rect-select\",\n ray_origin: null,\n ray_direction: null,\n screen_pos: screenBoxList,\n });\n }\n\n // Release drag lock.\n pointerInfo.isDragging = false;\n }}\n >\n {children}\n \n \n \n \n \n \n \n \n \n \n \n );\n}\n\nfunction AdaptiveDpr() {\n const setDpr = useThree((state) => state.setDpr);\n return (\n {\n const max = Math.min(refreshrate * 0.9, 85);\n const min = Math.max(max * 0.5, 38);\n return [min, max];\n }}\n onChange={({ factor, fps, refreshrate }) => {\n const dpr = window.devicePixelRatio * (0.2 + 0.8 * factor);\n console.log(\n `[Performance] Setting DPR to ${dpr}; FPS=${fps}/${refreshrate}`,\n );\n setDpr(dpr);\n }}\n />\n );\n}\n\n/* HTML Canvas, for drawing 2D. */\nfunction Viewer2DCanvas() {\n const viewer = React.useContext(ViewerContext)!;\n useEffect(() => {\n // Create a resize observer to resize the CSS canvas when the window is resized.\n const resizeObserver = new ResizeObserver((entries) => {\n const { width, height } = entries[0].contentRect;\n canvas.width = width;\n canvas.height = height;\n });\n\n // Observe the canvas.\n const canvas = viewer.canvas2dRef.current!;\n resizeObserver.observe(canvas);\n\n // Cleanup\n return () => resizeObserver.disconnect();\n });\n return (\n \n );\n}\n\n/* Background image with support for depth compositing. */\nfunction BackgroundImage() {\n // Create a fragment shader that composites depth using depth and rgb\n const vertShader = `\n varying vec2 vUv;\n\n void main() {\n vUv = uv;\n gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);\n }\n `.trim();\n const fragShader = `\n #include \n precision highp float;\n precision highp int;\n\n varying vec2 vUv;\n uniform sampler2D colorMap;\n uniform sampler2D depthMap;\n uniform float cameraNear;\n uniform float cameraFar;\n uniform bool enabled;\n uniform bool hasDepth;\n\n float readDepth(sampler2D depthMap, vec2 coord) {\n vec4 rgbPacked = texture(depthMap, coord);\n\n // For the k-th channel, coefficients are calculated as: 255 * 1e-5 * 2^(8 * k).\n // Note that: [0, 255] channels are scaled to [0, 1], and we multiply by 1e5 on the server side.\n float depth = rgbPacked.r * 0.00255 + rgbPacked.g * 0.6528 + rgbPacked.b * 167.1168;\n return depth;\n }\n\n void main() {\n if (!enabled) {\n // discard the pixel if we're not enabled\n discard;\n }\n vec4 color = texture(colorMap, vUv);\n gl_FragColor = vec4(color.rgb, 1.0);\n\n float bufDepth;\n if(hasDepth){\n float depth = readDepth(depthMap, vUv);\n bufDepth = viewZToPerspectiveDepth(-depth, cameraNear, cameraFar);\n } else {\n // If no depth enabled, set depth to 1.0 (infinity) to treat it like a background image.\n bufDepth = 1.0;\n }\n gl_FragDepth = bufDepth;\n }`.trim();\n // initialize the rgb texture with all white and depth at infinity\n const backgroundMaterial = new THREE.ShaderMaterial({\n fragmentShader: fragShader,\n vertexShader: vertShader,\n uniforms: {\n enabled: { value: false },\n depthMap: { value: null },\n colorMap: { value: null },\n cameraNear: { value: null },\n cameraFar: { value: null },\n hasDepth: { value: false },\n },\n });\n const { backgroundMaterialRef } = React.useContext(ViewerContext)!;\n backgroundMaterialRef.current = backgroundMaterial;\n const backgroundMesh = React.useRef(null);\n useFrame(({ camera }) => {\n // Logic ahead relies on perspective camera assumption.\n if (!(camera instanceof THREE.PerspectiveCamera)) {\n console.error(\n \"Camera is not a perspective camera, cannot render background image\",\n );\n return;\n }\n\n // Update the position of the mesh based on the camera position.\n const lookdir = camera.getWorldDirection(new THREE.Vector3());\n backgroundMesh.current!.position.set(\n camera.position.x,\n camera.position.y,\n camera.position.z,\n );\n backgroundMesh.current!.position.addScaledVector(lookdir, 1.0);\n backgroundMesh.current!.quaternion.copy(camera.quaternion);\n\n // Resize the mesh based on focal length.\n const f = camera.getFocalLength();\n backgroundMesh.current!.scale.set(\n camera.getFilmWidth() / f,\n camera.getFilmHeight() / f,\n 1.0,\n );\n\n // Set near/far uniforms.\n backgroundMaterial.uniforms.cameraNear.value = camera.near;\n backgroundMaterial.uniforms.cameraFar.value = camera.far;\n });\n\n return (\n \n \n \n );\n}\n\n/** Component for helping us set the scene reference. */\nfunction SceneContextSetter() {\n const { sceneRef, cameraRef } = React.useContext(ViewerContext)!;\n sceneRef.current = useThree((state) => state.scene);\n cameraRef.current = useThree(\n (state) => state.camera as THREE.PerspectiveCamera,\n );\n return <>;\n}\n\nexport function Root() {\n return (\n \n \n \n );\n}\n\n/** Logo. When clicked, opens an info modal. */\nfunction ViserLogo() {\n const [aboutModalOpened, { open: openAbout, close: closeAbout }] =\n useDisclosure(false);\n return (\n <>\n \n \n \n \n \n \n \n

Viser is a 3D visualization toolkit developed at UC Berkeley.

\n

\n \n Nerfstudio\n \n   •  \n \n GitHub\n \n   •  \n \n Documentation\n \n

\n \n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/AppTheme.ts", "content": "import {\n Checkbox,\n ColorInput,\n Select,\n TextInput,\n NumberInput,\n Paper,\n ActionIcon,\n Button,\n createTheme,\n} from \"@mantine/core\";\nimport { themeToVars } from \"@mantine/vanilla-extract\";\n\nexport const theme = createTheme({\n fontFamily: \"Inter\",\n autoContrast: true,\n components: {\n Checkbox: Checkbox.extend({\n defaultProps: {\n radius: \"xs\",\n },\n }),\n ColorInput: ColorInput.extend({\n defaultProps: {\n radius: \"xs\",\n },\n }),\n Select: Select.extend({\n defaultProps: {\n radius: \"sm\",\n },\n }),\n TextInput: TextInput.extend({\n defaultProps: {\n radius: \"xs\",\n },\n }),\n NumberInput: NumberInput.extend({\n defaultProps: {\n radius: \"xs\",\n },\n }),\n Paper: Paper.extend({\n defaultProps: {\n radius: \"xs\",\n shadow: \"0\",\n },\n }),\n ActionIcon: ActionIcon.extend({\n defaultProps: {\n variant: \"subtle\",\n color: \"gray\",\n radius: \"xs\",\n },\n }),\n Button: Button.extend({\n defaultProps: {\n radius: \"xs\",\n fw: 450,\n },\n }),\n },\n});\n\nexport const vars = themeToVars(theme);\n" }, { "path": "viser/src/viser/client/src/BrowserWarning.tsx", "content": "import { notifications } from \"@mantine/notifications\";\nimport { detect } from \"detect-browser\";\nimport { useEffect } from \"react\";\n\nexport function BrowserWarning() {\n useEffect(() => {\n const browser = detect();\n\n // Browser version are based loosely on support for SIMD, OffscreenCanvas.\n //\n // https://caniuse.com/?search=simd\n // https://caniuse.com/?search=OffscreenCanvas\n if (browser === null || browser.version === null) {\n console.log(\"Failed to detect browser\");\n notifications.show({\n title: \"Could not detect browser version\",\n message:\n \"Your browser version could not be detected. It may not be supported.\",\n autoClose: false,\n color: \"red\",\n });\n } else {\n const version = parseFloat(browser.version);\n console.log(`Detected ${browser.name} version ${version}`);\n if (\n (browser.name === \"chrome\" && version < 91) ||\n (browser.name === \"edge\" && version < 91) ||\n (browser.name === \"firefox\" && version < 89) ||\n (browser.name === \"opera\" && version < 77) ||\n (browser.name === \"safari\" && version < 16.4)\n )\n notifications.show({\n title: \"Unsuppported browser\",\n message: `Your browser (${\n browser.name.slice(0, 1).toUpperCase() + browser.name.slice(1)\n }/${\n browser.version\n }) is outdated, which may cause problems. Consider updating.`,\n autoClose: false,\n color: \"red\",\n });\n }\n });\n return null;\n}\n" }, { "path": "viser/src/viser/client/src/CameraControls.tsx", "content": "import { ViewerContext } from \"./App\";\nimport { CameraControls } from \"@react-three/drei\";\nimport { useThree } from \"@react-three/fiber\";\nimport * as holdEvent from \"hold-event\";\nimport React, { useContext, useRef } from \"react\";\nimport { PerspectiveCamera } from \"three\";\nimport * as THREE from \"three\";\nimport { computeT_threeworld_world } from \"./WorldTransformUtils\";\nimport { useThrottledMessageSender } from \"./WebsocketFunctions\";\n\nexport function SynchronizedCameraControls() {\n const viewer = useContext(ViewerContext)!;\n const camera = useThree((state) => state.camera as PerspectiveCamera);\n\n const sendCameraThrottled = useThrottledMessageSender(20);\n\n // Helper for resetting camera poses.\n const initialCameraRef = useRef<{\n camera: PerspectiveCamera;\n lookAt: THREE.Vector3;\n } | null>(null);\n\n viewer.resetCameraViewRef.current = () => {\n viewer.cameraControlRef.current!.setLookAt(\n initialCameraRef.current!.camera.position.x,\n initialCameraRef.current!.camera.position.y,\n initialCameraRef.current!.camera.position.z,\n initialCameraRef.current!.lookAt.x,\n initialCameraRef.current!.lookAt.y,\n initialCameraRef.current!.lookAt.z,\n true,\n );\n viewer.cameraRef.current!.up.set(\n initialCameraRef.current!.camera.up.x,\n initialCameraRef.current!.camera.up.y,\n initialCameraRef.current!.camera.up.z,\n );\n viewer.cameraControlRef.current!.updateCameraUp();\n };\n\n // Callback for sending cameras.\n // It makes the code more chaotic, but we preallocate a bunch of things to\n // minimize garbage collection!\n const R_threecam_cam = new THREE.Quaternion().setFromEuler(\n new THREE.Euler(Math.PI, 0.0, 0.0),\n );\n const R_world_threeworld = new THREE.Quaternion();\n const tmpMatrix4 = new THREE.Matrix4();\n const lookAt = new THREE.Vector3();\n const R_world_camera = new THREE.Quaternion();\n const t_world_camera = new THREE.Vector3();\n const scale = new THREE.Vector3();\n const sendCamera = React.useCallback(() => {\n const three_camera = camera;\n const camera_control = viewer.cameraControlRef.current;\n\n if (camera_control === null) {\n // Camera controls not yet ready, let's re-try later.\n setTimeout(sendCamera, 10);\n return;\n }\n\n // We put Z up to match the scene tree, and convert threejs camera convention\n // to the OpenCV one.\n const T_world_threeworld = computeT_threeworld_world(viewer).invert();\n const T_world_camera = T_world_threeworld.clone()\n .multiply(\n tmpMatrix4\n .makeRotationFromQuaternion(three_camera.quaternion)\n .setPosition(three_camera.position),\n )\n .multiply(tmpMatrix4.makeRotationFromQuaternion(R_threecam_cam));\n R_world_threeworld.setFromRotationMatrix(T_world_threeworld);\n\n camera_control.getTarget(lookAt).applyQuaternion(R_world_threeworld);\n const up = three_camera.up.clone().applyQuaternion(R_world_threeworld);\n\n //Store initial camera values\n if (initialCameraRef.current === null) {\n initialCameraRef.current = {\n camera: three_camera.clone(),\n lookAt: camera_control.getTarget(new THREE.Vector3()),\n };\n }\n\n T_world_camera.decompose(t_world_camera, R_world_camera, scale);\n\n sendCameraThrottled({\n type: \"ViewerCameraMessage\",\n wxyz: [\n R_world_camera.w,\n R_world_camera.x,\n R_world_camera.y,\n R_world_camera.z,\n ],\n position: t_world_camera.toArray(),\n aspect: three_camera.aspect,\n fov: (three_camera.fov * Math.PI) / 180.0,\n look_at: [lookAt.x, lookAt.y, lookAt.z],\n up_direction: [up.x, up.y, up.z],\n });\n }, [camera, sendCameraThrottled]);\n\n // Send camera for new connections.\n // We add a small delay to give the server time to add a callback.\n const connected = viewer.useGui((state) => state.websocketConnected);\n\n const initHeightOffset = parseFloat(\n new URLSearchParams(window.location.search).get(\"initHeightOffset\") ??\n \"0.05\",\n );\n // Add the height offset to the initial camera position\n React.useEffect(() => {\n const cameraControls = viewer.cameraControlRef.current!;\n const lookAt = new THREE.Vector3();\n cameraControls.getTarget(lookAt);\n\n viewer.cameraControlRef.current!.setLookAt(\n camera.position.x,\n camera.position.y + (camera.position.y / 0.1) * initHeightOffset, // init_scale: camera.position.y / 0.1 \n camera.position.z,\n lookAt.x,\n lookAt.y + (camera.position.y / 0.1) * initHeightOffset,\n lookAt.z,\n true,\n );\n viewer.sendCameraRef.current = sendCamera;\n if (!connected) return;\n setTimeout(() => sendCamera(), 50);\n }, [connected, sendCamera]);\n\n // Send camera for 3D viewport changes.\n const canvas = viewer.canvasRef.current!; // R3F canvas.\n React.useEffect(() => {\n // Create a resize observer to resize the CSS canvas when the window is resized.\n const resizeObserver = new ResizeObserver(() => {\n sendCamera();\n });\n resizeObserver.observe(canvas);\n\n // Cleanup.\n return () => resizeObserver.disconnect();\n }, [canvas]);\n\n // Keyboard controls.\n React.useEffect(() => {\n const cameraControls = viewer.cameraControlRef.current!;\n\n const wKey = new holdEvent.KeyboardKeyHold(\"KeyW\", 20);\n const aKey = new holdEvent.KeyboardKeyHold(\"KeyA\", 20);\n const sKey = new holdEvent.KeyboardKeyHold(\"KeyS\", 20);\n const dKey = new holdEvent.KeyboardKeyHold(\"KeyD\", 20);\n const qKey = new holdEvent.KeyboardKeyHold(\"KeyQ\", 20);\n const eKey = new holdEvent.KeyboardKeyHold(\"KeyE\", 20);\n\n // TODO: these event listeners are currently never removed, even if this\n // component gets unmounted.\n aKey.addEventListener(\"holding\", (event) => {\n cameraControls.truck(-0.0002 * event?.deltaTime, 0, true);\n });\n dKey.addEventListener(\"holding\", (event) => {\n cameraControls.truck(0.0002 * event?.deltaTime, 0, true);\n });\n wKey.addEventListener(\"holding\", (event) => {\n cameraControls.forward(0.0002 * event?.deltaTime, true);\n });\n sKey.addEventListener(\"holding\", (event) => {\n cameraControls.forward(-0.0002 * event?.deltaTime, true);\n });\n qKey.addEventListener(\"holding\", (event) => {\n cameraControls.elevate(0.0002 * event?.deltaTime, true);\n });\n eKey.addEventListener(\"holding\", (event) => {\n cameraControls.elevate(-0.0002 * event?.deltaTime, true);\n });\n\n const leftKey = new holdEvent.KeyboardKeyHold(\"ArrowLeft\", 20);\n const rightKey = new holdEvent.KeyboardKeyHold(\"ArrowRight\", 20);\n const upKey = new holdEvent.KeyboardKeyHold(\"ArrowUp\", 20);\n const downKey = new holdEvent.KeyboardKeyHold(\"ArrowDown\", 20);\n leftKey.addEventListener(\"holding\", (event) => {\n cameraControls.rotate(\n -0.05 * THREE.MathUtils.DEG2RAD * event?.deltaTime,\n 0,\n true,\n );\n });\n rightKey.addEventListener(\"holding\", (event) => {\n cameraControls.rotate(\n 0.05 * THREE.MathUtils.DEG2RAD * event?.deltaTime,\n 0,\n true,\n );\n });\n upKey.addEventListener(\"holding\", (event) => {\n cameraControls.rotate(\n 0,\n -0.05 * THREE.MathUtils.DEG2RAD * event?.deltaTime,\n true,\n );\n });\n downKey.addEventListener(\"holding\", (event) => {\n cameraControls.rotate(\n 0,\n 0.05 * THREE.MathUtils.DEG2RAD * event?.deltaTime,\n true,\n );\n });\n\n // TODO: we currently don't remove any event listeners. This is a bit messy\n // because KeyboardKeyHold attaches listeners directly to the\n // document/window; it's unclear if we can remove these.\n return () => {\n return;\n };\n }, [CameraControls]);\n\n return (\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/ClickUtils.tsx", "content": "import * as THREE from \"three\";\nimport { ViewerContextContents } from \"./App\";\n\n/** Turn a click event into a normalized device coordinate (NDC) vector.\n * Normalizes click coordinates to be between -1 and 1, with (0, 0) being the center of the screen.\n *\n * Returns null if input is not valid.\n */\nexport function ndcFromPointerXy(\n viewer: ViewerContextContents,\n xy: [number, number],\n): THREE.Vector2 | null {\n const mouseVector = new THREE.Vector2();\n mouseVector.x =\n 2 * ((xy[0] + 0.5) / viewer.canvasRef.current!.clientWidth) - 1;\n mouseVector.y =\n 1 - 2 * ((xy[1] + 0.5) / viewer.canvasRef.current!.clientHeight);\n return mouseVector.x < 1 &&\n mouseVector.x > -1 &&\n mouseVector.y < 1 &&\n mouseVector.y > -1\n ? mouseVector\n : null;\n}\n\n/** Turn a click event to normalized OpenCV coordinate (NDC) vector.\n * Normalizes click coordinates to be between (0, 0) as upper-left corner,\n * and (1, 1) as lower-right corner, with (0.5, 0.5) being the center of the screen.\n * Uses offsetX/Y, and clientWidth/Height to get the coordinates.\n */\nexport function opencvXyFromPointerXy(\n viewer: ViewerContextContents,\n xy: [number, number],\n): THREE.Vector2 {\n const mouseVector = new THREE.Vector2();\n mouseVector.x = (xy[0] + 0.5) / viewer.canvasRef.current!.clientWidth;\n mouseVector.y = (xy[1] + 0.5) / viewer.canvasRef.current!.clientHeight;\n return mouseVector;\n}\n" }, { "path": "viser/src/viser/client/src/ControlPanel/BottomPanel.tsx", "content": "import { Box, Collapse, Paper, useMantineColorScheme } from \"@mantine/core\";\nimport React from \"react\";\nimport { useDisclosure } from \"@mantine/hooks\";\n\nconst BottomPanelContext = React.createContext;\n expanded: boolean;\n toggleExpanded: () => void;\n}>(null);\n\n/** A bottom panel is used to display the controls on mobile devices. */\nexport default function BottomPanel({\n children,\n}: {\n children: string | React.ReactNode;\n}) {\n const panelWrapperRef = React.useRef(null);\n const [expanded, { toggle: toggleExpanded }] = useDisclosure(true);\n return (\n \n ({\n borderTopWidth: \"1px\",\n borderTopStyle: \"solid\",\n borderColor:\n useMantineColorScheme().colorScheme == \"dark\"\n ? theme.colors.dark[4]\n : theme.colors.gray[3],\n boxSizing: \"border-box\",\n width: \"100%\",\n zIndex: 10,\n position: \"fixed\",\n bottom: 0,\n left: 0,\n margin: 0,\n overflow: \"scroll\",\n minHeight: \"3.5em\",\n maxHeight: \"60%\",\n transition: \"height 0.3s linear\",\n })}\n ref={panelWrapperRef}\n >\n {children}\n \n \n );\n}\nBottomPanel.Handle = function BottomPanelHandle({\n children,\n}: {\n children: string | React.ReactNode;\n}) {\n const panelContext = React.useContext(BottomPanelContext)!;\n return (\n ({\n borderBottomWidth: panelContext.expanded ? \"1px\" : undefined,\n borderBottomStyle: \"solid\",\n borderColor:\n useMantineColorScheme().colorScheme == \"dark\"\n ? theme.colors.dark[4]\n : theme.colors.gray[3],\n cursor: \"pointer\",\n position: \"relative\",\n fontWeight: 400,\n userSelect: \"none\",\n display: \"flex\",\n alignItems: \"center\",\n padding: \"0 0.8em\",\n height: \"3.5em\",\n })}\n onClick={() => {\n panelContext.toggleExpanded();\n }}\n >\n {children}\n \n );\n};\n\n/** Contents of a panel. */\nBottomPanel.Contents = function BottomPanelContents({\n children,\n}: {\n children: string | React.ReactNode;\n}) {\n const panelContext = React.useContext(BottomPanelContext)!;\n return {children};\n};\n\n/** Hides contents when panel is collapsed. */\nBottomPanel.HideWhenCollapsed = function BottomPanelHideWhenCollapsed({\n children,\n}: {\n children: React.ReactNode;\n}) {\n const expanded = React.useContext(BottomPanelContext)?.expanded ?? true;\n return expanded ? children : null;\n};\n" }, { "path": "viser/src/viser/client/src/ControlPanel/ControlPanel.tsx", "content": "import { useDisclosure, useMediaQuery } from \"@mantine/hooks\";\nimport GeneratedGuiContainer from \"./Generated\";\nimport { ViewerContext } from \"../App\";\n\nimport QRCode from \"react-qr-code\";\nimport ServerControls from \"./ServerControls\";\nimport {\n ActionIcon,\n Anchor,\n Box,\n Button,\n Collapse,\n CopyButton,\n Flex,\n Loader,\n Modal,\n Stack,\n Text,\n TextInput,\n Tooltip,\n Transition,\n useMantineColorScheme,\n useMantineTheme,\n} from \"@mantine/core\";\nimport {\n IconAdjustments,\n IconCloudCheck,\n IconArrowBack,\n IconShare,\n IconCopy,\n IconCheck,\n IconPlugConnectedX,\n IconQrcode,\n IconQrcodeOff,\n} from \"@tabler/icons-react\";\nimport React from \"react\";\nimport BottomPanel from \"./BottomPanel\";\nimport FloatingPanel from \"./FloatingPanel\";\nimport { ThemeConfigurationMessage } from \"../WebsocketMessages\";\nimport SidebarPanel from \"./SidebarPanel\";\n\n// Must match constant in Python.\nconst ROOT_CONTAINER_ID = \"root\";\n\nexport default function ControlPanel(props: {\n control_layout: ThemeConfigurationMessage[\"control_layout\"];\n}) {\n const theme = useMantineTheme();\n const useMobileView = useMediaQuery(`(max-width: ${theme.breakpoints.xs})`);\n\n // TODO: will result in unnecessary re-renders.\n const viewer = React.useContext(ViewerContext)!;\n const showGenerated = viewer.useGui(\n (state) => \"root\" in state.guiIdSetFromContainerId,\n );\n const [showSettings, { toggle }] = useDisclosure(false);\n\n const controlWidthString = viewer.useGui(\n (state) => state.theme.control_width,\n );\n const controlWidth = (\n controlWidthString == \"small\"\n ? \"16em\"\n : controlWidthString == \"medium\"\n ? \"20em\"\n : controlWidthString == \"large\"\n ? \"24em\"\n : null\n )!;\n\n const generatedServerToggleButton = (\n {\n evt.stopPropagation();\n toggle();\n }}\n style={{\n display: showGenerated ? undefined : \"none\",\n transform: \"translateY(0.05em)\",\n }}\n >\n \n {showSettings ? (\n \n ) : (\n \n )}\n \n \n );\n\n const panelContents = (\n <>\n \n \n \n \n \n \n \n );\n\n if (useMobileView) {\n /* Mobile layout. */\n return (\n \n \n \n \n \n {generatedServerToggleButton}\n \n \n {panelContents}\n \n );\n } else if (props.control_layout === \"floating\") {\n /* Floating layout. */\n return (\n \n \n \n \n \n {generatedServerToggleButton}\n \n \n {panelContents}\n \n );\n } else {\n /* Sidebar view. */\n return (\n \n \n \n \n {generatedServerToggleButton}\n \n {panelContents}\n \n );\n }\n}\n\n/* Icon and label telling us the current status of the websocket connection. */\nfunction ConnectionStatus() {\n const { useGui } = React.useContext(ViewerContext)!;\n const connected = useGui((state) => state.websocketConnected);\n const label = useGui((state) => state.label);\n\n return (\n <>\n
{/* Spacer. */}\n \n {(styles) => (\n \n )}\n \n \n {(styles) => (\n \n )}\n \n \n {label !== \"\" ? label : connected ? \"Connected\" : \"Connecting...\"}\n \n \n );\n}\n\nfunction ShareButton() {\n const viewer = React.useContext(ViewerContext)!;\n const connected = viewer.useGui((state) => state.websocketConnected);\n const shareUrl = viewer.useGui((state) => state.shareUrl);\n const setShareUrl = viewer.useGui((state) => state.setShareUrl);\n\n const [doingSomething, setDoingSomething] = React.useState(false);\n\n const [shareModalOpened, { open: openShareModal, close: closeShareModal }] =\n useDisclosure(false);\n\n const [showQrCode, { toggle: toggleShowQrcode }] = useDisclosure();\n\n // Turn off loader when share URL is set.\n React.useEffect(() => {\n if (shareUrl !== null) {\n setDoingSomething(false);\n }\n }, [shareUrl]);\n React.useEffect(() => {\n if (!connected && shareModalOpened) closeShareModal();\n }, [connected, shareModalOpened]);\n\n if (viewer.useGui((state) => state.theme).show_share_button === false)\n return null;\n\n const colorScheme = useMantineColorScheme().colorScheme;\n return (\n <>\n \n {\n evt.stopPropagation();\n openShareModal();\n }}\n style={{\n transform: \"translateY(0.05em)\",\n }}\n disabled={!connected}\n >\n \n \n \n evt.stopPropagation()}\n onMouseDown={(evt) => evt.stopPropagation()}\n onMouseMove={(evt) => evt.stopPropagation()}\n onMouseUp={(evt) => evt.stopPropagation()}\n styles={{ title: { fontWeight: 600 } }}\n >\n {shareUrl === null ? (\n <>\n {/* */}\n {doingSomething ? (\n \n \n \n ) : (\n \n \n Create a public, shareable URL to this Viser instance.\n \n {\n viewer.sendMessageRef.current({\n type: \"ShareUrlRequest\",\n });\n setDoingSomething(true); // Loader state will help with debouncing.\n }}\n >\n Request Share URL\n \n \n )}\n \n ) : (\n <>\n Share URL is connected.\n \n \n \n \n {({ copied, copy }) => (\n \n ) : (\n \n )\n }\n onClick={copy}\n variant={copied ? \"outline\" : \"filled\"}\n >\n {copied ? \"Copied!\" : \"Copy URL\"}\n \n )}\n \n : }\n onClick={toggleShowQrcode}\n >\n QR Code\n \n \n {\n viewer.sendMessageRef.current({\n type: \"ShareUrlDisconnect\",\n });\n setShareUrl(null);\n }}\n >\n \n \n \n \n \n \n \n \n \n )}\n \n This feature is experimental. Problems? Consider{\" \"}\n \n reporting on GitHub\n \n .\n \n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/ControlPanel/FloatingPanel.tsx", "content": "// @refresh reset\n\nimport {\n Box,\n Collapse,\n Paper,\n ScrollArea,\n useMantineColorScheme,\n} from \"@mantine/core\";\nimport React from \"react\";\nimport { isMouseEvent, isTouchEvent, mouseEvents, touchEvents } from \"../Utils\";\nimport { useDisclosure } from \"@mantine/hooks\";\n\nconst FloatingPanelContext = React.createContext;\n expanded: boolean;\n width: string;\n maxHeight: number;\n toggleExpanded: () => void;\n dragHandler: (\n event:\n | React.TouchEvent\n | React.MouseEvent,\n ) => void;\n dragInfo: React.MutableRefObject<{\n dragging: boolean;\n startPosX: number;\n startPosY: number;\n startClientX: number;\n startClientY: number;\n }>;\n}>(null);\n\n/** A floating panel for displaying controls. */\nexport default function FloatingPanel({\n children,\n width,\n}: {\n children: string | React.ReactNode;\n width: string;\n}) {\n const panelWrapperRef = React.useRef(null);\n const [expanded, { toggle: toggleExpanded }] = useDisclosure(true);\n const [maxHeight, setMaxHeight] = React.useState(800);\n\n // Things to track for dragging.\n const dragInfo = React.useRef({\n dragging: false,\n startPosX: 0,\n startPosY: 0,\n startClientX: 0,\n startClientY: 0,\n });\n\n // Logic for \"fixing\" panel locations, which keeps the control panel within\n // the bounds of the parent div.\n //\n // For `unfixedOffset`, we use a negative sign to indicate that the panel is\n // positioned relative to the right/bottom bound of the parent.\n const unfixedOffset = React.useRef<{ x?: number; y?: number }>({});\n const computePanelOffset = (\n panelPosition: number,\n panelSize: number,\n parentSize: number,\n ) =>\n Math.abs(panelPosition + panelSize / 2.0) <\n Math.abs(panelPosition - parentSize + panelSize / 2.0)\n ? panelPosition\n : panelPosition - parentSize;\n\n const panelBoundaryPad = 15;\n function setPanelLocation(x: number, y: number) {\n const panel = panelWrapperRef.current;\n if (panel === null) return [x, y];\n\n const parent = panel.parentElement;\n if (parent === null) return [x, y];\n\n let newX = x;\n let newY = y;\n\n newX = Math.min(\n newX,\n parent.clientWidth - panel.clientWidth - panelBoundaryPad,\n );\n newX = Math.max(newX, panelBoundaryPad);\n newY = Math.min(\n newY,\n parent.clientHeight - panel.clientHeight - panelBoundaryPad,\n );\n newY = Math.max(newY, panelBoundaryPad);\n\n panel.style.top = `${newY.toString()}px`;\n panel.style.left = `${newX.toString()}px`;\n\n return [\n computePanelOffset(newX, panel.clientWidth, parent.clientWidth),\n computePanelOffset(newY, panel.clientHeight, parent.clientHeight),\n ];\n }\n\n // Fix locations on resize.\n React.useEffect(() => {\n const panel = panelWrapperRef.current;\n if (panel === null) return;\n\n const parent = panel.parentElement;\n if (parent === null) return;\n\n const observer = new ResizeObserver(() => {\n if (unfixedOffset.current.x === undefined)\n unfixedOffset.current.x = computePanelOffset(\n panel.offsetLeft,\n panel.clientWidth,\n parent.clientWidth,\n );\n if (unfixedOffset.current.y === undefined)\n unfixedOffset.current.y = computePanelOffset(\n panel.offsetTop,\n panel.clientHeight,\n parent.clientHeight,\n );\n\n const newMaxHeight =\n parent.clientHeight - panelBoundaryPad * 2 - 2.5 * 16;\n maxHeight !== newMaxHeight && setMaxHeight(newMaxHeight);\n\n let newX = unfixedOffset.current.x;\n let newY = unfixedOffset.current.y;\n while (newX < 0) newX += parent.clientWidth;\n while (newY < 0) newY += parent.clientHeight;\n setPanelLocation(newX, newY);\n });\n observer.observe(panel);\n observer.observe(parent);\n return () => {\n observer.disconnect();\n };\n });\n\n const dragHandler = (\n event:\n | React.TouchEvent\n | React.MouseEvent,\n ) => {\n const state = dragInfo.current;\n const panel = panelWrapperRef.current;\n if (!panel) return;\n if (event.type == \"touchstart\") {\n event = event as React.TouchEvent;\n state.startClientX = event.touches[0].clientX;\n state.startClientY = event.touches[0].clientY;\n } else {\n event = event as React.MouseEvent;\n state.startClientX = event.clientX;\n state.startClientY = event.clientY;\n }\n state.startPosX = panel.offsetLeft;\n state.startPosY = panel.offsetTop;\n const eventNames = event.type == \"touchstart\" ? touchEvents : mouseEvents;\n function dragListener(event: MouseEvent | TouchEvent) {\n // Minimum motion.\n let deltaX = 0;\n let deltaY = 0;\n if (isTouchEvent(event)) {\n event = event as TouchEvent;\n deltaX = event.touches[0].clientX - state.startClientX;\n deltaY = event.touches[0].clientY - state.startClientY;\n } else if (isMouseEvent(event)) {\n event = event as MouseEvent;\n deltaX = event.clientX - state.startClientX;\n deltaY = event.clientY - state.startClientY;\n }\n if (Math.abs(deltaX) <= 3 && Math.abs(deltaY) <= 3) return;\n\n state.dragging = true;\n const newX = state.startPosX + deltaX;\n const newY = state.startPosY + deltaY;\n [unfixedOffset.current.x, unfixedOffset.current.y] = setPanelLocation(\n newX,\n newY,\n );\n }\n window.addEventListener(eventNames.move, dragListener);\n window.addEventListener(\n eventNames.end,\n () => {\n if (event.type == \"touchstart\") {\n state.dragging = false;\n }\n window.removeEventListener(eventNames.move, dragListener);\n },\n { once: true },\n );\n };\n\n return (\n \n \n {children}\n \n \n );\n}\n\n/** Handle object helps us hide, show, and drag our panel.*/\nFloatingPanel.Handle = function FloatingPanelHandle({\n children,\n}: {\n children: string | React.ReactNode;\n}) {\n const panelContext = React.useContext(FloatingPanelContext)!;\n\n return (\n ({\n borderRadius: \"0.2em 0.2em 0 0\",\n lineHeight: \"1.5em\",\n cursor: \"pointer\",\n position: \"relative\",\n fontWeight: 400,\n userSelect: \"none\",\n display: \"flex\",\n alignItems: \"center\",\n padding: \"0 0.75em\",\n height: \"2.75em\",\n borderBottomWidth: panelContext.expanded ? \"1px\" : 0,\n borderBottomStyle: \"solid\",\n borderColor:\n useMantineColorScheme().colorScheme == \"dark\"\n ? theme.colors.dark[4]\n : theme.colors.gray[3],\n })}\n onClick={() => {\n const state = panelContext.dragInfo.current;\n if (state.dragging) {\n state.dragging = false;\n return;\n }\n panelContext.toggleExpanded();\n }}\n onTouchStart={(event) => {\n panelContext.dragHandler(event);\n }}\n onMouseDown={(event) => {\n panelContext.dragHandler(event);\n }}\n >\n {children}\n \n );\n};\n/** Contents of a panel. */\nFloatingPanel.Contents = function FloatingPanelContents({\n children,\n}: {\n children: string | React.ReactNode;\n}) {\n const context = React.useContext(FloatingPanelContext)!;\n return (\n \n \n \n {children}\n \n \n \n );\n};\n\n/** Hides contents when floating panel is collapsed. */\nFloatingPanel.HideWhenCollapsed = function FloatingPanelHideWhenCollapsed({\n children,\n}: {\n children: React.ReactNode;\n}) {\n const expanded = React.useContext(FloatingPanelContext)?.expanded ?? true;\n return expanded ? children : null;\n};\n" }, { "path": "viser/src/viser/client/src/ControlPanel/Generated.tsx", "content": "import { ViewerContext } from \"../App\";\nimport { useThrottledMessageSender } from \"../WebsocketFunctions\";\nimport { GuiComponentContext } from \"./GuiComponentContext\";\n\nimport { Box } from \"@mantine/core\";\nimport React from \"react\";\nimport ButtonComponent from \"../components/Button\";\nimport SliderComponent from \"../components/Slider\";\nimport NumberInputComponent from \"../components/NumberInput\";\nimport TextInputComponent from \"../components/TextInput\";\nimport CheckboxComponent from \"../components/Checkbox\";\nimport Vector2Component from \"../components/Vector2\";\nimport Vector3Component from \"../components/Vector3\";\nimport DropdownComponent from \"../components/Dropdown\";\nimport RgbComponent from \"../components/Rgb\";\nimport RgbaComponent from \"../components/Rgba\";\nimport ButtonGroupComponent from \"../components/ButtonGroup\";\nimport MarkdownComponent from \"../components/Markdown\";\nimport PlotlyComponent from \"../components/PlotlyComponent\";\nimport TabGroupComponent from \"../components/TabGroup\";\nimport FolderComponent from \"../components/Folder\";\nimport MultiSliderComponent from \"../components/MultiSlider\";\nimport UploadButtonComponent from \"../components/UploadButton\";\nimport ProgressBarComponent from \"../components/ProgressBar\";\n\n/** Root of generated inputs. */\nexport default function GeneratedGuiContainer({\n containerId,\n}: {\n containerId: string;\n}) {\n const viewer = React.useContext(ViewerContext)!;\n const updateGuiProps = viewer.useGui((state) => state.updateGuiProps);\n const messageSender = useThrottledMessageSender(50);\n\n function setValue(id: string, value: NonNullable) {\n updateGuiProps(id, { value: value });\n messageSender({\n type: \"GuiUpdateMessage\",\n id: id,\n updates: { value: value },\n });\n }\n return (\n \n \n \n );\n}\n\nfunction GuiContainer({ containerId }: { containerId: string }) {\n const viewer = React.useContext(ViewerContext)!;\n\n const guiIdSet =\n viewer.useGui((state) => state.guiIdSetFromContainerId[containerId]) ?? {};\n\n // Render each GUI element in this container.\n const guiIdArray = [...Object.keys(guiIdSet)];\n const guiOrderFromId = viewer!.useGui((state) => state.guiOrderFromId);\n if (guiIdSet === undefined) return null;\n\n let guiIdOrderPairArray = guiIdArray.map((id) => ({\n id: id,\n order: guiOrderFromId[id],\n }));\n guiIdOrderPairArray = guiIdOrderPairArray.sort((a, b) => a.order - b.order);\n const out = (\n \n {guiIdOrderPairArray.map((pair) => (\n \n ))}\n \n );\n return out;\n}\n\n/** A single generated GUI element. */\nfunction GeneratedInput(props: { guiId: string }) {\n const viewer = React.useContext(ViewerContext)!;\n const conf = viewer.useGui((state) => state.guiConfigFromId[props.guiId]);\n switch (conf.type) {\n case \"GuiAddFolderMessage\":\n return ;\n case \"GuiAddTabGroupMessage\":\n return ;\n case \"GuiAddMarkdownMessage\":\n return ;\n case \"GuiAddPlotlyMessage\":\n return ;\n case \"GuiAddButtonMessage\":\n return ;\n case \"GuiAddUploadButtonMessage\":\n return ;\n case \"GuiAddSliderMessage\":\n return ;\n case \"GuiAddMultiSliderMessage\":\n return ;\n case \"GuiAddNumberMessage\":\n return ;\n case \"GuiAddTextMessage\":\n return ;\n case \"GuiAddCheckboxMessage\":\n return ;\n case \"GuiAddVector2Message\":\n return ;\n case \"GuiAddVector3Message\":\n return ;\n case \"GuiAddDropdownMessage\":\n return ;\n case \"GuiAddRgbMessage\":\n return ;\n case \"GuiAddRgbaMessage\":\n return ;\n case \"GuiAddButtonGroupMessage\":\n return ;\n case \"GuiAddProgressBarMessage\":\n return ;\n default:\n assertNeverType(conf);\n }\n}\n\nfunction assertNeverType(x: never): never {\n throw new Error(\"Unexpected object: \" + (x as any).type);\n}\n" }, { "path": "viser/src/viser/client/src/ControlPanel/GuiComponentContext.tsx", "content": "import * as React from \"react\";\nimport * as Messages from \"../WebsocketMessages\";\n\ninterface GuiComponentContext {\n folderDepth: number;\n setValue: (id: string, value: NonNullable) => void;\n messageSender: (message: Messages.Message) => void;\n GuiContainer: React.FC<{ containerId: string }>;\n}\n\nexport const GuiComponentContext = React.createContext({\n folderDepth: 0,\n setValue: () => undefined,\n messageSender: () => undefined,\n GuiContainer: () => {\n throw new Error(\"GuiComponentContext not initialized\");\n },\n});\n" }, { "path": "viser/src/viser/client/src/ControlPanel/GuiState.tsx", "content": "import * as Messages from \"../WebsocketMessages\";\nimport React from \"react\";\nimport { create } from \"zustand\";\nimport { ColorTranslator } from \"colortranslator\";\n\nimport { immer } from \"zustand/middleware/immer\";\n\nexport type GuiConfig = Messages.GuiAddComponentMessage;\n\nexport function isGuiConfig(message: Messages.Message): message is GuiConfig {\n return message.type.startsWith(\"GuiAdd\");\n}\n\ninterface GuiState {\n theme: Messages.ThemeConfigurationMessage;\n label: string;\n server: string;\n shareUrl: string | null;\n websocketConnected: boolean;\n backgroundAvailable: boolean;\n guiIdSetFromContainerId: {\n [containerId: string]: { [id: string]: true } | undefined;\n };\n modals: Messages.GuiModalMessage[];\n guiOrderFromId: { [id: string]: number };\n guiConfigFromId: { [id: string]: GuiConfig };\n uploadsInProgress: {\n [id: string]: {\n notificationId: string;\n uploadedBytes: number;\n totalBytes: number;\n filename: string;\n };\n };\n}\n\ninterface GuiActions {\n setTheme: (theme: Messages.ThemeConfigurationMessage) => void;\n setShareUrl: (share_url: string | null) => void;\n addGui: (config: GuiConfig) => void;\n addModal: (config: Messages.GuiModalMessage) => void;\n removeModal: (id: string) => void;\n updateGuiProps: (id: string, updates: { [key: string]: any }) => void;\n removeGui: (id: string) => void;\n resetGui: () => void;\n updateUploadState: (\n state: (\n | { uploadedBytes: number; totalBytes: number }\n | GuiState[\"uploadsInProgress\"][string]\n ) & { componentId: string },\n ) => void;\n}\n\nconst cleanGuiState: GuiState = {\n theme: {\n type: \"ThemeConfigurationMessage\",\n titlebar_content: null,\n control_layout: \"floating\",\n control_width: \"medium\",\n dark_mode: false,\n show_logo: true,\n show_share_button: true,\n colors: null,\n },\n label: \"\",\n server: \"ws://localhost:8080\", // Currently this will always be overridden.\n shareUrl: null,\n websocketConnected: false,\n backgroundAvailable: false,\n guiIdSetFromContainerId: {},\n modals: [],\n guiOrderFromId: {},\n guiConfigFromId: {},\n uploadsInProgress: {},\n};\n\nexport function computeRelativeLuminance(color: string) {\n const colorTrans = new ColorTranslator(color);\n\n // Coefficients are from:\n // https://en.wikipedia.org/wiki/Relative_luminance#Relative_luminance_and_%22gamma_encoded%22_colorspaces\n return (\n ((0.2126 * colorTrans.R + 0.7152 * colorTrans.G + 0.0722 * colorTrans.B) /\n 255.0) *\n 100.0\n );\n}\n\nexport function useGuiState(initialServer: string) {\n return React.useState(() =>\n create(\n immer((set) => ({\n ...cleanGuiState,\n server: initialServer,\n setTheme: (theme) =>\n set((state) => {\n state.theme = theme;\n }),\n setShareUrl: (share_url) =>\n set((state) => {\n state.shareUrl = share_url;\n }),\n addGui: (guiConfig) =>\n set((state) => {\n state.guiOrderFromId[guiConfig.id] = guiConfig.order;\n state.guiConfigFromId[guiConfig.id] = guiConfig;\n if (!(guiConfig.container_id in state.guiIdSetFromContainerId)) {\n state.guiIdSetFromContainerId[guiConfig.container_id] = {};\n }\n state.guiIdSetFromContainerId[guiConfig.container_id]![\n guiConfig.id\n ] = true;\n }),\n addModal: (modalConfig) =>\n set((state) => {\n state.modals.push(modalConfig);\n }),\n removeModal: (id) =>\n set((state) => {\n state.modals = state.modals.filter((m) => m.id !== id);\n }),\n removeGui: (id) =>\n set((state) => {\n const guiConfig = state.guiConfigFromId[id];\n\n delete state.guiIdSetFromContainerId[guiConfig.container_id]![id];\n delete state.guiOrderFromId[id];\n delete state.guiConfigFromId[id];\n if (\n Object.keys(\n state.guiIdSetFromContainerId[guiConfig.container_id]!,\n ).length == 0\n )\n delete state.guiIdSetFromContainerId[guiConfig.container_id];\n }),\n resetGui: () =>\n set((state) => {\n // No need to overwrite the theme or label. The former especially\n // can be jarring.\n // state.theme = cleanGuiState.theme;\n // state.label = cleanGuiState.label;\n\n // This feels brittle, could be cleaned up...\n state.shareUrl = null;\n state.guiIdSetFromContainerId = {};\n state.modals = [];\n state.guiOrderFromId = {};\n state.guiConfigFromId = {};\n state.uploadsInProgress = {};\n }),\n updateUploadState: (state) =>\n set((globalState) => {\n const { componentId, ...rest } = state;\n globalState.uploadsInProgress[componentId] = {\n ...globalState.uploadsInProgress[componentId],\n ...rest,\n };\n }),\n updateGuiProps: (id, updates) => {\n set((state) => {\n const config = state.guiConfigFromId[id];\n if (config === undefined) {\n console.error(\"Tried to update non-existent component\", id);\n return;\n }\n\n // Double-check that key exists.\n Object.keys(updates).forEach((key) => {\n if (!(key in config))\n console.error(\n `Tried to update nonexistent property '${key}' of GUI element ${id}!`,\n );\n });\n\n state.guiConfigFromId[id] = {\n ...config,\n ...updates,\n } as GuiConfig;\n });\n },\n })),\n ),\n )[0];\n}\n\n/** Type corresponding to a zustand-style useGuiState hook. */\nexport type UseGui = ReturnType;\n" }, { "path": "viser/src/viser/client/src/ControlPanel/SceneTreeTable.css.ts", "content": "import { style } from \"@vanilla-extract/css\";\nimport { vars } from \"../AppTheme\";\n\nexport const tableWrapper = style({\n border: \"1px solid\",\n borderColor: vars.colors.defaultBorder,\n borderRadius: vars.radius.xs,\n});\nexport const icon = style({\n opacity: 0.5,\n height: \"1em\",\n width: \"1em\",\n transform: \"translateY(0.1em)\",\n});\n\nexport const tableRow = style({\n display: \"flex\",\n alignItems: \"center\",\n gap: \"0.4em\",\n padding: \"0 0.25em\",\n lineHeight: \"2.25em\",\n fontSize: \"0.875em\",\n});\n" }, { "path": "viser/src/viser/client/src/ControlPanel/SceneTreeTable.tsx", "content": "import { ViewerContext } from \"../App\";\nimport { Box, Stack, Tooltip } from \"@mantine/core\";\nimport {\n IconCaretDown,\n IconCaretRight,\n IconEye,\n IconEyeOff,\n} from \"@tabler/icons-react\";\nimport React from \"react\";\nimport {\n icon as caretIcon,\n tableRow,\n tableWrapper,\n} from \"./SceneTreeTable.css\";\nimport { useDisclosure } from \"@mantine/hooks\";\n\n/* Table for seeing an overview of the scene tree, toggling visibility, etc. * */\nexport default function SceneTreeTable() {\n const viewer = React.useContext(ViewerContext)!;\n const childrenName = viewer.useSceneTree(\n (state) => state.nodeFromName[\"\"]!.children,\n );\n return (\n \n {childrenName.map((name) => (\n \n ))}\n \n );\n}\n\nconst SceneTreeTableRow = React.memo(function SceneTreeTableRow(props: {\n nodeName: string;\n isParentVisible: boolean;\n indentCount: number;\n}) {\n const viewer = React.useContext(ViewerContext)!;\n const childrenName = viewer.useSceneTree(\n (state) => state.nodeFromName[props.nodeName]!.children,\n );\n const expandable = childrenName.length > 0;\n\n const [expanded, { toggle: toggleExpanded }] = useDisclosure(false);\n\n function setOverrideVisibility(name: string, visible: boolean | undefined) {\n const attr = viewer.nodeAttributesFromName.current;\n attr[name]!.overrideVisibility = visible;\n rerenderTable();\n }\n const setLabelVisibility = viewer.useSceneTree(\n (state) => state.setLabelVisibility,\n );\n\n // For performance, scene node visibility is stored in a ref instead of the\n // zustand state. This means that re-renders for the table need to be\n // triggered manually when visibilities are updated.\n const [, setTime] = React.useState(Date.now());\n function rerenderTable() {\n setTime(Date.now());\n }\n React.useEffect(() => {\n const interval = setInterval(rerenderTable, 200);\n return () => {\n clearInterval(interval);\n };\n }, []);\n\n const attrs = viewer.nodeAttributesFromName.current[props.nodeName];\n const isVisible =\n (attrs?.overrideVisibility === undefined\n ? attrs?.visibility\n : attrs.overrideVisibility) ?? true;\n const isVisibleEffective = isVisible && props.isParentVisible;\n const VisibleIcon = isVisible ? IconEye : IconEyeOff;\n\n return (\n <>\n setLabelVisibility(props.nodeName, true)}\n onMouseOut={() => setLabelVisibility(props.nodeName, false)}\n >\n \n {expanded ? (\n \n ) : (\n \n )}\n \n \n {\n evt.stopPropagation();\n setOverrideVisibility(props.nodeName, !isVisible);\n }}\n />\n \n \n {props.nodeName\n .split(\"/\")\n .filter((part) => part.length > 0)\n .map((part, index, all) => (\n // We set userSelect to prevent users from accidentally\n // selecting text when dragging over the hide/show icons.\n \n \n {index === all.length - 1 ? \"/\" : `/${part}`}\n \n {index === all.length - 1 ? part : \"\"}\n \n ))}\n \n \n {expanded\n ? childrenName.map((name) => (\n \n ))\n : null}\n \n );\n});\n" }, { "path": "viser/src/viser/client/src/ControlPanel/ServerControls.tsx", "content": "import { ViewerContext } from \"../App\";\nimport {\n Box,\n Button,\n Divider,\n Stack,\n Switch,\n Text,\n TextInput,\n} from \"@mantine/core\";\nimport { IconHomeMove, IconPhoto } from \"@tabler/icons-react\";\nimport { Stats } from \"@react-three/drei\";\nimport React from \"react\";\nimport SceneTreeTable from \"./SceneTreeTable\";\n\nexport default function ServerControls() {\n const viewer = React.useContext(ViewerContext)!;\n const [showStats, setShowStats] = React.useState(false);\n\n function triggerBlur(event: React.KeyboardEvent) {\n if (event.key !== \"Enter\") return;\n event.currentTarget.blur();\n event.currentTarget.focus();\n }\n const MemoizedTable = React.memo(SceneTreeTable);\n\n return (\n <>\n {showStats ? : null}\n \n state.server)}\n onBlur={(event) =>\n viewer.useGui.setState({ server: event.currentTarget.value })\n }\n onKeyDown={triggerBlur}\n styles={{\n input: {\n minHeight: \"1.75rem\",\n height: \"1.75rem\",\n padding: \"0 0.5em\",\n },\n }}\n />\n state.label)}\n onBlur={(event) =>\n viewer.useGui.setState({ label: event.currentTarget.value })\n }\n onKeyDown={triggerBlur}\n styles={{\n input: {\n minHeight: \"1.75rem\",\n height: \"1.75rem\",\n padding: \"0 0.5em\",\n },\n }}\n mb=\"0.375em\"\n />\n {\n const supportsFileSystemAccess =\n \"showSaveFilePicker\" in window &&\n (() => {\n try {\n return window.self === window.top;\n } catch {\n return false;\n }\n })();\n\n if (supportsFileSystemAccess) {\n // File System Access API is supported. (eg Chrome)\n const fileHandlePromise = window.showSaveFilePicker({\n suggestedName: \"render.png\",\n types: [\n {\n accept: { \"image/png\": [\".png\"] },\n },\n ],\n });\n viewer.canvasRef.current?.toBlob(async (blob) => {\n if (blob === null) {\n console.error(\"Export failed\");\n return;\n }\n\n const handle = await fileHandlePromise;\n const writableStream = await handle.createWritable();\n await writableStream.write(blob);\n await writableStream.close();\n });\n } else {\n // File System Access API is not supported. (eg Firefox)\n viewer.canvasRef.current?.toBlob((blob) => {\n if (blob === null) {\n console.error(\"Export failed\");\n return;\n }\n const href = URL.createObjectURL(blob);\n\n // Download a file by creating a link and then clicking it.\n const link = document.createElement(\"a\");\n link.href = href;\n const filename = \"render.png\";\n link.download = filename;\n document.body.appendChild(link);\n link.click();\n document.body.removeChild(link);\n URL.revokeObjectURL(href);\n });\n }\n }}\n fullWidth\n leftSection={}\n style={{ height: \"1.875rem\" }}\n >\n Export Canvas\n \n {\n viewer.resetCameraViewRef.current!();\n }}\n fullWidth\n leftSection={}\n style={{ height: \"1.875rem\" }}\n >\n Reset View\n \n {\n setShowStats(event.currentTarget.checked);\n }}\n size=\"sm\"\n />\n \n \n \n Scene tree\n \n \n \n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/ControlPanel/SidebarPanel.tsx", "content": "// @refresh reset\n\nimport {\n ActionIcon,\n Box,\n Paper,\n ScrollArea,\n Tooltip,\n useMantineColorScheme,\n} from \"@mantine/core\";\nimport React from \"react\";\nimport { useDisclosure } from \"@mantine/hooks\";\nimport { IconChevronLeft, IconChevronRight } from \"@tabler/icons-react\";\n\nexport const SidebarPanelContext = React.createContext void;\n}>(null);\n\n/** A fixed or collapsible side panel for displaying controls. */\nexport default function SidebarPanel({\n children,\n collapsible,\n width,\n}: {\n children: string | React.ReactNode;\n collapsible: boolean;\n width: string;\n}) {\n const [collapsed, { toggle: toggleCollapsed }] = useDisclosure(false);\n\n const collapsedView = (\n ({\n /* Animate in when collapsed. */\n position: \"absolute\",\n top: 0,\n right: collapsed ? \"0em\" : \"-3em\",\n transitionProperty: \"right\",\n transitionDuration: \"0.5s\",\n transitionDelay: \"0.25s\",\n /* Visuals. */\n borderBottomLeftRadius: \"0.5em\",\n backgroundColor:\n useMantineColorScheme().colorScheme == \"dark\"\n ? theme.colors.dark[5]\n : theme.colors.gray[2],\n padding: \"0.5em\",\n })}\n >\n {\n evt.stopPropagation();\n toggleCollapsed();\n }}\n >\n \n {}\n \n \n \n );\n\n return (\n \n {collapsedView}\n {/* Using an
\n );\n } else {\n return (\n \n \n {paused ? (\n setPaused(false)}\n height=\"1.125em\"\n width=\"1.125em\"\n />\n ) : (\n setPaused(true)}\n height=\"1.125em\"\n width=\"1.125em\"\n />\n )}\n \n \n updateCurrentTime(\n typeof value === \"number\" ? value : parseFloat(value),\n )\n }\n />\n \n \n (val === null ? null : setPlaybackSpeed(val))}\n radius=\"xs\"\n data={[\"0.5x\", \"1x\", \"2x\", \"4x\", \"8x\"]}\n styles={{\n wrapper: { width: \"3.25em\" },\n }}\n comboboxProps={{ zIndex: 5, width: \"5.25em\" }}\n />\n \n \n );\n }\n}\n" }, { "path": "viser/src/viser/client/src/Markdown.tsx", "content": "import React from \"react\";\nimport * as runtime from \"react/jsx-runtime\";\nimport * as provider from \"@mdx-js/react\";\nimport { evaluate } from \"@mdx-js/mdx\";\nimport { type MDXComponents } from \"mdx/types\";\nimport { ReactNode, useEffect, useState } from \"react\";\nimport remarkGfm from \"remark-gfm\";\nimport rehypeColorChips from \"rehype-color-chips\";\nimport {\n Anchor,\n Blockquote,\n Code,\n Image,\n List,\n ListProps,\n Table,\n Text,\n Title,\n TitleOrder,\n} from \"@mantine/core\";\nimport { visit } from \"unist-util-visit\";\nimport { Transformer } from \"unified\";\nimport { Root } from \"hast\";\n\n// Custom Rehype to clean up code blocks (Mantine makes these annoying to style)\n// Adds \"block\" to any code non-inline code block, which gets directly passed into\n// the Mantine Code component.\nfunction rehypeCodeblock(): void | Transformer {\n return (tree) => {\n visit(tree, \"element\", (node, _i, parent) => {\n if (node.tagName !== \"code\") return;\n if (parent && parent.type === \"element\" && parent.tagName === \"pre\") {\n node.properties = { block: true, ...node.properties };\n }\n });\n };\n}\n\n// Custom classes to pipe MDX into Mantine Components\n// Some of them separate the children into a separate prop since Mantine requires a child\n// and MDX always makes children optional, so destructuring props doesn't work\nfunction MdxText(props: React.ComponentPropsWithoutRef) {\n return ;\n}\n\nfunction MdxAnchor(props: React.ComponentPropsWithoutRef) {\n return ;\n}\n\nfunction MdxTitle(\n props: React.ComponentPropsWithoutRef,\n order: TitleOrder,\n) {\n return ;\n}\n\nfunction MdxList(\n props: Omit, \"children\" | \"type\">,\n children: React.ComponentPropsWithoutRef[\"children\"],\n type: ListProps[\"type\"],\n) {\n // Account for GFM Checkboxes\n if (props.className == \"contains-task-list\") {\n return (\n \n {children}\n \n );\n }\n return (\n \n {children}\n \n );\n}\n\nfunction MdxListItem(\n props: Omit, \"children\">,\n children: React.ComponentPropsWithoutRef[\"children\"],\n) {\n return {children};\n}\n\n// A possible improvement is to use Mantine Prism to add code highlighting support.\nfunction MdxCode(\n props: Omit, \"children\">,\n children: React.ComponentPropsWithoutRef[\"children\"],\n) {\n return {children};\n}\n\nfunction MdxBlockquote(\n props: React.ComponentPropsWithoutRef,\n) {\n return
;\n}\n\nfunction MdxCite(\n props: React.DetailedHTMLProps<\n React.HTMLAttributes,\n HTMLElement\n >,\n) {\n return (\n \n );\n}\n\nfunction MdxTable(props: React.ComponentPropsWithoutRef) {\n return ;\n}\n\nfunction MdxImage(props: React.ComponentPropsWithoutRef) {\n return ;\n}\n\nconst components: MDXComponents = {\n p: (props) => MdxText(props),\n a: (props) => MdxAnchor(props),\n h1: (props) => MdxTitle(props, 1),\n h2: (props) => MdxTitle(props, 2),\n h3: (props) => MdxTitle(props, 3),\n h4: (props) => MdxTitle(props, 4),\n h5: (props) => MdxTitle(props, 5),\n h6: (props) => MdxTitle(props, 6),\n ul: (props) => MdxList(props, props.children ?? \"\", \"unordered\"),\n ol: (props) => MdxList(props, props.children ?? \"\", \"ordered\"),\n li: (props) => MdxListItem(props, props.children ?? \"\"),\n code: (props) => MdxCode(props, props.children ?? \"\"),\n pre: (props) => <>{props.children},\n blockquote: (props) => MdxBlockquote(props),\n Cite: (props) => MdxCite(props),\n table: (props) => MdxTable(props),\n img: (props) => MdxImage(props),\n \"*\": () => <>,\n};\n\nasync function parseMarkdown(markdown: string) {\n // @ts-ignore (necessary since JSX runtime isn't properly typed according to the internet)\n const { default: Content } = await evaluate(markdown, {\n ...runtime,\n ...provider,\n development: false,\n remarkPlugins: [remarkGfm],\n rehypePlugins: [rehypeCodeblock, rehypeColorChips],\n });\n return Content;\n}\n\n/**\n * Parses and renders markdown on the client. This is generally a bad practice.\n * NOTE: Only run on markdown you trust.\n * It might be worth looking into sandboxing all markdown so that it can't run JS.\n */\nexport default function Markdown(props: { children?: string }) {\n const [child, setChild] = useState(null);\n\n useEffect(() => {\n try {\n parseMarkdown(props.children ?? \"\").then((Content) => {\n setChild();\n });\n } catch {\n setChild(Error Parsing Markdown...);\n }\n }, [props.children]);\n\n return child;\n}\n" }, { "path": "viser/src/viser/client/src/MessageHandler.tsx", "content": "import { CatmullRomLine, CubicBezierLine, Grid, Html } from \"@react-three/drei\";\nimport { useContextBridge } from \"its-fine\";\nimport { notifications } from \"@mantine/notifications\";\n\nimport React, { useContext } from \"react\";\nimport * as THREE from \"three\";\nimport { TextureLoader } from \"three\";\n\nimport { ViewerContext } from \"./App\";\nimport { SceneNode } from \"./SceneTree\";\nimport {\n CameraFrustum,\n CoordinateFrame,\n InstancedAxes,\n GlbAsset,\n OutlinesIfHovered,\n PointCloud,\n} from \"./ThreeAssets\";\nimport {\n FileTransferPart,\n FileTransferStart,\n Message,\n} from \"./WebsocketMessages\";\nimport { PivotControls } from \"@react-three/drei\";\nimport { isTexture, makeThrottledMessageSender } from \"./WebsocketFunctions\";\nimport { isGuiConfig } from \"./ControlPanel/GuiState\";\nimport { useFrame } from \"@react-three/fiber\";\nimport GeneratedGuiContainer from \"./ControlPanel/Generated\";\nimport { Paper, Progress } from \"@mantine/core\";\nimport { IconCheck } from \"@tabler/icons-react\";\nimport { computeT_threeworld_world } from \"./WorldTransformUtils\";\nimport { SplatObject } from \"./Splatting/GaussianSplats\";\n\n/** Convert raw RGB color buffers to linear color buffers. **/\nfunction threeColorBufferFromUint8Buffer(colors: ArrayBuffer) {\n return new THREE.Float32BufferAttribute(\n new Float32Array(new Uint8Array(colors)).map((value) => {\n value = value / 255.0;\n if (value <= 0.04045) {\n return value / 12.92;\n } else {\n return Math.pow((value + 0.055) / 1.055, 2.4);\n }\n }),\n 3,\n );\n}\n\n/** Returns a handler for all incoming messages. */\nfunction useMessageHandler() {\n const viewer = useContext(ViewerContext)!;\n const ContextBridge = useContextBridge();\n\n // We could reduce the redundancy here if we wanted to.\n // https://github.com/nerfstudio-project/viser/issues/39\n const removeSceneNode = viewer.useSceneTree((state) => state.removeSceneNode);\n const resetScene = viewer.useSceneTree((state) => state.resetScene);\n const addSceneNode = viewer.useSceneTree((state) => state.addSceneNode);\n const resetGui = viewer.useGui((state) => state.resetGui);\n const setTheme = viewer.useGui((state) => state.setTheme);\n const setShareUrl = viewer.useGui((state) => state.setShareUrl);\n const addGui = viewer.useGui((state) => state.addGui);\n const addModal = viewer.useGui((state) => state.addModal);\n const removeModal = viewer.useGui((state) => state.removeModal);\n const removeGui = viewer.useGui((state) => state.removeGui);\n const updateGuiProps = viewer.useGui((state) => state.updateGuiProps);\n const setClickable = viewer.useSceneTree((state) => state.setClickable);\n const updateUploadState = viewer.useGui((state) => state.updateUploadState);\n\n // Same as addSceneNode, but make a parent in the form of a dummy coordinate\n // frame if it doesn't exist yet.\n function addSceneNodeMakeParents(node: SceneNode) {\n // Make sure scene node is in attributes.\n const attrs = viewer.nodeAttributesFromName.current;\n attrs[node.name] = {\n overrideVisibility: attrs[node.name]?.overrideVisibility,\n };\n\n // Don't update the pose of the object until we've made a new one!\n attrs[node.name]!.poseUpdateState = \"waitForMakeObject\";\n\n // Make sure parents exists.\n const nodeFromName = viewer.useSceneTree.getState().nodeFromName;\n const parentName = node.name.split(\"/\").slice(0, -1).join(\"/\");\n if (!(parentName in nodeFromName)) {\n addSceneNodeMakeParents(\n new SceneNode(parentName, (ref) => (\n \n )),\n );\n }\n addSceneNode(node);\n }\n\n const fileDownloadHandler = useFileDownloadHandler();\n\n // Return message handler.\n return (message: Message) => {\n if (isGuiConfig(message)) {\n addGui(message);\n return;\n }\n\n switch (message.type) {\n // Set the share URL.\n case \"ShareUrlUpdated\": {\n setShareUrl(message.share_url);\n return;\n }\n // Request a render.\n case \"GetRenderRequestMessage\": {\n viewer.getRenderRequest.current = message;\n viewer.getRenderRequestState.current = \"triggered\";\n return;\n }\n // Set the GUI panel label.\n case \"SetGuiPanelLabelMessage\": {\n viewer.useGui.setState({ label: message.label ?? \"\" });\n return;\n }\n // Configure the theme.\n case \"ThemeConfigurationMessage\": {\n setTheme(message);\n return;\n }\n\n // Run some arbitrary Javascript.\n // This is used for plotting, where the Python server will send over a\n // copy of plotly.min.js for the currently-installed version of plotly.\n case \"RunJavascriptMessage\": {\n eval(message.source);\n return;\n }\n\n // Add a notification.\n case \"NotificationMessage\": {\n if (message.mode === \"show\") {\n notifications.show({\n id: message.id,\n title: message.title,\n message: message.body,\n withCloseButton: message.with_close_button,\n loading: message.loading,\n autoClose: message.auto_close,\n color: message.color ?? undefined,\n });\n } else if (message.mode === \"update\") {\n notifications.update({\n id: message.id,\n title: message.title,\n message: message.body,\n withCloseButton: message.with_close_button,\n loading: message.loading,\n autoClose: message.auto_close,\n color: message.color ?? undefined,\n });\n }\n return;\n }\n\n // Remove a specific notification.\n case \"RemoveNotificationMessage\": {\n notifications.hide(message.id);\n return;\n }\n // Enable/disable whether scene pointer events are sent.\n case \"ScenePointerEnableMessage\": {\n // Update scene click enable state.\n viewer.scenePointerInfo.current!.enabled = message.enable\n ? message.event_type\n : false;\n\n // Update cursor to indicate whether the scene can be clicked.\n viewer.canvasRef.current!.style.cursor = message.enable\n ? \"pointer\"\n : \"auto\";\n return;\n }\n\n // Add a coordinate frame.\n case \"FrameMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(message.name, (ref) => (\n \n )),\n );\n return;\n }\n\n // Add axes to visualize.\n case \"BatchedAxesMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => (\n // Minor naming discrepancy: I think \"batched\" will be clearer to\n // folks on the Python side, but instanced is somewhat more\n // precise.\n \n ),\n undefined,\n undefined,\n undefined,\n // Compute click instance index from instance ID. Each visualized\n // frame has 1 instance for each of 3 line segments.\n (instanceId) => Math.floor(instanceId! / 3),\n ),\n );\n return;\n }\n\n case \"GridMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(message.name, (ref) => (\n \n \n \n )),\n );\n return;\n }\n\n // Add a point cloud.\n case \"PointCloudMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(message.name, (ref) => (\n val / 255.0,\n )}\n />\n )),\n );\n return;\n }\n\n case \"GuiModalMessage\": {\n addModal(message);\n return;\n }\n\n case \"GuiCloseModalMessage\": {\n removeModal(message.id);\n return;\n }\n\n // Add mesh\n case \"SkinnedMeshMessage\":\n case \"MeshMessage\": {\n const geometry = new THREE.BufferGeometry();\n\n const generateGradientMap = (shades: 3 | 5) => {\n const texture = new THREE.DataTexture(\n Uint8Array.from(\n shades == 3\n ? [0, 0, 0, 255, 128, 128, 128, 255, 255, 255, 255, 255]\n : [\n 0, 0, 0, 255, 64, 64, 64, 255, 128, 128, 128, 255, 192, 192,\n 192, 255, 255, 255, 255, 255,\n ],\n ),\n shades,\n 1,\n THREE.RGBAFormat,\n );\n\n texture.needsUpdate = true;\n return texture;\n };\n const standardArgs = {\n color: message.color ?? undefined,\n vertexColors: message.vertex_colors !== null,\n wireframe: message.wireframe,\n transparent: message.opacity !== null,\n opacity: message.opacity ?? 1.0,\n // Flat shading only makes sense for non-wireframe materials.\n flatShading: message.flat_shading && !message.wireframe,\n side: {\n front: THREE.FrontSide,\n back: THREE.BackSide,\n double: THREE.DoubleSide,\n }[message.side],\n };\n const assertUnreachable = (x: never): never => {\n throw new Error(`Should never get here! ${x}`);\n };\n const material =\n message.material == \"standard\" || message.wireframe\n ? new THREE.MeshStandardMaterial(standardArgs)\n : message.material == \"toon3\"\n ? new THREE.MeshToonMaterial({\n gradientMap: generateGradientMap(3),\n ...standardArgs,\n })\n : message.material == \"toon5\"\n ? new THREE.MeshToonMaterial({\n gradientMap: generateGradientMap(5),\n ...standardArgs,\n })\n : assertUnreachable(message.material);\n geometry.setAttribute(\n \"position\",\n new THREE.Float32BufferAttribute(\n new Float32Array(\n message.vertices.buffer.slice(\n message.vertices.byteOffset,\n message.vertices.byteOffset + message.vertices.byteLength,\n ),\n ),\n 3,\n ),\n );\n if (message.vertex_colors !== null) {\n geometry.setAttribute(\n \"color\",\n threeColorBufferFromUint8Buffer(message.vertex_colors),\n );\n }\n\n geometry.setIndex(\n new THREE.Uint32BufferAttribute(\n new Uint32Array(\n message.faces.buffer.slice(\n message.faces.byteOffset,\n message.faces.byteOffset + message.faces.byteLength,\n ),\n ),\n 1,\n ),\n );\n geometry.computeVertexNormals();\n geometry.computeBoundingSphere();\n const cleanupMesh = () => {\n // TODO: we can switch to the react-three-fiber ,\n // , etc components to avoid manual\n // disposal.\n geometry.dispose();\n material.dispose();\n };\n if (message.type === \"MeshMessage\")\n // Normal mesh.\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => {\n return (\n \n \n \n );\n },\n cleanupMesh,\n ),\n );\n else if (message.type === \"SkinnedMeshMessage\") {\n // Skinned mesh.\n const bones: THREE.Bone[] = [];\n for (let i = 0; i < message.bone_wxyzs!.length; i++) {\n bones.push(new THREE.Bone());\n }\n\n const xyzw_quat = new THREE.Quaternion();\n const boneInverses: THREE.Matrix4[] = [];\n viewer.skinnedMeshState.current[message.name] = {\n initialized: false,\n poses: [],\n };\n bones.forEach((bone, i) => {\n const wxyz = message.bone_wxyzs[i];\n const position = message.bone_positions[i];\n xyzw_quat.set(wxyz[1], wxyz[2], wxyz[3], wxyz[0]);\n\n const boneInverse = new THREE.Matrix4();\n boneInverse.makeRotationFromQuaternion(xyzw_quat);\n boneInverse.setPosition(position[0], position[1], position[2]);\n boneInverse.invert();\n boneInverses.push(boneInverse);\n\n bone.quaternion.copy(xyzw_quat);\n bone.position.set(position[0], position[1], position[2]);\n bone.matrixAutoUpdate = false;\n bone.matrixWorldAutoUpdate = false;\n\n viewer.skinnedMeshState.current[message.name].poses.push({\n wxyz: wxyz,\n position: position,\n });\n });\n const skeleton = new THREE.Skeleton(bones, boneInverses);\n\n geometry.setAttribute(\n \"skinIndex\",\n new THREE.Uint16BufferAttribute(\n new Uint16Array(\n message.skin_indices.buffer.slice(\n message.skin_indices.byteOffset,\n message.skin_indices.byteOffset +\n message.skin_indices.byteLength,\n ),\n ),\n 4,\n ),\n );\n geometry.setAttribute(\n \"skinWeight\",\n new THREE.Float32BufferAttribute(\n new Float32Array(\n message.skin_weights!.buffer.slice(\n message.skin_weights!.byteOffset,\n message.skin_weights!.byteOffset +\n message.skin_weights!.byteLength,\n ),\n ),\n 4,\n ),\n );\n\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => {\n return (\n \n \n \n );\n },\n () => {\n delete viewer.skinnedMeshState.current[message.name];\n skeleton.dispose();\n cleanupMesh();\n },\n false,\n // everyFrameCallback: update bone transforms.\n () => {\n const parentNode = viewer.nodeRefFromName.current[message.name];\n if (parentNode === undefined) return;\n\n const state = viewer.skinnedMeshState.current[message.name];\n bones.forEach((bone, i) => {\n if (!state.initialized) {\n parentNode.add(bone);\n }\n const wxyz = state.initialized\n ? state.poses[i].wxyz\n : message.bone_wxyzs[i];\n const position = state.initialized\n ? state.poses[i].position\n : message.bone_positions[i];\n\n xyzw_quat.set(wxyz[1], wxyz[2], wxyz[3], wxyz[0]);\n bone.matrix.makeRotationFromQuaternion(xyzw_quat);\n bone.matrix.setPosition(\n position[0],\n position[1],\n position[2],\n );\n bone.updateMatrixWorld();\n });\n\n if (!state.initialized) {\n state.initialized = true;\n }\n },\n ),\n );\n }\n return;\n }\n // Set the bone poses.\n case \"SetBoneOrientationMessage\": {\n const bonePoses = viewer.skinnedMeshState.current;\n bonePoses[message.name].poses[message.bone_index].wxyz = message.wxyz;\n break;\n }\n case \"SetBonePositionMessage\": {\n const bonePoses = viewer.skinnedMeshState.current;\n bonePoses[message.name].poses[message.bone_index].position =\n message.position;\n break;\n }\n // Add a camera frustum.\n case \"CameraFrustumMessage\": {\n let texture = undefined;\n if (\n message.image_media_type !== null &&\n message.image_binary !== null\n ) {\n const image_url = URL.createObjectURL(\n new Blob([message.image_binary]),\n );\n texture = new TextureLoader().load(image_url, () =>\n URL.revokeObjectURL(image_url),\n );\n }\n\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => (\n \n ),\n () => texture?.dispose(),\n ),\n );\n return;\n }\n case \"TransformControlsMessage\": {\n const name = message.name;\n const sendDragMessage = makeThrottledMessageSender(viewer, 50);\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => (\n e.stopPropagation()}>\n {\n const attrs = viewer.nodeAttributesFromName.current;\n if (attrs[message.name] === undefined) {\n attrs[message.name] = {};\n }\n\n const wxyz = new THREE.Quaternion();\n wxyz.setFromRotationMatrix(l);\n const position = new THREE.Vector3().setFromMatrixPosition(\n l,\n );\n\n const nodeAttributes = attrs[message.name]!;\n nodeAttributes.wxyz = [wxyz.w, wxyz.x, wxyz.y, wxyz.z];\n nodeAttributes.position = position.toArray();\n sendDragMessage({\n type: \"TransformControlsUpdateMessage\",\n name: name,\n wxyz: nodeAttributes.wxyz,\n position: nodeAttributes.position,\n });\n }}\n />\n \n ),\n undefined,\n true, // unmountWhenInvisible\n ),\n );\n return;\n }\n case \"SetCameraLookAtMessage\": {\n const cameraControls = viewer.cameraControlRef.current!;\n\n const T_threeworld_world = computeT_threeworld_world(viewer);\n const target = new THREE.Vector3(\n message.look_at[0],\n message.look_at[1],\n message.look_at[2],\n );\n target.applyMatrix4(T_threeworld_world);\n cameraControls.setTarget(target.x, target.y, target.z, false);\n return;\n }\n case \"SetCameraUpDirectionMessage\": {\n const camera = viewer.cameraRef.current!;\n const cameraControls = viewer.cameraControlRef.current!;\n const T_threeworld_world = computeT_threeworld_world(viewer);\n const updir = new THREE.Vector3(\n message.position[0],\n message.position[1],\n message.position[2],\n )\n .normalize()\n .applyQuaternion(\n new THREE.Quaternion().setFromRotationMatrix(T_threeworld_world),\n );\n camera.up.set(updir.x, updir.y, updir.z);\n\n // Back up position.\n const prevPosition = new THREE.Vector3();\n cameraControls.getPosition(prevPosition);\n\n cameraControls.updateCameraUp();\n\n // Restore position, which can get unexpectedly mutated in updateCameraUp().\n cameraControls.setPosition(\n prevPosition.x,\n prevPosition.y,\n prevPosition.z,\n false,\n );\n return;\n }\n case \"SetCameraPositionMessage\": {\n const cameraControls = viewer.cameraControlRef.current!;\n\n // Set the camera position. Due to the look-at, note that this will\n // shift the orientation as-well.\n const position_cmd = new THREE.Vector3(\n message.position[0],\n message.position[1],\n message.position[2],\n );\n\n const T_threeworld_world = computeT_threeworld_world(viewer);\n position_cmd.applyMatrix4(T_threeworld_world);\n\n cameraControls.setPosition(\n position_cmd.x,\n position_cmd.y,\n position_cmd.z,\n );\n return;\n }\n case \"SetCameraFovMessage\": {\n const camera = viewer.cameraRef.current!;\n // tan(fov / 2.0) = 0.5 * film height / focal length\n // focal length = 0.5 * film height / tan(fov / 2.0)\n camera.setFocalLength(\n (0.5 * camera.getFilmHeight()) / Math.tan(message.fov / 2.0),\n );\n viewer.sendCameraRef.current !== null && viewer.sendCameraRef.current();\n return;\n }\n case \"SetOrientationMessage\": {\n const attr = viewer.nodeAttributesFromName.current;\n if (attr[message.name] === undefined) attr[message.name] = {};\n attr[message.name]!.wxyz = message.wxyz;\n if (attr[message.name]!.poseUpdateState == \"updated\")\n attr[message.name]!.poseUpdateState = \"needsUpdate\";\n break;\n }\n case \"SetPositionMessage\": {\n const attr = viewer.nodeAttributesFromName.current;\n if (attr[message.name] === undefined) attr[message.name] = {};\n attr[message.name]!.position = message.position;\n if (attr[message.name]!.poseUpdateState == \"updated\")\n attr[message.name]!.poseUpdateState = \"needsUpdate\";\n break;\n }\n case \"SetSceneNodeVisibilityMessage\": {\n const attr = viewer.nodeAttributesFromName.current;\n if (attr[message.name] === undefined) attr[message.name] = {};\n attr[message.name]!.visibility = message.visible;\n break;\n }\n // Add a background image.\n case \"BackgroundImageMessage\": {\n const rgb_url = URL.createObjectURL(\n new Blob([message.rgb_bytes], {\n type: message.media_type,\n }),\n );\n new TextureLoader().load(rgb_url, (texture) => {\n URL.revokeObjectURL(rgb_url);\n const oldBackgroundTexture =\n viewer.backgroundMaterialRef.current!.uniforms.colorMap.value;\n viewer.backgroundMaterialRef.current!.uniforms.colorMap.value =\n texture;\n if (isTexture(oldBackgroundTexture)) oldBackgroundTexture.dispose();\n\n viewer.useGui.setState({ backgroundAvailable: true });\n });\n viewer.backgroundMaterialRef.current!.uniforms.enabled.value = true;\n viewer.backgroundMaterialRef.current!.uniforms.hasDepth.value =\n message.depth_bytes !== null;\n\n if (message.depth_bytes !== null) {\n // If depth is available set the texture\n const depth_url = URL.createObjectURL(\n new Blob([message.depth_bytes], {\n type: message.media_type,\n }),\n );\n new TextureLoader().load(depth_url, (texture) => {\n URL.revokeObjectURL(depth_url);\n const oldDepthTexture =\n viewer.backgroundMaterialRef.current?.uniforms.depthMap.value;\n viewer.backgroundMaterialRef.current!.uniforms.depthMap.value =\n texture;\n if (isTexture(oldDepthTexture)) oldDepthTexture.dispose();\n });\n }\n return;\n }\n // Add a 2D label.\n case \"LabelMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => {\n // We wrap with because Html doesn't implement THREE.Object3D.\n return (\n \n \n \n \n {message.text}\n \n \n \n \n );\n },\n undefined,\n true,\n ),\n );\n return;\n }\n case \"Gui3DMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => {\n // We wrap with because Html doesn't implement\n // THREE.Object3D. The initial position is intended to be\n // off-screen; it will be overwritten with the actual position\n // after the component is mounted.\n return (\n \n \n \n {\n evt.stopPropagation();\n }}\n >\n \n \n \n \n \n \n \n );\n },\n undefined,\n true,\n ),\n );\n return;\n }\n // Add an image.\n case \"ImageMessage\": {\n // This current implementation may flicker when the image is updated,\n // because the texture is not necessarily done loading before the\n // component is mounted. We could fix this by passing an `onLoad`\n // callback into `TextureLoader`, but this would require work because\n // `addSceneNodeMakeParents` needs to be called immediately: it\n // overwrites position/wxyz attributes, and we don't want this to\n // happen after later messages are received.\n const image_url = URL.createObjectURL(\n new Blob([message.data], {\n type: message.media_type,\n }),\n );\n const texture = new TextureLoader().load(\n image_url,\n () => URL.revokeObjectURL(image_url), // Revoke URL on load.\n );\n addSceneNodeMakeParents(\n new SceneNode(\n message.name,\n (ref) => {\n return (\n \n \n \n \n \n \n \n );\n },\n () => texture.dispose(),\n ),\n );\n return;\n }\n // Remove a scene node and its children by name.\n case \"RemoveSceneNodeMessage\": {\n console.log(\"Removing scene node:\", message.name);\n removeSceneNode(message.name);\n const attrs = viewer.nodeAttributesFromName.current;\n delete attrs[message.name];\n return;\n }\n // Set the clickability of a particular scene node.\n case \"SetSceneNodeClickableMessage\": {\n // This setTimeout is totally unnecessary, but can help surface some race\n // conditions.\n setTimeout(() => setClickable(message.name, message.clickable), 50);\n return;\n }\n // Reset the entire scene, removing all scene nodes.\n case \"ResetSceneMessage\": {\n resetScene();\n\n const oldBackground = viewer.sceneRef.current?.background;\n viewer.sceneRef.current!.background = null;\n if (isTexture(oldBackground)) oldBackground.dispose();\n\n viewer.useGui.setState({ backgroundAvailable: false });\n // Disable the depth texture rendering\n viewer.backgroundMaterialRef.current!.uniforms.enabled.value = false;\n return;\n }\n // Reset the GUI state.\n case \"ResetGuiMessage\": {\n resetGui();\n return;\n }\n // Update props of a GUI component\n case \"GuiUpdateMessage\": {\n updateGuiProps(message.id, message.updates);\n return;\n }\n // Remove a GUI input.\n case \"GuiRemoveMessage\": {\n removeGui(message.id);\n return;\n }\n // Add a glTF/GLB asset.\n case \"GlbMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(message.name, (ref) => {\n return (\n \n );\n }),\n );\n return;\n }\n case \"CatmullRomSplineMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(message.name, (ref) => {\n return (\n \n \n \n );\n }),\n );\n return;\n }\n case \"CubicBezierSplineMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(message.name, (ref) => {\n return (\n \n {[...Array(message.positions.length - 1).keys()].map((i) => (\n \n ))}\n \n );\n }),\n );\n return;\n }\n case \"GaussianSplatsMessage\": {\n addSceneNodeMakeParents(\n new SceneNode(message.name, (ref) => {\n return (\n \n );\n }),\n );\n return;\n }\n case \"FileTransferStart\":\n case \"FileTransferPart\": {\n fileDownloadHandler(message);\n return;\n }\n case \"FileTransferPartAck\": {\n updateUploadState({\n componentId: message.source_component_id!,\n uploadedBytes: message.transferred_bytes,\n totalBytes: message.total_bytes,\n });\n return;\n }\n default: {\n console.log(\"Received message did not match any known types:\", message);\n return;\n }\n }\n };\n}\n\nfunction useFileDownloadHandler() {\n const downloadStatesRef = React.useRef<{\n [uuid: string]: {\n metadata: FileTransferStart;\n notificationId: string;\n parts: Uint8Array[];\n bytesDownloaded: number;\n displayFilesize: string;\n };\n }>({});\n\n return (message: FileTransferStart | FileTransferPart) => {\n const notificationId = \"download-\" + message.transfer_uuid;\n\n // Create or update download state.\n switch (message.type) {\n case \"FileTransferStart\": {\n let displaySize = message.size_bytes;\n const displayUnits = [\"B\", \"K\", \"M\", \"G\", \"T\", \"P\"];\n let displayUnitIndex = 0;\n while (\n displaySize >= 100 &&\n displayUnitIndex < displayUnits.length - 1\n ) {\n displaySize /= 1024;\n displayUnitIndex += 1;\n }\n downloadStatesRef.current[message.transfer_uuid] = {\n metadata: message,\n notificationId: notificationId,\n parts: [],\n bytesDownloaded: 0,\n displayFilesize: `${displaySize.toFixed(1)}${\n displayUnits[displayUnitIndex]\n }`,\n };\n break;\n }\n case \"FileTransferPart\": {\n const downloadState = downloadStatesRef.current[message.transfer_uuid];\n if (message.part != downloadState.parts.length) {\n console.error(\n \"A file download message was dropped; this should never happen!\",\n );\n }\n downloadState.parts.push(message.content);\n downloadState.bytesDownloaded += message.content.length;\n break;\n }\n }\n\n // Show notification.\n const downloadState = downloadStatesRef.current[message.transfer_uuid];\n const progressValue =\n (100.0 * downloadState.bytesDownloaded) /\n downloadState.metadata.size_bytes;\n const isDone =\n downloadState.bytesDownloaded == downloadState.metadata.size_bytes;\n\n (downloadState.bytesDownloaded == 0\n ? notifications.show\n : notifications.update)({\n title:\n (isDone ? \"Downloaded \" : \"Downloading \") +\n `${downloadState.metadata.filename} (${downloadState.displayFilesize})`,\n message: ,\n id: notificationId,\n autoClose: isDone,\n withCloseButton: isDone,\n loading: !isDone,\n icon: isDone ? : undefined,\n });\n\n // If done: download file and clear state.\n if (isDone) {\n const link = document.createElement(\"a\");\n link.href = window.URL.createObjectURL(\n new Blob(downloadState.parts, {\n type: downloadState.metadata.mime_type,\n }),\n );\n link.download = downloadState.metadata.filename;\n link.click();\n link.remove();\n delete downloadStatesRef.current[message.transfer_uuid];\n }\n };\n}\n\nexport function FrameSynchronizedMessageHandler() {\n const handleMessage = useMessageHandler();\n const viewer = useContext(ViewerContext)!;\n const messageQueueRef = viewer.messageQueueRef;\n\n useFrame(() => {\n // Send a render along if it was requested!\n if (viewer.getRenderRequestState.current === \"triggered\") {\n viewer.getRenderRequestState.current = \"pause\";\n } else if (viewer.getRenderRequestState.current === \"pause\") {\n const sourceCanvas = viewer.canvasRef.current!;\n\n const targetWidth = viewer.getRenderRequest.current!.width;\n const targetHeight = viewer.getRenderRequest.current!.height;\n\n // We'll save a render to an intermediate canvas with the requested dimensions.\n const renderBufferCanvas = new OffscreenCanvas(targetWidth, targetHeight);\n const ctx = renderBufferCanvas.getContext(\"2d\")!;\n ctx.reset();\n // Use a white background for JPEGs, which don't have an alpha channel.\n if (viewer.getRenderRequest.current?.format === \"image/jpeg\") {\n ctx.fillStyle = \"white\";\n ctx.fillRect(0, 0, renderBufferCanvas.width, renderBufferCanvas.height);\n }\n\n // Determine offsets for the source canvas. We'll always center our renders.\n // https://developer.mozilla.org/en-US/docs/Web/API/CanvasRenderingContext2D/drawImage\n let sourceWidth = sourceCanvas.width;\n let sourceHeight = sourceCanvas.height;\n\n const sourceAspect = sourceWidth / sourceHeight;\n const targetAspect = targetWidth / targetHeight;\n\n if (sourceAspect > targetAspect) {\n // The source is wider than the target.\n // We need to shrink the width.\n sourceWidth = Math.round(targetAspect * sourceHeight);\n } else if (sourceAspect < targetAspect) {\n // The source is narrower than the target.\n // We need to shrink the height.\n sourceHeight = Math.round(sourceWidth / targetAspect);\n }\n\n console.log(\n `Sending render; requested aspect ratio was ${targetAspect} (dimensinos: ${targetWidth}/${targetHeight}), copying from aspect ratio ${\n sourceWidth / sourceHeight\n } (dimensions: ${sourceWidth}/${sourceHeight}).`,\n );\n\n ctx.drawImage(\n sourceCanvas,\n (sourceCanvas.width - sourceWidth) / 2.0,\n (sourceCanvas.height - sourceHeight) / 2.0,\n sourceWidth,\n sourceHeight,\n 0,\n 0,\n targetWidth,\n targetHeight,\n );\n\n viewer.getRenderRequestState.current = \"in_progress\";\n\n // Encode the image, the send it.\n renderBufferCanvas\n .convertToBlob({\n type: viewer.getRenderRequest.current!.format,\n quality: viewer.getRenderRequest.current!.quality / 100.0,\n })\n .then(async (blob) => {\n if (blob === null) {\n console.error(\"Render failed\");\n viewer.getRenderRequestState.current = \"ready\";\n return;\n }\n const payload = new Uint8Array(await blob.arrayBuffer());\n viewer.sendMessageRef.current({\n type: \"GetRenderResponseMessage\",\n payload: payload,\n });\n viewer.getRenderRequestState.current = \"ready\";\n });\n }\n\n // Handle messages, but only if we're not trying to render something.\n if (viewer.getRenderRequestState.current === \"ready\") {\n // Handle messages before every frame.\n // Place this directly in ws.onmessage can cause race conditions!\n //\n // If a render is requested, note that we don't handle any more messages\n // until the render is done.\n const requestRenderIndex = messageQueueRef.current.findIndex(\n (message) => message.type === \"GetRenderRequestMessage\",\n );\n const numMessages =\n requestRenderIndex !== -1\n ? requestRenderIndex + 1\n : messageQueueRef.current.length;\n const processBatch = messageQueueRef.current.splice(0, numMessages);\n processBatch.forEach(handleMessage);\n }\n });\n\n return null;\n}\n" }, { "path": "viser/src/viser/client/src/Modal.tsx", "content": "import { ViewerContext } from \"./App\";\nimport { GuiModalMessage } from \"./WebsocketMessages\";\nimport GeneratedGuiContainer from \"./ControlPanel/Generated\";\nimport { Modal } from \"@mantine/core\";\nimport { useContext } from \"react\";\n\nexport function ViserModal() {\n const viewer = useContext(ViewerContext)!;\n\n const modalList = viewer.useGui((state) => state.modals);\n const modals = modalList.map((conf, index) => {\n return ;\n });\n\n return modals;\n}\n\nfunction GeneratedModal({\n conf,\n index,\n}: {\n conf: GuiModalMessage;\n index: number;\n}) {\n return (\n {\n // To make memory management easier, we should only close modals from\n // the server.\n // Otherwise, the client would need to communicate to the server that\n // the modal was deleted and contained GUI elements were cleared.\n }}\n withCloseButton={false}\n centered\n zIndex={100 + index}\n >\n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/Outlines.tsx", "content": "/** This is a modified version of drei's component. The primary\n * change is to add support for ref forwarding. https://github.com/pmndrs/drei\n * */\n\nimport * as THREE from \"three\";\nimport * as React from \"react\";\nimport {\n extend,\n applyProps,\n ReactThreeFiber,\n useThree,\n} from \"@react-three/fiber\";\nimport { toCreasedNormals } from \"three-stdlib\";\nimport { version } from \"@react-three/drei/helpers/constants\";\nimport { shaderMaterial } from \"@react-three/drei\";\n\nconst OutlinesMaterial = /* @__PURE__ */ shaderMaterial(\n {\n screenspace: false,\n color: /* @__PURE__ */ new THREE.Color(\"black\"),\n opacity: 1,\n thickness: 0.05,\n size: /* @__PURE__ */ new THREE.Vector2(),\n },\n `#include \n #include \n #include \n uniform float thickness;\n uniform float screenspace;\n uniform vec2 size;\n void main() {\n #if defined (USE_SKINNING)\n\t #include \n #include \n #include \n #include \n #include \n #endif\n #include \n\t #include \n\t #include \n #include \n vec4 tNormal = vec4(normal, 0.0);\n vec4 tPosition = vec4(transformed, 1.0);\n #ifdef USE_INSTANCING\n tNormal = instanceMatrix * tNormal;\n tPosition = instanceMatrix * tPosition;\n #endif\n if (screenspace == 0.0) {\n vec3 newPosition = tPosition.xyz + tNormal.xyz * thickness;\n gl_Position = projectionMatrix * modelViewMatrix * vec4(newPosition, 1.0);\n } else {\n vec4 clipPosition = projectionMatrix * modelViewMatrix * tPosition;\n vec4 clipNormal = projectionMatrix * modelViewMatrix * tNormal;\n vec2 offset = normalize(clipNormal.xy) * thickness / size * clipPosition.w * 2.0;\n clipPosition.xy += offset;\n gl_Position = clipPosition;\n }\n }`,\n `uniform vec3 color;\n uniform float opacity;\n void main(){\n gl_FragColor = vec4(color, opacity);\n #include \n #include <${version >= 154 ? \"colorspace_fragment\" : \"encodings_fragment\"}>\n }`,\n);\n\ntype OutlinesProps = JSX.IntrinsicElements[\"group\"] & {\n /** Outline color, default: black */\n color?: ReactThreeFiber.Color;\n /** Line thickness is independent of zoom, default: false */\n screenspace?: boolean;\n /** Outline opacity, default: 1 */\n opacity?: number;\n /** Outline transparency, default: false */\n transparent?: boolean;\n /** Outline thickness, default 0.05 */\n thickness?: number;\n /** Geometry crease angle (0 === no crease), default: Math.PI */\n angle?: number;\n toneMapped?: boolean;\n polygonOffset?: boolean;\n polygonOffsetFactor?: number;\n renderOrder?: number;\n};\n\nexport const Outlines = React.forwardRef(\n function Outlines(\n {\n color = \"black\",\n opacity = 1,\n transparent = false,\n screenspace = false,\n toneMapped = true,\n polygonOffset = false,\n polygonOffsetFactor = 0,\n renderOrder = 0,\n thickness = 0.05,\n angle = Math.PI,\n ...props\n },\n ref,\n ) {\n const localRef = React.useRef(null);\n\n const [material] = React.useState(\n () => new OutlinesMaterial({ side: THREE.BackSide }),\n );\n const gl = useThree((state) => state.gl);\n const contextSize = gl.getDrawingBufferSize(new THREE.Vector2());\n React.useMemo(() => extend({ OutlinesMaterial }), []);\n\n const oldAngle = React.useRef(0);\n const oldGeometry = React.useRef();\n React.useLayoutEffect(() => {\n const group = localRef.current;\n if (!group) return;\n\n const parent = group.parent as THREE.Mesh &\n THREE.SkinnedMesh &\n THREE.InstancedMesh;\n if (parent && parent.geometry) {\n if (\n oldAngle.current !== angle ||\n oldGeometry.current !== parent.geometry\n ) {\n oldAngle.current = angle;\n oldGeometry.current = parent.geometry;\n\n // Remove old mesh\n let mesh = group.children[0] as any;\n if (mesh) {\n if (angle) mesh.geometry.dispose();\n group.remove(mesh);\n }\n\n if (parent.skeleton) {\n mesh = new THREE.SkinnedMesh();\n mesh.material = material;\n mesh.bind(parent.skeleton, parent.bindMatrix);\n group.add(mesh);\n } else if (parent.isInstancedMesh) {\n mesh = new THREE.InstancedMesh(\n parent.geometry,\n material,\n parent.count,\n );\n mesh.instanceMatrix = parent.instanceMatrix;\n group.add(mesh);\n } else {\n mesh = new THREE.Mesh();\n mesh.material = material;\n group.add(mesh);\n }\n mesh.geometry = angle\n ? toCreasedNormals(parent.geometry, angle)\n : parent.geometry;\n }\n }\n });\n\n React.useLayoutEffect(() => {\n const group = localRef.current;\n if (!group) return;\n\n const mesh = group.children[0] as THREE.Mesh<\n THREE.BufferGeometry,\n THREE.Material\n >;\n if (mesh) {\n mesh.renderOrder = renderOrder;\n applyProps(mesh.material as any, {\n transparent,\n thickness,\n color,\n opacity,\n size: contextSize,\n screenspace,\n toneMapped,\n polygonOffset,\n polygonOffsetFactor,\n });\n }\n });\n\n React.useEffect(() => {\n return () => {\n // Dispose everything on unmount\n const group = localRef.current;\n if (!group) return;\n\n const mesh = group.children[0] as THREE.Mesh<\n THREE.BufferGeometry,\n THREE.Material\n >;\n if (mesh) {\n if (angle) mesh.geometry.dispose();\n group.remove(mesh);\n }\n };\n }, []);\n\n return (\n {\n localRef.current = obj;\n if (typeof ref === \"function\") ref(obj!);\n else if (ref) ref.current = obj;\n }}\n {...props}\n />\n );\n },\n);\n" }, { "path": "viser/src/viser/client/src/SceneTree.tsx", "content": "import { useCursor } from \"@react-three/drei\";\nimport { createPortal, useFrame } from \"@react-three/fiber\";\nimport React from \"react\";\nimport * as THREE from \"three\";\n\nimport { ViewerContext } from \"./App\";\nimport { useThrottledMessageSender } from \"./WebsocketFunctions\";\nimport { Html } from \"@react-three/drei\";\nimport { immerable } from \"immer\";\nimport { useSceneTreeState } from \"./SceneTreeState\";\nimport { ErrorBoundary } from \"react-error-boundary\";\nimport { rayToViserCoords } from \"./WorldTransformUtils\";\nimport { HoverableContext } from \"./ThreeAssets\";\nimport { opencvXyFromPointerXy } from \"./ClickUtils\";\n\nexport type MakeObject = (\n ref: React.Ref,\n) => React.ReactNode;\n\n/** Scenes will consist of nodes, which form a tree. */\nexport class SceneNode {\n [immerable] = true;\n\n public children: string[];\n public clickable: boolean;\n\n constructor(\n public readonly name: string,\n public readonly makeObject: MakeObject,\n public readonly cleanup?: () => void,\n /** unmountWhenInvisible is used to unmount components when they\n * should be hidden.\n *\n * https://github.com/pmndrs/drei/issues/1323\n */\n public readonly unmountWhenInvisible?: boolean,\n public readonly everyFrameCallback?: () => void,\n /** For click events on instanced nodes, like batched axes, we want to keep track of which. */\n public readonly computeClickInstanceIndexFromInstanceId?: (\n instanceId: number | undefined,\n ) => number | null,\n ) {\n this.children = [];\n this.clickable = false;\n }\n}\n\n/** Type corresponding to a zustand-style useSceneTree hook. */\nexport type UseSceneTree = ReturnType;\n\nfunction SceneNodeThreeChildren(props: {\n name: string;\n parent: THREE.Object3D;\n}) {\n const viewer = React.useContext(ViewerContext)!;\n\n const [children, setChildren] = React.useState(\n viewer.useSceneTree.getState().nodeFromName[props.name]?.children ?? [],\n );\n\n React.useEffect(() => {\n let updateQueued = false;\n return viewer.useSceneTree.subscribe((state) => {\n // Do nothing if an update is already queued.\n if (updateQueued) return;\n\n // Do nothing if children haven't changed.\n const newChildren = state.nodeFromName[props.name]?.children;\n if (\n newChildren === undefined ||\n newChildren === children || // Note that this won't check for elementwise equality!\n (newChildren.length === 0 && children.length == 0)\n )\n return;\n\n // Queue a (throttled) children update.\n updateQueued = true;\n setTimeout(\n () => {\n updateQueued = false;\n const newChildren =\n viewer.useSceneTree.getState().nodeFromName[props.name]!.children!;\n setChildren(newChildren);\n },\n // Throttle more when we have a lot of children...\n newChildren.length <= 16 ? 10 : newChildren.length <= 128 ? 50 : 200,\n );\n });\n }, []);\n\n // Create a group of children inside of the parent object.\n return createPortal(\n \n {children &&\n children.map((child_id) => (\n \n ))}\n \n ,\n props.parent,\n );\n}\n\n/** Component for updating attributes of a scene node. */\nfunction SceneNodeLabel(props: { name: string }) {\n const viewer = React.useContext(ViewerContext)!;\n const labelVisible = viewer.useSceneTree(\n (state) => state.labelVisibleFromName[props.name],\n );\n return labelVisible ? (\n \n \n {props.name}\n \n \n ) : null;\n}\n\n/** Component containing the three.js object and children for a particular scene node. */\nexport function SceneNodeThreeObject(props: {\n name: string;\n parent: THREE.Object3D | null;\n}) {\n const viewer = React.useContext(ViewerContext)!;\n const makeObject = viewer.useSceneTree(\n (state) => state.nodeFromName[props.name]?.makeObject,\n );\n const unmountWhenInvisible = viewer.useSceneTree(\n (state) => state.nodeFromName[props.name]?.unmountWhenInvisible,\n );\n const everyFrameCallback = viewer.useSceneTree(\n (state) => state.nodeFromName[props.name]?.everyFrameCallback,\n );\n const computeClickInstanceIndexFromInstanceId = viewer.useSceneTree(\n (state) =>\n state.nodeFromName[props.name]?.computeClickInstanceIndexFromInstanceId,\n );\n const [unmount, setUnmount] = React.useState(false);\n const clickable =\n viewer.useSceneTree((state) => state.nodeFromName[props.name]?.clickable) ??\n false;\n const [obj, setRef] = React.useState(null);\n\n // Update global registry of node objects.\n // This is used for updating bone transforms in skinned meshes.\n React.useEffect(() => {\n if (obj !== null) viewer.nodeRefFromName.current[props.name] = obj;\n }, [obj]);\n\n // Create object + children.\n //\n // For not-fully-understood reasons, wrapping makeObject with useMemo() fixes\n // stability issues (eg breaking runtime errors) associated with\n // PivotControls.\n const objNode = React.useMemo(() => {\n if (makeObject === undefined) return null;\n\n // Pose will need to be updated.\n const attrs = viewer.nodeAttributesFromName.current;\n if (!(props.name in attrs)) {\n attrs[props.name] = {};\n }\n attrs[props.name]!.poseUpdateState = \"needsUpdate\";\n\n return makeObject(setRef);\n }, [makeObject]);\n const children =\n obj === null ? null : (\n \n );\n\n // Helper for transient visibility checks. Checks the .visible attribute of\n // both this object and ancestors.\n //\n // This is used for (1) suppressing click events and (2) unmounting when\n // unmountWhenInvisible is true. The latter is used for components.\n function isDisplayed() {\n // We avoid checking obj.visible because obj may be unmounted when\n // unmountWhenInvisible=true.\n const attrs = viewer.nodeAttributesFromName.current[props.name];\n const visibility =\n (attrs?.overrideVisibility === undefined\n ? attrs?.visibility\n : attrs.overrideVisibility) ?? true;\n if (visibility === false) return false;\n if (props.parent === null) return true;\n\n // Check visibility of parents + ancestors.\n let visible = props.parent.visible;\n if (visible) {\n props.parent.traverseAncestors((ancestor) => {\n visible = visible && ancestor.visible;\n });\n }\n return visible;\n }\n\n // Pose needs to be updated whenever component is remounted.\n React.useEffect(() => {\n const attrs = viewer.nodeAttributesFromName.current[props.name];\n if (attrs !== undefined) attrs.poseUpdateState = \"needsUpdate\";\n });\n\n // Update attributes on a per-frame basis. Currently does redundant work,\n // although this shouldn't be a bottleneck.\n useFrame(\n () => {\n const attrs = viewer.nodeAttributesFromName.current[props.name];\n everyFrameCallback && everyFrameCallback();\n\n // Unmount when invisible.\n // Examples: components, PivotControls.\n //\n // This is a workaround for situations where just setting `visible` doesn't\n // work (like ), or to prevent invisible elements from being\n // interacted with ().\n //\n // https://github.com/pmndrs/drei/issues/1323\n if (unmountWhenInvisible) {\n const displayed = isDisplayed();\n if (displayed && unmount) {\n if (obj !== null) obj.visible = false;\n setUnmount(false);\n }\n if (!displayed && !unmount) {\n setUnmount(true);\n }\n }\n\n if (obj === null) return;\n if (attrs === undefined) return;\n\n const visibility =\n (attrs?.overrideVisibility === undefined\n ? attrs?.visibility\n : attrs.overrideVisibility) ?? true;\n obj.visible = visibility;\n\n if (attrs.poseUpdateState == \"needsUpdate\") {\n attrs.poseUpdateState = \"updated\";\n const wxyz = attrs.wxyz;\n if (wxyz !== undefined) {\n obj.quaternion.set(wxyz[1], wxyz[2], wxyz[3], wxyz[0]);\n }\n const position = attrs.position;\n if (position !== undefined) {\n obj.position.set(position[0], position[1], position[2]);\n }\n\n // Update matrices if necessary. This is necessary for PivotControls.\n if (!obj.matrixAutoUpdate) obj.updateMatrix();\n if (!obj.matrixWorldAutoUpdate) obj.updateMatrixWorld();\n }\n },\n // Other useFrame hooks may depend on transforms + visibility. So it's best\n // to call this hook early.\n -10000,\n );\n\n // Clicking logic.\n const sendClicksThrottled = useThrottledMessageSender(50);\n const [hovered, setHovered] = React.useState(false);\n useCursor(hovered);\n const hoveredRef = React.useRef(false);\n if (!clickable && hovered) setHovered(false);\n\n const dragInfo = React.useRef({\n dragging: false,\n startClientX: 0,\n startClientY: 0,\n });\n\n if (objNode === undefined || unmount) {\n return <>{children};\n } else if (clickable) {\n return (\n <>\n {\n // This sometimes (but very rarely) catches a race condition when\n // we remove scene nodes. I would guess it's related to portaling,\n // but the issue is unnoticeable with ErrorBoundary in-place so not\n // debugging further for now...\n console.error(\n \"There was an error rendering a scene node object:\",\n objNode,\n );\n return null;\n }}\n >\n {\n if (!isDisplayed()) return;\n e.stopPropagation();\n const state = dragInfo.current;\n const canvasBbox =\n viewer.canvasRef.current!.getBoundingClientRect();\n state.startClientX = e.clientX - canvasBbox.left;\n state.startClientY = e.clientY - canvasBbox.top;\n state.dragging = false;\n }}\n onPointerMove={(e) => {\n if (!isDisplayed()) return;\n e.stopPropagation();\n const state = dragInfo.current;\n const canvasBbox =\n viewer.canvasRef.current!.getBoundingClientRect();\n const deltaX = e.clientX - canvasBbox.left - state.startClientX;\n const deltaY = e.clientY - canvasBbox.top - state.startClientY;\n // Minimum motion.\n if (Math.abs(deltaX) <= 3 && Math.abs(deltaY) <= 3) return;\n state.dragging = true;\n }}\n onPointerUp={(e) => {\n if (!isDisplayed()) return;\n e.stopPropagation();\n const state = dragInfo.current;\n if (state.dragging) return;\n // Convert ray to viser coordinates.\n const ray = rayToViserCoords(viewer, e.ray);\n\n // Send OpenCV image coordinates to the server (normalized).\n const canvasBbox =\n viewer.canvasRef.current!.getBoundingClientRect();\n const mouseVectorOpenCV = opencvXyFromPointerXy(viewer, [\n e.clientX - canvasBbox.left,\n e.clientY - canvasBbox.top,\n ]);\n\n sendClicksThrottled({\n type: \"SceneNodeClickMessage\",\n name: props.name,\n instance_index:\n computeClickInstanceIndexFromInstanceId === undefined\n ? null\n : computeClickInstanceIndexFromInstanceId(e.instanceId),\n // Note that the threejs up is +Y, but we expose a +Z up.\n ray_origin: [ray.origin.x, ray.origin.y, ray.origin.z],\n ray_direction: [\n ray.direction.x,\n ray.direction.y,\n ray.direction.z,\n ],\n screen_pos: [mouseVectorOpenCV.x, mouseVectorOpenCV.y],\n });\n }}\n onPointerOver={(e) => {\n if (!isDisplayed()) return;\n e.stopPropagation();\n setHovered(true);\n hoveredRef.current = true;\n }}\n onPointerOut={() => {\n if (!isDisplayed()) return;\n setHovered(false);\n hoveredRef.current = false;\n }}\n >\n \n {objNode}\n \n \n {children}\n \n \n );\n } else {\n return (\n <>\n {/* This does nothing, but switching between clickable vs not\n causes strange transform behavior without it. */}\n {objNode}\n {children}\n \n );\n }\n}\n" }, { "path": "viser/src/viser/client/src/SceneTreeState.tsx", "content": "import React from \"react\";\nimport { MakeObject, SceneNode } from \"./SceneTree\";\nimport { CoordinateFrame } from \"./ThreeAssets\";\nimport * as THREE from \"three\";\nimport { create } from \"zustand\";\nimport { subscribeWithSelector } from \"zustand/middleware\";\nimport { immer } from \"zustand/middleware/immer\";\n\ninterface SceneTreeState {\n nodeFromName: { [name: string]: undefined | SceneNode };\n // Putting this into SceneNode makes the scene tree table much harder to implement.\n labelVisibleFromName: { [name: string]: boolean };\n}\nexport interface SceneTreeActions extends SceneTreeState {\n setClickable(name: string, clickable: boolean): void;\n addSceneNode(nodes: SceneNode): void;\n removeSceneNode(name: string): void;\n resetScene(): void;\n setLabelVisibility(name: string, labelVisibility: boolean): void;\n}\n\n// Create default scene tree state.\nconst rootAxesTemplate: MakeObject = (ref) => (\n \n);\n\nconst rootNodeTemplate = new SceneNode(\"\", (ref) => (\n \n)) as SceneNode;\n\nconst rootAxesNode = new SceneNode(\n \"/WorldAxes\",\n rootAxesTemplate,\n) as SceneNode;\nrootNodeTemplate.children.push(\"/WorldAxes\");\n\n/** Declare a scene state, and return a hook for accessing it. Note that we put\neffort into avoiding a global state! */\nexport function useSceneTreeState(\n nodeRefFromName: React.MutableRefObject<{\n [name: string]: undefined | THREE.Object3D;\n }>,\n) {\n return React.useState(() =>\n create(\n subscribeWithSelector(\n immer((set) => ({\n nodeFromName: { \"\": rootNodeTemplate, \"/WorldAxes\": rootAxesNode },\n labelVisibleFromName: {},\n setClickable: (name, clickable) =>\n set((state) => {\n const node = state.nodeFromName[name];\n if (node !== undefined) node.clickable = clickable;\n }),\n addSceneNode: (node) =>\n set((state) => {\n const existingNode = state.nodeFromName[node.name];\n if (existingNode !== undefined) {\n // Node already exists.\n delete nodeRefFromName.current[node.name];\n existingNode.cleanup && existingNode.cleanup(); // Free resources.\n state.nodeFromName[node.name] = {\n ...node,\n children: existingNode.children,\n };\n } else {\n // Node doesn't exist yet!\n // TODO: this assumes the parent exists. We could probably merge this with addSceneNodeMakeParents.\n const parent_name = node.name.split(\"/\").slice(0, -1).join(\"/\");\n state.nodeFromName[node.name] = node;\n state.nodeFromName[parent_name]!.children.push(node.name);\n }\n }),\n removeSceneNode: (name) =>\n set((state) => {\n if (!(name in state.nodeFromName)) {\n console.log(\"Skipping scene node removal for \" + name);\n return;\n }\n\n // Remove this scene node and all children.\n const removeNames: string[] = [];\n function findChildrenRecursive(name: string) {\n removeNames.push(name);\n state.nodeFromName[name]!.children.forEach(\n findChildrenRecursive,\n );\n }\n findChildrenRecursive(name);\n\n removeNames.forEach((removeName) => {\n const node = state.nodeFromName[removeName]!;\n node.cleanup && node.cleanup(); // Free resources.\n delete state.nodeFromName[removeName];\n delete nodeRefFromName.current[removeName];\n });\n\n // Remove node from parent's children list.\n const parent_name = name.split(\"/\").slice(0, -1).join(\"/\");\n state.nodeFromName[parent_name]!.children = state.nodeFromName[\n parent_name\n ]!.children.filter((child_name) => child_name !== name);\n }),\n resetScene: () =>\n set((state) => {\n // For scene resets: we need to retain the object references created for the root and world frame nodes.\n for (const key of Object.keys(state.nodeFromName)) {\n if (key !== \"\" && key !== \"/WorldAxes\")\n delete state.nodeFromName[key];\n }\n Object.values(state.nodeFromName).forEach((node) => {\n // Free resources.\n if (node === undefined || node.cleanup === undefined) return;\n node.cleanup();\n });\n state.nodeFromName[\"\"] = rootNodeTemplate;\n state.nodeFromName[\"/WorldAxes\"] = rootAxesNode;\n }),\n setLabelVisibility: (name, labelVisibility) =>\n set((state) => {\n state.labelVisibleFromName[name] = labelVisibility;\n }),\n })),\n ),\n ),\n )[0];\n}\n" }, { "path": "viser/src/viser/client/src/SearchParamsUtils.tsx", "content": "/** Utilities for interacting with the URL search parameters.\n *\n * This lets us specify the websocket server + port from the URL. */\n\nexport const searchParamKey = \"websocket\";\n\nexport function syncSearchParamServer(server: string) {\n const searchParams = new URLSearchParams(window.location.search);\n // No need to update the URL bar if the websocket port matches the HTTP port.\n // So if we navigate to http://localhost:8081, this should by default connect to ws://localhost:8081.\n const isDefaultServer =\n window.location.host.includes(\n server.replace(\"ws://\", \"\").replace(\"/\", \"\"),\n ) ||\n window.location.host.includes(\n server.replace(\"wss://\", \"\").replace(\"/\", \"\"),\n );\n if (isDefaultServer && searchParams.has(searchParamKey)) {\n searchParams.delete(searchParamKey);\n } else if (!isDefaultServer) {\n searchParams.set(searchParamKey, server);\n }\n window.history.replaceState(\n null,\n \"Viser\",\n // We could use URLSearchParams.toString() to build this string, but that\n // would escape it. We're going to just not escape the string. :)\n searchParams.size === 0\n ? window.location.href.split(\"?\")[0]\n : \"?\" +\n Array.from(searchParams.entries())\n .map(([k, v]) => `${k}=${v}`)\n .join(\"&\"),\n );\n}\n" }, { "path": "viser/src/viser/client/src/Splatting/GaussianSplats.tsx", "content": "/** Gaussian splatting implementation for viser.\n *\n * This borrows heavily from existing open-source implementations. Particularly\n * useful references:\n * - https://github.com/quadjr/aframe-gaussian-splatting\n * - https://github.com/antimatter15/splat\n * - https://github.com/pmndrs/drei\n * - https://github.com/vincent-lecrubier-skydio/react-three-fiber-gaussian-splat\n *\n * Usage should look like:\n *\n * \n * \n * \n * \n * \n *\n * Where `buffer` contains serialized Gaussian attributes. SplatObjects are\n * globally sorted by a worker (with some help from WebAssembly + SIMD\n * intrinsics), and then rendered as a single threejs mesh. Unlike other R3F\n * implementations that we're aware of, this enables correct compositing\n * between multiple splat objects.\n */\n\nimport React from \"react\";\nimport * as THREE from \"three\";\nimport SplatSortWorker from \"./SplatSortWorker?worker\";\nimport { useFrame, useThree } from \"@react-three/fiber\";\nimport { shaderMaterial } from \"@react-three/drei\";\nimport { SorterWorkerIncoming } from \"./SplatSortWorker\";\nimport { create } from \"zustand\";\nimport { Object3D } from \"three\";\n\n/**Global splat state.*/\ninterface SplatState {\n groupBufferFromId: { [id: string]: Uint32Array };\n nodeRefFromId: React.MutableRefObject<{\n [name: string]: undefined | Object3D;\n }>;\n setBuffer: (id: string, buffer: Uint32Array) => void;\n removeBuffer: (id: string) => void;\n}\n\n/**Hook for creating global splat state.*/\nfunction useGaussianSplatStore() {\n const nodeRefFromId = React.useRef({});\n return React.useState(() =>\n create((set) => ({\n groupBufferFromId: {},\n nodeRefFromId: nodeRefFromId,\n setBuffer: (id, buffer) => {\n return set((state) => ({\n groupBufferFromId: { ...state.groupBufferFromId, [id]: buffer },\n }));\n },\n removeBuffer: (id) => {\n return set((state) => {\n // eslint-disable-next-line @typescript-eslint/no-unused-vars\n const { [id]: _, ...buffers } = state.groupBufferFromId;\n return { groupBufferFromId: buffers };\n });\n },\n })),\n )[0];\n}\nconst GaussianSplatsContext = React.createContext | null>(null);\n\n/**Provider for creating splat rendering context.*/\nexport function SplatRenderContext({\n children,\n}: {\n children: React.ReactNode;\n}) {\n const store = useGaussianSplatStore();\n return (\n \n \n {children}\n \n );\n}\n\nconst GaussianSplatMaterial = /* @__PURE__ */ shaderMaterial(\n {\n numGaussians: 0,\n focal: 100.0,\n viewport: [640, 480],\n near: 1.0,\n far: 100.0,\n depthTest: true,\n depthWrite: false,\n transparent: true,\n textureBuffer: null,\n textureT_camera_groups: null,\n transitionInState: 0.0,\n },\n `precision highp usampler2D; // Most important: ints must be 32-bit.\n precision mediump float;\n\n // Index from the splat sorter.\n attribute uint sortedIndex;\n\n // Buffers for splat data; each Gaussian gets 4 floats and 4 int32s. We just\n // copy quadjr for this.\n uniform usampler2D textureBuffer;\n\n // We could also use a uniform to store transforms, but this would be more\n // limiting in terms of the # of groups we can have.\n uniform sampler2D textureT_camera_groups;\n\n // Various other uniforms...\n uniform uint numGaussians;\n uniform vec2 focal;\n uniform vec2 viewport;\n uniform float near;\n uniform float far;\n\n // Fade in state between [0, 1].\n uniform float transitionInState;\n\n out vec4 vRgba;\n out vec2 vPosition;\n\n // Function to fetch and construct the i-th transform matrix using texelFetch\n mat4 getGroupTransform(uint i) {\n // Calculate the base index for the i-th transform.\n uint baseIndex = i * 3u;\n\n // Fetch the texels that represent the first 3 rows of the transform. We\n // choose to use row-major here, since it lets us exclude the fourth row of\n // the matrix.\n vec4 row0 = texelFetch(textureT_camera_groups, ivec2(baseIndex + 0u, 0), 0);\n vec4 row1 = texelFetch(textureT_camera_groups, ivec2(baseIndex + 1u, 0), 0);\n vec4 row2 = texelFetch(textureT_camera_groups, ivec2(baseIndex + 2u, 0), 0);\n\n // Construct the mat4 with the fetched rows.\n mat4 transform = mat4(row0, row1, row2, vec4(0.0, 0.0, 0.0, 1.0));\n return transpose(transform);\n }\n\n void main () {\n // Get position + scale from float buffer.\n ivec2 texSize = textureSize(textureBuffer, 0);\n uint texStart = sortedIndex << 1u;\n ivec2 texPos0 = ivec2(texStart % uint(texSize.x), texStart / uint(texSize.x));\n\n\n // Fetch from textures.\n uvec4 floatBufferData = texelFetch(textureBuffer, texPos0, 0);\n mat4 T_camera_group = getGroupTransform(floatBufferData.w);\n\n // Any early return will discard the fragment.\n gl_Position = vec4(0.0, 0.0, 2.0, 1.0);\n\n // Get center wrt camera. modelViewMatrix is T_cam_world.\n vec3 center = uintBitsToFloat(floatBufferData.xyz);\n vec4 c_cam = T_camera_group * vec4(center, 1);\n if (-c_cam.z < near || -c_cam.z > far)\n return;\n vec4 pos2d = projectionMatrix * c_cam;\n float clip = 1.1 * pos2d.w;\n if (pos2d.x < -clip || pos2d.x > clip || pos2d.y < -clip || pos2d.y > clip)\n return;\n\n // Read covariance terms.\n ivec2 texPos1 = ivec2((texStart + 1u) % uint(texSize.x), (texStart + 1u) / uint(texSize.x));\n uvec4 intBufferData = texelFetch(textureBuffer, texPos1, 0);\n\n // Get covariance terms from int buffer.\n uint rgbaUint32 = intBufferData.w;\n vec2 chol01 = unpackHalf2x16(intBufferData.x);\n vec2 chol23 = unpackHalf2x16(intBufferData.y);\n vec2 chol45 = unpackHalf2x16(intBufferData.z);\n\n // Transition in.\n float startTime = 0.8 * float(sortedIndex) / float(numGaussians);\n float cov_scale = smoothstep(startTime, startTime + 0.2, transitionInState);\n\n // Do the actual splatting.\n mat3 chol = mat3(\n chol01.x, chol01.y, chol23.x,\n 0., chol23.y, chol45.x,\n 0., 0., chol45.y\n );\n mat3 cov3d = chol * transpose(chol) * cov_scale;\n mat3 J = mat3(\n // Matrices are column-major.\n focal.x / c_cam.z, 0., 0.0,\n 0., focal.y / c_cam.z, 0.0,\n -(focal.x * c_cam.x) / (c_cam.z * c_cam.z), -(focal.y * c_cam.y) / (c_cam.z * c_cam.z), 0.\n );\n mat3 A = J * mat3(T_camera_group);\n mat3 cov_proj = A * cov3d * transpose(A);\n float diag1 = cov_proj[0][0] + 0.3;\n float offDiag = cov_proj[0][1];\n float diag2 = cov_proj[1][1] + 0.3;\n\n // Eigendecomposition.\n float mid = 0.5 * (diag1 + diag2);\n float radius = length(vec2((diag1 - diag2) / 2.0, offDiag));\n float lambda1 = mid + radius;\n float lambda2 = mid - radius;\n if (lambda2 < 0.0)\n return;\n vec2 diagonalVector = normalize(vec2(offDiag, lambda1 - diag1));\n vec2 v1 = min(sqrt(2.0 * lambda1), 1024.0) * diagonalVector;\n vec2 v2 = min(sqrt(2.0 * lambda2), 1024.0) * vec2(diagonalVector.y, -diagonalVector.x);\n\n vRgba = vec4(\n float(rgbaUint32 & uint(0xFF)) / 255.0,\n float((rgbaUint32 >> uint(8)) & uint(0xFF)) / 255.0,\n float((rgbaUint32 >> uint(16)) & uint(0xFF)) / 255.0,\n float(rgbaUint32 >> uint(24)) / 255.0\n );\n\n // Throw the Gaussian off the screen if it's too close, too far, or too small.\n float weightedDeterminant = vRgba.a * (diag1 * diag2 - offDiag * offDiag);\n if (weightedDeterminant < 0.5)\n return;\n vPosition = position.xy;\n\n gl_Position = vec4(\n vec2(pos2d) / pos2d.w\n + position.x * v1 / viewport * 2.0\n + position.y * v2 / viewport * 2.0, pos2d.z / pos2d.w, 1.);\n }\n`,\n `precision mediump float;\n\n uniform vec2 viewport;\n uniform vec2 focal;\n\n in vec4 vRgba;\n in vec2 vPosition;\n\n void main () {\n float A = -dot(vPosition, vPosition);\n if (A < -4.0) discard;\n float B = exp(A) * vRgba.a;\n if (B < 0.01) discard; // alphaTest.\n gl_FragColor = vec4(vRgba.rgb, B);\n }`,\n);\n\nexport const SplatObject = React.forwardRef<\n THREE.Group,\n {\n buffer: Uint32Array;\n }\n>(function SplatObject({ buffer }, ref) {\n const splatContext = React.useContext(GaussianSplatsContext)!;\n const setBuffer = splatContext((state) => state.setBuffer);\n const removeBuffer = splatContext((state) => state.removeBuffer);\n const nodeRefFromId = splatContext((state) => state.nodeRefFromId);\n const name = React.useMemo(() => crypto.randomUUID(), [buffer]);\n\n const [obj, setRef] = React.useState(null);\n\n React.useEffect(() => {\n if (obj === null) return;\n setBuffer(name, buffer);\n if (ref !== null) {\n if (\"current\" in ref) {\n ref.current = obj;\n } else {\n ref(obj);\n }\n }\n nodeRefFromId.current[name] = obj;\n return () => {\n removeBuffer(name);\n delete nodeRefFromId.current[name];\n };\n }, [obj]);\n\n return ;\n});\n\n/** External interface. Component should be added to the root of canvas. */\nfunction SplatRenderer() {\n const splatContext = React.useContext(GaussianSplatsContext)!;\n const groupBufferFromId = splatContext((state) => state.groupBufferFromId);\n const nodeRefFromId = splatContext((state) => state.nodeRefFromId);\n\n // Consolidate Gaussian groups into a single buffer.\n const merged = mergeGaussianGroups(groupBufferFromId);\n const meshProps = useGaussianMeshProps(\n merged.gaussianBuffer,\n merged.numGroups,\n );\n\n // Create sorting worker.\n const sortWorker = new SplatSortWorker();\n let initializedBufferTexture = false;\n sortWorker.onmessage = (e) => {\n // Update rendering order.\n const sortedIndices = e.data.sortedIndices as Uint32Array;\n meshProps.sortedIndexAttribute.set(sortedIndices);\n meshProps.sortedIndexAttribute.needsUpdate = true;\n\n // Trigger initial render.\n if (!initializedBufferTexture) {\n meshProps.material.uniforms.numGaussians.value = merged.numGaussians;\n meshProps.textureBuffer.needsUpdate = true;\n initializedBufferTexture = true;\n }\n };\n function postToWorker(message: SorterWorkerIncoming) {\n sortWorker.postMessage(message);\n }\n\n postToWorker({\n setBuffer: merged.gaussianBuffer,\n setGroupIndices: merged.groupIndices,\n });\n\n // Cleanup.\n React.useEffect(() => {\n return () => {\n meshProps.textureBuffer.dispose();\n meshProps.geometry.dispose();\n meshProps.material.dispose();\n postToWorker({ close: true });\n };\n });\n\n // Per-frame updates. This is in charge of synchronizing transforms and\n // triggering sorting.\n //\n // We pre-allocate matrices to make life easier for the garbage collector.\n const meshRef = React.useRef(null);\n const tmpT_camera_group = new THREE.Matrix4();\n const Tz_camera_groups = new Float32Array(merged.numGroups * 4);\n const prevRowMajorT_camera_groups = meshProps.rowMajorT_camera_groups\n .slice()\n .fill(0);\n const prevVisibles: boolean[] = [];\n useFrame((state, delta) => {\n const mesh = meshRef.current;\n if (mesh === null || sortWorker === null) return;\n\n // Update camera parameter uniforms.\n const dpr = state.viewport.dpr;\n const fovY =\n ((state.camera as THREE.PerspectiveCamera).fov * Math.PI) / 180.0;\n const fovX = 2 * Math.atan(Math.tan(fovY / 2) * state.viewport.aspect);\n const fy = (dpr * state.size.height) / (2 * Math.tan(fovY / 2));\n const fx = (dpr * state.size.width) / (2 * Math.tan(fovX / 2));\n\n if (meshProps.material === undefined) return;\n\n const uniforms = meshProps.material.uniforms;\n uniforms.transitionInState.value = Math.min(\n uniforms.transitionInState.value + delta * 2.0,\n 1.0,\n );\n uniforms.focal.value = [fx, fy];\n uniforms.near.value = state.camera.near;\n uniforms.far.value = state.camera.far;\n uniforms.viewport.value = [state.size.width * dpr, state.size.height * dpr];\n\n // Update group transforms.\n const T_camera_world = state.camera.matrixWorldInverse;\n const groupVisibles: boolean[] = [];\n let visibilitiesChanged = false;\n for (const [groupIndex, name] of Object.keys(groupBufferFromId).entries()) {\n const node = nodeRefFromId.current[name];\n if (node === undefined) continue;\n tmpT_camera_group.copy(T_camera_world).multiply(node.matrixWorld);\n const colMajorElements = tmpT_camera_group.elements;\n Tz_camera_groups.set(\n [\n colMajorElements[2],\n colMajorElements[6],\n colMajorElements[10],\n colMajorElements[14],\n ],\n groupIndex * 4,\n );\n const rowMajorElements = tmpT_camera_group.transpose().elements;\n meshProps.rowMajorT_camera_groups.set(\n rowMajorElements.slice(0, 12),\n groupIndex * 12,\n );\n\n // Determine visibility. If the parent has unmountWhenInvisible=true, the\n // first frame after showing a hidden parent can have visible=true with\n // an incorrect matrixWorld transform. There might be a better fix, but\n // `prevVisible` is an easy workaround for this.\n let visibleNow = node.visible && node.parent !== null;\n if (visibleNow) {\n node.traverseAncestors((ancestor) => {\n visibleNow = visibleNow && ancestor.visible;\n });\n }\n groupVisibles.push(visibleNow && prevVisibles[groupIndex] === true);\n if (prevVisibles[groupIndex] !== visibleNow) {\n prevVisibles[groupIndex] = visibleNow;\n visibilitiesChanged = true;\n }\n }\n\n const groupsMovedWrtCam = !meshProps.rowMajorT_camera_groups.every(\n (v, i) => v === prevRowMajorT_camera_groups[i],\n );\n\n if (groupsMovedWrtCam) {\n // Gaussians need to be re-sorted.\n postToWorker({\n setTz_camera_groups: Tz_camera_groups,\n });\n }\n if (groupsMovedWrtCam || visibilitiesChanged) {\n // If a group is not visible, we'll throw it off the screen with some Big\n // Numbers. It's important that this only impacts the coordinates used\n // for the shader and not for the sorter; that way when we \"show\" a group\n // of Gaussians the correct rendering order is immediately available.\n for (const [i, visible] of groupVisibles.entries()) {\n if (!visible) {\n meshProps.rowMajorT_camera_groups[i * 12 + 3] = 1e10;\n meshProps.rowMajorT_camera_groups[i * 12 + 7] = 1e10;\n meshProps.rowMajorT_camera_groups[i * 12 + 11] = 1e10;\n }\n }\n prevRowMajorT_camera_groups.set(meshProps.rowMajorT_camera_groups);\n meshProps.textureT_camera_groups.needsUpdate = true;\n }\n }, -100 /* This should be called early to reduce group transform artifacts. */);\n\n return (\n \n );\n}\n\n/**Consolidate groups of Gaussians into a single buffer, to make it possible\n * for them to be sorted globally.*/\nfunction mergeGaussianGroups(groupBufferFromName: {\n [name: string]: Uint32Array;\n}) {\n // Create geometry. Each Gaussian will be rendered as a quad.\n let totalBufferLength = 0;\n for (const buffer of Object.values(groupBufferFromName)) {\n totalBufferLength += buffer.length;\n }\n const numGaussians = totalBufferLength / 8;\n const gaussianBuffer = new Uint32Array(totalBufferLength);\n const groupIndices = new Uint32Array(numGaussians);\n\n let offset = 0;\n for (const [groupIndex, groupBuffer] of Object.values(\n groupBufferFromName,\n ).entries()) {\n groupIndices.fill(\n groupIndex,\n offset / 8,\n (offset + groupBuffer.length) / 8,\n );\n gaussianBuffer.set(groupBuffer, offset);\n\n // Each Gaussian is allocated\n // - 12 bytes for center x, y, z (float32)\n // - 4 bytes for group index (uint32); we're filling this in now\n //\n // - 12 bytes for covariance (6 terms, float16)\n // - 4 bytes for RGBA (uint8)\n for (let i = 0; i < groupBuffer.length; i += 8) {\n gaussianBuffer[offset + i + 3] = groupIndex;\n }\n offset += groupBuffer.length;\n }\n\n const numGroups = Object.keys(groupBufferFromName).length;\n return { numGaussians, gaussianBuffer, numGroups, groupIndices };\n}\n\n/**Hook to generate properties for rendering Gaussians via a three.js mesh.*/\nfunction useGaussianMeshProps(gaussianBuffer: Uint32Array, numGroups: number) {\n const numGaussians = gaussianBuffer.length / 8;\n const maxTextureSize = useThree((state) => state.gl).capabilities\n .maxTextureSize;\n\n // Create instanced geometry.\n const geometry = new THREE.InstancedBufferGeometry();\n geometry.instanceCount = numGaussians;\n geometry.setIndex(\n new THREE.BufferAttribute(new Uint32Array([0, 2, 1, 0, 3, 2]), 1),\n );\n geometry.setAttribute(\n \"position\",\n new THREE.BufferAttribute(\n new Float32Array([-2, -2, 2, -2, 2, 2, -2, 2]),\n 2,\n ),\n );\n\n // Rendering order for Gaussians.\n const sortedIndexAttribute = new THREE.InstancedBufferAttribute(\n new Uint32Array(numGaussians),\n 1,\n );\n sortedIndexAttribute.setUsage(THREE.DynamicDrawUsage);\n geometry.setAttribute(\"sortedIndex\", sortedIndexAttribute);\n\n // Create texture buffers.\n const textureWidth = Math.min(numGaussians * 2, maxTextureSize);\n const textureHeight = Math.ceil((numGaussians * 2) / textureWidth);\n const bufferPadded = new Uint32Array(textureWidth * textureHeight * 4);\n bufferPadded.set(gaussianBuffer);\n const textureBuffer = new THREE.DataTexture(\n bufferPadded,\n textureWidth,\n textureHeight,\n THREE.RGBAIntegerFormat,\n THREE.UnsignedIntType,\n );\n textureBuffer.internalFormat = \"RGBA32UI\";\n textureBuffer.needsUpdate = true;\n\n const rowMajorT_camera_groups = new Float32Array(numGroups * 12);\n const textureT_camera_groups = new THREE.DataTexture(\n rowMajorT_camera_groups,\n (numGroups * 12) / 4,\n 1,\n THREE.RGBAFormat,\n THREE.FloatType,\n );\n textureT_camera_groups.internalFormat = \"RGBA32F\";\n textureT_camera_groups.needsUpdate = true;\n\n const material = new GaussianSplatMaterial({\n // @ts-ignore\n textureBuffer: textureBuffer,\n textureT_camera_groups: textureT_camera_groups,\n numGaussians: 0,\n transitionInState: 0.0,\n });\n\n return {\n geometry,\n material,\n textureBuffer,\n sortedIndexAttribute,\n textureT_camera_groups,\n rowMajorT_camera_groups,\n };\n}\n" }, { "path": "viser/src/viser/client/src/Splatting/SplatSortWorker.ts", "content": "/** Worker for sorting splats.\n */\n\nimport MakeSorterModulePromise from \"./WasmSorter/Sorter.mjs\";\n\nexport type SorterWorkerIncoming =\n | {\n setBuffer: Uint32Array;\n setGroupIndices: Uint32Array;\n }\n | {\n setTz_camera_groups: Float32Array;\n }\n | { close: true };\n\n{\n let sorter: any = null;\n let Tz_camera_groups: Float32Array | null = null;\n let sortRunning = false;\n const throttledSort = () => {\n if (sorter === null || Tz_camera_groups === null) {\n setTimeout(throttledSort, 1);\n return;\n }\n if (sortRunning) return;\n\n sortRunning = true;\n const lastView = Tz_camera_groups;\n\n // Important: we clone the output so we can transfer the buffer to the main\n // thread. Compared to relying on postMessage for copying, this reduces\n // backlog artifacts.\n const sortedIndices = (\n sorter.sort(Tz_camera_groups) as Uint32Array\n ).slice();\n\n // @ts-ignore\n self.postMessage({ sortedIndices: sortedIndices }, [sortedIndices.buffer]);\n\n setTimeout(() => {\n sortRunning = false;\n if (Tz_camera_groups === null) return;\n if (\n !lastView.every(\n // Cast is needed because of closure...\n (val, i) => val === (Tz_camera_groups as Float32Array)[i],\n )\n ) {\n throttledSort();\n }\n }, 0);\n };\n\n const SorterModulePromise = MakeSorterModulePromise();\n\n self.onmessage = async (e) => {\n const data = e.data as SorterWorkerIncoming;\n if (\"setBuffer\" in data) {\n // Instantiate sorter with buffers populated.\n sorter = new (await SorterModulePromise).Sorter(\n data.setBuffer,\n data.setGroupIndices,\n );\n } else if (\"setTz_camera_groups\" in data) {\n // Update object transforms.\n Tz_camera_groups = data.setTz_camera_groups;\n throttledSort();\n } else if (\"close\" in data) {\n // Done!\n self.close();\n }\n };\n}\n" }, { "path": "viser/src/viser/client/src/Splatting/WasmSorter/Sorter.mjs", "content": "\nvar Module = (() => {\n var _scriptName = import.meta.url;\n \n return (\nasync function(moduleArg = {}) {\n var moduleRtn;\n\nvar Module=moduleArg;var readyPromiseResolve,readyPromiseReject;var readyPromise=new Promise((resolve,reject)=>{readyPromiseResolve=resolve;readyPromiseReject=reject});var ENVIRONMENT_IS_WEB=typeof window==\"object\";var ENVIRONMENT_IS_WORKER=typeof importScripts==\"function\";var ENVIRONMENT_IS_NODE=typeof process==\"object\"&&typeof process.versions==\"object\"&&typeof process.versions.node==\"string\";if(ENVIRONMENT_IS_NODE){const{createRequire:createRequire}=await import(\"module\");var require=createRequire(import.meta.url)}var moduleOverrides=Object.assign({},Module);var arguments_=[];var thisProgram=\"./this.program\";var quit_=(status,toThrow)=>{throw toThrow};var scriptDirectory=\"\";function locateFile(path){if(Module[\"locateFile\"]){return Module[\"locateFile\"](path,scriptDirectory)}return scriptDirectory+path}var read_,readAsync,readBinary;if(ENVIRONMENT_IS_NODE){var fs=require(\"fs\");var nodePath=require(\"path\");scriptDirectory=require(\"url\").fileURLToPath(new URL(\"./\",import.meta.url));read_=(filename,binary)=>{filename=isFileURI(filename)?new URL(filename):nodePath.normalize(filename);return fs.readFileSync(filename,binary?undefined:\"utf8\")};readBinary=filename=>{var ret=read_(filename,true);if(!ret.buffer){ret=new Uint8Array(ret)}return ret};readAsync=(filename,onload,onerror,binary=true)=>{filename=isFileURI(filename)?new URL(filename):nodePath.normalize(filename);fs.readFile(filename,binary?undefined:\"utf8\",(err,data)=>{if(err)onerror(err);else onload(binary?data.buffer:data)})};if(!Module[\"thisProgram\"]&&process.argv.length>1){thisProgram=process.argv[1].replace(/\\\\/g,\"/\")}arguments_=process.argv.slice(2);quit_=(status,toThrow)=>{process.exitCode=status;throw toThrow}}else if(ENVIRONMENT_IS_WEB||ENVIRONMENT_IS_WORKER){if(ENVIRONMENT_IS_WORKER){scriptDirectory=self.location.href}else if(typeof document!=\"undefined\"&&document.currentScript){scriptDirectory=document.currentScript.src}if(_scriptName){scriptDirectory=_scriptName}if(scriptDirectory.startsWith(\"blob:\")){scriptDirectory=\"\"}else{scriptDirectory=scriptDirectory.substr(0,scriptDirectory.replace(/[?#].*/,\"\").lastIndexOf(\"/\")+1)}{read_=url=>{var xhr=new XMLHttpRequest;xhr.open(\"GET\",url,false);xhr.send(null);return xhr.responseText};if(ENVIRONMENT_IS_WORKER){readBinary=url=>{var xhr=new XMLHttpRequest;xhr.open(\"GET\",url,false);xhr.responseType=\"arraybuffer\";xhr.send(null);return new Uint8Array(xhr.response)}}readAsync=(url,onload,onerror)=>{if(isFileURI(url)){var xhr=new XMLHttpRequest;xhr.open(\"GET\",url,true);xhr.responseType=\"arraybuffer\";xhr.onload=()=>{if(xhr.status==200||xhr.status==0&&xhr.response){onload(xhr.response);return}onerror()};xhr.onerror=onerror;xhr.send(null);return}fetch(url,{credentials:\"same-origin\"}).then(response=>{if(response.ok){return response.arrayBuffer()}return Promise.reject(new Error(response.status+\" : \"+response.url))}).then(onload,onerror)}}}else{}var out=Module[\"print\"]||console.log.bind(console);var err=Module[\"printErr\"]||console.error.bind(console);Object.assign(Module,moduleOverrides);moduleOverrides=null;if(Module[\"arguments\"])arguments_=Module[\"arguments\"];if(Module[\"thisProgram\"])thisProgram=Module[\"thisProgram\"];if(Module[\"quit\"])quit_=Module[\"quit\"];var wasmBinary;if(Module[\"wasmBinary\"])wasmBinary=Module[\"wasmBinary\"];var wasmMemory;var ABORT=false;var EXITSTATUS;var HEAP8,HEAPU8,HEAP16,HEAPU16,HEAP32,HEAPU32,HEAPF32,HEAPF64;function updateMemoryViews(){var b=wasmMemory.buffer;Module[\"HEAP8\"]=HEAP8=new Int8Array(b);Module[\"HEAP16\"]=HEAP16=new Int16Array(b);Module[\"HEAPU8\"]=HEAPU8=new Uint8Array(b);Module[\"HEAPU16\"]=HEAPU16=new Uint16Array(b);Module[\"HEAP32\"]=HEAP32=new Int32Array(b);Module[\"HEAPU32\"]=HEAPU32=new Uint32Array(b);Module[\"HEAPF32\"]=HEAPF32=new Float32Array(b);Module[\"HEAPF64\"]=HEAPF64=new Float64Array(b)}var __ATPRERUN__=[];var __ATINIT__=[];var __ATPOSTRUN__=[];var runtimeInitialized=false;function preRun(){if(Module[\"preRun\"]){if(typeof Module[\"preRun\"]==\"function\")Module[\"preRun\"]=[Module[\"preRun\"]];while(Module[\"preRun\"].length){addOnPreRun(Module[\"preRun\"].shift())}}callRuntimeCallbacks(__ATPRERUN__)}function initRuntime(){runtimeInitialized=true;callRuntimeCallbacks(__ATINIT__)}function postRun(){if(Module[\"postRun\"]){if(typeof Module[\"postRun\"]==\"function\")Module[\"postRun\"]=[Module[\"postRun\"]];while(Module[\"postRun\"].length){addOnPostRun(Module[\"postRun\"].shift())}}callRuntimeCallbacks(__ATPOSTRUN__)}function addOnPreRun(cb){__ATPRERUN__.unshift(cb)}function addOnInit(cb){__ATINIT__.unshift(cb)}function addOnPostRun(cb){__ATPOSTRUN__.unshift(cb)}var runDependencies=0;var runDependencyWatcher=null;var dependenciesFulfilled=null;function addRunDependency(id){runDependencies++;Module[\"monitorRunDependencies\"]?.(runDependencies)}function removeRunDependency(id){runDependencies--;Module[\"monitorRunDependencies\"]?.(runDependencies);if(runDependencies==0){if(runDependencyWatcher!==null){clearInterval(runDependencyWatcher);runDependencyWatcher=null}if(dependenciesFulfilled){var callback=dependenciesFulfilled;dependenciesFulfilled=null;callback()}}}function abort(what){Module[\"onAbort\"]?.(what);what=\"Aborted(\"+what+\")\";err(what);ABORT=true;EXITSTATUS=1;what+=\". Build with -sASSERTIONS for more info.\";var e=new WebAssembly.RuntimeError(what);readyPromiseReject(e);throw e}var dataURIPrefix=\"data:application/octet-stream;base64,\";var isDataURI=filename=>filename.startsWith(dataURIPrefix);var isFileURI=filename=>filename.startsWith(\"file://\");function findWasmBinary(){if(Module[\"locateFile\"]){var f=\"Sorter.wasm\";if(!isDataURI(f)){return locateFile(f)}return f}return new URL(\"Sorter.wasm\",import.meta.url).href}var wasmBinaryFile;function getBinarySync(file){if(file==wasmBinaryFile&&wasmBinary){return new Uint8Array(wasmBinary)}if(readBinary){return readBinary(file)}throw\"both async and sync fetching of the wasm failed\"}function getBinaryPromise(binaryFile){if(!wasmBinary){return new Promise((resolve,reject)=>{readAsync(binaryFile,response=>resolve(new Uint8Array(response)),error=>{try{resolve(getBinarySync(binaryFile))}catch(e){reject(e)}})})}return Promise.resolve().then(()=>getBinarySync(binaryFile))}function instantiateArrayBuffer(binaryFile,imports,receiver){return getBinaryPromise(binaryFile).then(binary=>WebAssembly.instantiate(binary,imports)).then(receiver,reason=>{err(`failed to asynchronously prepare wasm: ${reason}`);abort(reason)})}function instantiateAsync(binary,binaryFile,imports,callback){if(!binary&&typeof WebAssembly.instantiateStreaming==\"function\"&&!isDataURI(binaryFile)&&!isFileURI(binaryFile)&&!ENVIRONMENT_IS_NODE&&typeof fetch==\"function\"){return fetch(binaryFile,{credentials:\"same-origin\"}).then(response=>{var result=WebAssembly.instantiateStreaming(response,imports);return result.then(callback,function(reason){err(`wasm streaming compile failed: ${reason}`);err(\"falling back to ArrayBuffer instantiation\");return instantiateArrayBuffer(binaryFile,imports,callback)})})}return instantiateArrayBuffer(binaryFile,imports,callback)}function getWasmImports(){return{a:wasmImports}}function createWasm(){var info=getWasmImports();function receiveInstance(instance,module){wasmExports=instance.exports;wasmMemory=wasmExports[\"z\"];updateMemoryViews();wasmTable=wasmExports[\"C\"];addOnInit(wasmExports[\"A\"]);removeRunDependency(\"wasm-instantiate\");return wasmExports}addRunDependency(\"wasm-instantiate\");function receiveInstantiationResult(result){receiveInstance(result[\"instance\"])}if(Module[\"instantiateWasm\"]){try{return Module[\"instantiateWasm\"](info,receiveInstance)}catch(e){err(`Module.instantiateWasm callback failed with error: ${e}`);readyPromiseReject(e)}}if(!wasmBinaryFile)wasmBinaryFile=findWasmBinary();instantiateAsync(wasmBinary,wasmBinaryFile,info,receiveInstantiationResult).catch(readyPromiseReject);return{}}var callRuntimeCallbacks=callbacks=>{while(callbacks.length>0){callbacks.shift()(Module)}};var noExitRuntime=Module[\"noExitRuntime\"]||true;class ExceptionInfo{constructor(excPtr){this.excPtr=excPtr;this.ptr=excPtr-24}set_type(type){HEAPU32[this.ptr+4>>2]=type}get_type(){return HEAPU32[this.ptr+4>>2]}set_destructor(destructor){HEAPU32[this.ptr+8>>2]=destructor}get_destructor(){return HEAPU32[this.ptr+8>>2]}set_caught(caught){caught=caught?1:0;HEAP8[this.ptr+12]=caught}get_caught(){return HEAP8[this.ptr+12]!=0}set_rethrown(rethrown){rethrown=rethrown?1:0;HEAP8[this.ptr+13]=rethrown}get_rethrown(){return HEAP8[this.ptr+13]!=0}init(type,destructor){this.set_adjusted_ptr(0);this.set_type(type);this.set_destructor(destructor)}set_adjusted_ptr(adjustedPtr){HEAPU32[this.ptr+16>>2]=adjustedPtr}get_adjusted_ptr(){return HEAPU32[this.ptr+16>>2]}get_exception_ptr(){var isPointer=___cxa_is_pointer_type(this.get_type());if(isPointer){return HEAPU32[this.excPtr>>2]}var adjusted=this.get_adjusted_ptr();if(adjusted!==0)return adjusted;return this.excPtr}}var exceptionLast=0;var uncaughtExceptionCount=0;var ___cxa_throw=(ptr,type,destructor)=>{var info=new ExceptionInfo(ptr);info.init(type,destructor);exceptionLast=ptr;uncaughtExceptionCount++;throw exceptionLast};var __abort_js=()=>{abort(\"\")};var __embind_register_bigint=(primitiveType,name,size,minRange,maxRange)=>{};var embind_init_charCodes=()=>{var codes=new Array(256);for(var i=0;i<256;++i){codes[i]=String.fromCharCode(i)}embind_charCodes=codes};var embind_charCodes;var readLatin1String=ptr=>{var ret=\"\";var c=ptr;while(HEAPU8[c]){ret+=embind_charCodes[HEAPU8[c++]]}return ret};var awaitingDependencies={};var registeredTypes={};var typeDependencies={};var BindingError;var throwBindingError=message=>{throw new BindingError(message)};var InternalError;var throwInternalError=message=>{throw new InternalError(message)};var whenDependentTypesAreResolved=(myTypes,dependentTypes,getTypeConverters)=>{myTypes.forEach(function(type){typeDependencies[type]=dependentTypes});function onComplete(typeConverters){var myTypeConverters=getTypeConverters(typeConverters);if(myTypeConverters.length!==myTypes.length){throwInternalError(\"Mismatched type converter count\")}for(var i=0;i{if(registeredTypes.hasOwnProperty(dt)){typeConverters[i]=registeredTypes[dt]}else{unregisteredTypes.push(dt);if(!awaitingDependencies.hasOwnProperty(dt)){awaitingDependencies[dt]=[]}awaitingDependencies[dt].push(()=>{typeConverters[i]=registeredTypes[dt];++registered;if(registered===unregisteredTypes.length){onComplete(typeConverters)}})}});if(0===unregisteredTypes.length){onComplete(typeConverters)}};function sharedRegisterType(rawType,registeredInstance,options={}){var name=registeredInstance.name;if(!rawType){throwBindingError(`type \"${name}\" must have a positive integer typeid pointer`)}if(registeredTypes.hasOwnProperty(rawType)){if(options.ignoreDuplicateRegistrations){return}else{throwBindingError(`Cannot register type '${name}' twice`)}}registeredTypes[rawType]=registeredInstance;delete typeDependencies[rawType];if(awaitingDependencies.hasOwnProperty(rawType)){var callbacks=awaitingDependencies[rawType];delete awaitingDependencies[rawType];callbacks.forEach(cb=>cb())}}function registerType(rawType,registeredInstance,options={}){if(!(\"argPackAdvance\"in registeredInstance)){throw new TypeError(\"registerType registeredInstance requires argPackAdvance\")}return sharedRegisterType(rawType,registeredInstance,options)}var GenericWireTypeSize=8;var __embind_register_bool=(rawType,name,trueValue,falseValue)=>{name=readLatin1String(name);registerType(rawType,{name:name,fromWireType:function(wt){return!!wt},toWireType:function(destructors,o){return o?trueValue:falseValue},argPackAdvance:GenericWireTypeSize,readValueFromPointer:function(pointer){return this[\"fromWireType\"](HEAPU8[pointer])},destructorFunction:null})};var shallowCopyInternalPointer=o=>({count:o.count,deleteScheduled:o.deleteScheduled,preservePointerOnDelete:o.preservePointerOnDelete,ptr:o.ptr,ptrType:o.ptrType,smartPtr:o.smartPtr,smartPtrType:o.smartPtrType});var throwInstanceAlreadyDeleted=obj=>{function getInstanceTypeName(handle){return handle.$$.ptrType.registeredClass.name}throwBindingError(getInstanceTypeName(obj)+\" instance already deleted\")};var finalizationRegistry=false;var detachFinalizer=handle=>{};var runDestructor=$$=>{if($$.smartPtr){$$.smartPtrType.rawDestructor($$.smartPtr)}else{$$.ptrType.registeredClass.rawDestructor($$.ptr)}};var releaseClassHandle=$$=>{$$.count.value-=1;var toDelete=0===$$.count.value;if(toDelete){runDestructor($$)}};var downcastPointer=(ptr,ptrClass,desiredClass)=>{if(ptrClass===desiredClass){return ptr}if(undefined===desiredClass.baseClass){return null}var rv=downcastPointer(ptr,ptrClass,desiredClass.baseClass);if(rv===null){return null}return desiredClass.downcast(rv)};var registeredPointers={};var getInheritedInstanceCount=()=>Object.keys(registeredInstances).length;var getLiveInheritedInstances=()=>{var rv=[];for(var k in registeredInstances){if(registeredInstances.hasOwnProperty(k)){rv.push(registeredInstances[k])}}return rv};var deletionQueue=[];var flushPendingDeletes=()=>{while(deletionQueue.length){var obj=deletionQueue.pop();obj.$$.deleteScheduled=false;obj[\"delete\"]()}};var delayFunction;var setDelayFunction=fn=>{delayFunction=fn;if(deletionQueue.length&&delayFunction){delayFunction(flushPendingDeletes)}};var init_embind=()=>{Module[\"getInheritedInstanceCount\"]=getInheritedInstanceCount;Module[\"getLiveInheritedInstances\"]=getLiveInheritedInstances;Module[\"flushPendingDeletes\"]=flushPendingDeletes;Module[\"setDelayFunction\"]=setDelayFunction};var registeredInstances={};var getBasestPointer=(class_,ptr)=>{if(ptr===undefined){throwBindingError(\"ptr should not be undefined\")}while(class_.baseClass){ptr=class_.upcast(ptr);class_=class_.baseClass}return ptr};var getInheritedInstance=(class_,ptr)=>{ptr=getBasestPointer(class_,ptr);return registeredInstances[ptr]};var makeClassHandle=(prototype,record)=>{if(!record.ptrType||!record.ptr){throwInternalError(\"makeClassHandle requires ptr and ptrType\")}var hasSmartPtrType=!!record.smartPtrType;var hasSmartPtr=!!record.smartPtr;if(hasSmartPtrType!==hasSmartPtr){throwInternalError(\"Both smartPtrType and smartPtr must be specified\")}record.count={value:1};return attachFinalizer(Object.create(prototype,{$$:{value:record,writable:true}}))};function RegisteredPointer_fromWireType(ptr){var rawPointer=this.getPointee(ptr);if(!rawPointer){this.destructor(ptr);return null}var registeredInstance=getInheritedInstance(this.registeredClass,rawPointer);if(undefined!==registeredInstance){if(0===registeredInstance.$$.count.value){registeredInstance.$$.ptr=rawPointer;registeredInstance.$$.smartPtr=ptr;return registeredInstance[\"clone\"]()}else{var rv=registeredInstance[\"clone\"]();this.destructor(ptr);return rv}}function makeDefaultHandle(){if(this.isSmartPointer){return makeClassHandle(this.registeredClass.instancePrototype,{ptrType:this.pointeeType,ptr:rawPointer,smartPtrType:this,smartPtr:ptr})}else{return makeClassHandle(this.registeredClass.instancePrototype,{ptrType:this,ptr:ptr})}}var actualType=this.registeredClass.getActualType(rawPointer);var registeredPointerRecord=registeredPointers[actualType];if(!registeredPointerRecord){return makeDefaultHandle.call(this)}var toType;if(this.isConst){toType=registeredPointerRecord.constPointerType}else{toType=registeredPointerRecord.pointerType}var dp=downcastPointer(rawPointer,this.registeredClass,toType.registeredClass);if(dp===null){return makeDefaultHandle.call(this)}if(this.isSmartPointer){return makeClassHandle(toType.registeredClass.instancePrototype,{ptrType:toType,ptr:dp,smartPtrType:this,smartPtr:ptr})}else{return makeClassHandle(toType.registeredClass.instancePrototype,{ptrType:toType,ptr:dp})}}var attachFinalizer=handle=>{if(\"undefined\"===typeof FinalizationRegistry){attachFinalizer=handle=>handle;return handle}finalizationRegistry=new FinalizationRegistry(info=>{releaseClassHandle(info.$$)});attachFinalizer=handle=>{var $$=handle.$$;var hasSmartPtr=!!$$.smartPtr;if(hasSmartPtr){var info={$$:$$};finalizationRegistry.register(handle,info,handle)}return handle};detachFinalizer=handle=>finalizationRegistry.unregister(handle);return attachFinalizer(handle)};var init_ClassHandle=()=>{Object.assign(ClassHandle.prototype,{isAliasOf(other){if(!(this instanceof ClassHandle)){return false}if(!(other instanceof ClassHandle)){return false}var leftClass=this.$$.ptrType.registeredClass;var left=this.$$.ptr;other.$$=other.$$;var rightClass=other.$$.ptrType.registeredClass;var right=other.$$.ptr;while(leftClass.baseClass){left=leftClass.upcast(left);leftClass=leftClass.baseClass}while(rightClass.baseClass){right=rightClass.upcast(right);rightClass=rightClass.baseClass}return leftClass===rightClass&&left===right},clone(){if(!this.$$.ptr){throwInstanceAlreadyDeleted(this)}if(this.$$.preservePointerOnDelete){this.$$.count.value+=1;return this}else{var clone=attachFinalizer(Object.create(Object.getPrototypeOf(this),{$$:{value:shallowCopyInternalPointer(this.$$)}}));clone.$$.count.value+=1;clone.$$.deleteScheduled=false;return clone}},delete(){if(!this.$$.ptr){throwInstanceAlreadyDeleted(this)}if(this.$$.deleteScheduled&&!this.$$.preservePointerOnDelete){throwBindingError(\"Object already scheduled for deletion\")}detachFinalizer(this);releaseClassHandle(this.$$);if(!this.$$.preservePointerOnDelete){this.$$.smartPtr=undefined;this.$$.ptr=undefined}},isDeleted(){return!this.$$.ptr},deleteLater(){if(!this.$$.ptr){throwInstanceAlreadyDeleted(this)}if(this.$$.deleteScheduled&&!this.$$.preservePointerOnDelete){throwBindingError(\"Object already scheduled for deletion\")}deletionQueue.push(this);if(deletionQueue.length===1&&delayFunction){delayFunction(flushPendingDeletes)}this.$$.deleteScheduled=true;return this}})};function ClassHandle(){}var createNamedFunction=(name,body)=>Object.defineProperty(body,\"name\",{value:name});var ensureOverloadTable=(proto,methodName,humanName)=>{if(undefined===proto[methodName].overloadTable){var prevFunc=proto[methodName];proto[methodName]=function(...args){if(!proto[methodName].overloadTable.hasOwnProperty(args.length)){throwBindingError(`Function '${humanName}' called with an invalid number of arguments (${args.length}) - expects one of (${proto[methodName].overloadTable})!`)}return proto[methodName].overloadTable[args.length].apply(this,args)};proto[methodName].overloadTable=[];proto[methodName].overloadTable[prevFunc.argCount]=prevFunc}};var exposePublicSymbol=(name,value,numArguments)=>{if(Module.hasOwnProperty(name)){if(undefined===numArguments||undefined!==Module[name].overloadTable&&undefined!==Module[name].overloadTable[numArguments]){throwBindingError(`Cannot register public name '${name}' twice`)}ensureOverloadTable(Module,name,name);if(Module.hasOwnProperty(numArguments)){throwBindingError(`Cannot register multiple overloads of a function with the same number of arguments (${numArguments})!`)}Module[name].overloadTable[numArguments]=value}else{Module[name]=value;if(undefined!==numArguments){Module[name].numArguments=numArguments}}};var char_0=48;var char_9=57;var makeLegalFunctionName=name=>{if(undefined===name){return\"_unknown\"}name=name.replace(/[^a-zA-Z0-9_]/g,\"$\");var f=name.charCodeAt(0);if(f>=char_0&&f<=char_9){return`_${name}`}return name};function RegisteredClass(name,constructor,instancePrototype,rawDestructor,baseClass,getActualType,upcast,downcast){this.name=name;this.constructor=constructor;this.instancePrototype=instancePrototype;this.rawDestructor=rawDestructor;this.baseClass=baseClass;this.getActualType=getActualType;this.upcast=upcast;this.downcast=downcast;this.pureVirtualFunctions=[]}var upcastPointer=(ptr,ptrClass,desiredClass)=>{while(ptrClass!==desiredClass){if(!ptrClass.upcast){throwBindingError(`Expected null or instance of ${desiredClass.name}, got an instance of ${ptrClass.name}`)}ptr=ptrClass.upcast(ptr);ptrClass=ptrClass.baseClass}return ptr};function constNoSmartPtrRawPointerToWireType(destructors,handle){if(handle===null){if(this.isReference){throwBindingError(`null is not a valid ${this.name}`)}return 0}if(!handle.$$){throwBindingError(`Cannot pass \"${embindRepr(handle)}\" as a ${this.name}`)}if(!handle.$$.ptr){throwBindingError(`Cannot pass deleted object as a pointer of type ${this.name}`)}var handleClass=handle.$$.ptrType.registeredClass;var ptr=upcastPointer(handle.$$.ptr,handleClass,this.registeredClass);return ptr}function genericPointerToWireType(destructors,handle){var ptr;if(handle===null){if(this.isReference){throwBindingError(`null is not a valid ${this.name}`)}if(this.isSmartPointer){ptr=this.rawConstructor();if(destructors!==null){destructors.push(this.rawDestructor,ptr)}return ptr}else{return 0}}if(!handle||!handle.$$){throwBindingError(`Cannot pass \"${embindRepr(handle)}\" as a ${this.name}`)}if(!handle.$$.ptr){throwBindingError(`Cannot pass deleted object as a pointer of type ${this.name}`)}if(!this.isConst&&handle.$$.ptrType.isConst){throwBindingError(`Cannot convert argument of type ${handle.$$.smartPtrType?handle.$$.smartPtrType.name:handle.$$.ptrType.name} to parameter type ${this.name}`)}var handleClass=handle.$$.ptrType.registeredClass;ptr=upcastPointer(handle.$$.ptr,handleClass,this.registeredClass);if(this.isSmartPointer){if(undefined===handle.$$.smartPtr){throwBindingError(\"Passing raw pointer to smart pointer is illegal\")}switch(this.sharingPolicy){case 0:if(handle.$$.smartPtrType===this){ptr=handle.$$.smartPtr}else{throwBindingError(`Cannot convert argument of type ${handle.$$.smartPtrType?handle.$$.smartPtrType.name:handle.$$.ptrType.name} to parameter type ${this.name}`)}break;case 1:ptr=handle.$$.smartPtr;break;case 2:if(handle.$$.smartPtrType===this){ptr=handle.$$.smartPtr}else{var clonedHandle=handle[\"clone\"]();ptr=this.rawShare(ptr,Emval.toHandle(()=>clonedHandle[\"delete\"]()));if(destructors!==null){destructors.push(this.rawDestructor,ptr)}}break;default:throwBindingError(\"Unsupporting sharing policy\")}}return ptr}function nonConstNoSmartPtrRawPointerToWireType(destructors,handle){if(handle===null){if(this.isReference){throwBindingError(`null is not a valid ${this.name}`)}return 0}if(!handle.$$){throwBindingError(`Cannot pass \"${embindRepr(handle)}\" as a ${this.name}`)}if(!handle.$$.ptr){throwBindingError(`Cannot pass deleted object as a pointer of type ${this.name}`)}if(handle.$$.ptrType.isConst){throwBindingError(`Cannot convert argument of type ${handle.$$.ptrType.name} to parameter type ${this.name}`)}var handleClass=handle.$$.ptrType.registeredClass;var ptr=upcastPointer(handle.$$.ptr,handleClass,this.registeredClass);return ptr}function readPointer(pointer){return this[\"fromWireType\"](HEAPU32[pointer>>2])}var init_RegisteredPointer=()=>{Object.assign(RegisteredPointer.prototype,{getPointee(ptr){if(this.rawGetPointee){ptr=this.rawGetPointee(ptr)}return ptr},destructor(ptr){this.rawDestructor?.(ptr)},argPackAdvance:GenericWireTypeSize,readValueFromPointer:readPointer,fromWireType:RegisteredPointer_fromWireType})};function RegisteredPointer(name,registeredClass,isReference,isConst,isSmartPointer,pointeeType,sharingPolicy,rawGetPointee,rawConstructor,rawShare,rawDestructor){this.name=name;this.registeredClass=registeredClass;this.isReference=isReference;this.isConst=isConst;this.isSmartPointer=isSmartPointer;this.pointeeType=pointeeType;this.sharingPolicy=sharingPolicy;this.rawGetPointee=rawGetPointee;this.rawConstructor=rawConstructor;this.rawShare=rawShare;this.rawDestructor=rawDestructor;if(!isSmartPointer&®isteredClass.baseClass===undefined){if(isConst){this[\"toWireType\"]=constNoSmartPtrRawPointerToWireType;this.destructorFunction=null}else{this[\"toWireType\"]=nonConstNoSmartPtrRawPointerToWireType;this.destructorFunction=null}}else{this[\"toWireType\"]=genericPointerToWireType}}var replacePublicSymbol=(name,value,numArguments)=>{if(!Module.hasOwnProperty(name)){throwInternalError(\"Replacing nonexistent public symbol\")}if(undefined!==Module[name].overloadTable&&undefined!==numArguments){Module[name].overloadTable[numArguments]=value}else{Module[name]=value;Module[name].argCount=numArguments}};var dynCallLegacy=(sig,ptr,args)=>{sig=sig.replace(/p/g,\"i\");var f=Module[\"dynCall_\"+sig];return f(ptr,...args)};var wasmTableMirror=[];var wasmTable;var getWasmTableEntry=funcPtr=>{var func=wasmTableMirror[funcPtr];if(!func){if(funcPtr>=wasmTableMirror.length)wasmTableMirror.length=funcPtr+1;wasmTableMirror[funcPtr]=func=wasmTable.get(funcPtr)}return func};var dynCall=(sig,ptr,args=[])=>{if(sig.includes(\"j\")){return dynCallLegacy(sig,ptr,args)}var rtn=getWasmTableEntry(ptr)(...args);return rtn};var getDynCaller=(sig,ptr)=>(...args)=>dynCall(sig,ptr,args);var embind__requireFunction=(signature,rawFunction)=>{signature=readLatin1String(signature);function makeDynCaller(){if(signature.includes(\"j\")){return getDynCaller(signature,rawFunction)}return getWasmTableEntry(rawFunction)}var fp=makeDynCaller();if(typeof fp!=\"function\"){throwBindingError(`unknown function pointer with signature ${signature}: ${rawFunction}`)}return fp};var extendError=(baseErrorType,errorName)=>{var errorClass=createNamedFunction(errorName,function(message){this.name=errorName;this.message=message;var stack=new Error(message).stack;if(stack!==undefined){this.stack=this.toString()+\"\\n\"+stack.replace(/^Error(:[^\\n]*)?\\n/,\"\")}});errorClass.prototype=Object.create(baseErrorType.prototype);errorClass.prototype.constructor=errorClass;errorClass.prototype.toString=function(){if(this.message===undefined){return this.name}else{return`${this.name}: ${this.message}`}};return errorClass};var UnboundTypeError;var getTypeName=type=>{var ptr=___getTypeName(type);var rv=readLatin1String(ptr);_free(ptr);return rv};var throwUnboundTypeError=(message,types)=>{var unboundTypes=[];var seen={};function visit(type){if(seen[type]){return}if(registeredTypes[type]){return}if(typeDependencies[type]){typeDependencies[type].forEach(visit);return}unboundTypes.push(type);seen[type]=true}types.forEach(visit);throw new UnboundTypeError(`${message}: `+unboundTypes.map(getTypeName).join([\", \"]))};var __embind_register_class=(rawType,rawPointerType,rawConstPointerType,baseClassRawType,getActualTypeSignature,getActualType,upcastSignature,upcast,downcastSignature,downcast,name,destructorSignature,rawDestructor)=>{name=readLatin1String(name);getActualType=embind__requireFunction(getActualTypeSignature,getActualType);upcast&&=embind__requireFunction(upcastSignature,upcast);downcast&&=embind__requireFunction(downcastSignature,downcast);rawDestructor=embind__requireFunction(destructorSignature,rawDestructor);var legalFunctionName=makeLegalFunctionName(name);exposePublicSymbol(legalFunctionName,function(){throwUnboundTypeError(`Cannot construct ${name} due to unbound types`,[baseClassRawType])});whenDependentTypesAreResolved([rawType,rawPointerType,rawConstPointerType],baseClassRawType?[baseClassRawType]:[],base=>{base=base[0];var baseClass;var basePrototype;if(baseClassRawType){baseClass=base.registeredClass;basePrototype=baseClass.instancePrototype}else{basePrototype=ClassHandle.prototype}var constructor=createNamedFunction(name,function(...args){if(Object.getPrototypeOf(this)!==instancePrototype){throw new BindingError(\"Use 'new' to construct \"+name)}if(undefined===registeredClass.constructor_body){throw new BindingError(name+\" has no accessible constructor\")}var body=registeredClass.constructor_body[args.length];if(undefined===body){throw new BindingError(`Tried to invoke ctor of ${name} with invalid number of parameters (${args.length}) - expected (${Object.keys(registeredClass.constructor_body).toString()}) parameters instead!`)}return body.apply(this,args)});var instancePrototype=Object.create(basePrototype,{constructor:{value:constructor}});constructor.prototype=instancePrototype;var registeredClass=new RegisteredClass(name,constructor,instancePrototype,rawDestructor,baseClass,getActualType,upcast,downcast);if(registeredClass.baseClass){registeredClass.baseClass.__derivedClasses??=[];registeredClass.baseClass.__derivedClasses.push(registeredClass)}var referenceConverter=new RegisteredPointer(name,registeredClass,true,false,false);var pointerConverter=new RegisteredPointer(name+\"*\",registeredClass,false,false,false);var constPointerConverter=new RegisteredPointer(name+\" const*\",registeredClass,false,true,false);registeredPointers[rawType]={pointerType:pointerConverter,constPointerType:constPointerConverter};replacePublicSymbol(legalFunctionName,constructor);return[referenceConverter,pointerConverter,constPointerConverter]})};var heap32VectorToArray=(count,firstElement)=>{var array=[];for(var i=0;i>2])}return array};var runDestructors=destructors=>{while(destructors.length){var ptr=destructors.pop();var del=destructors.pop();del(ptr)}};function usesDestructorStack(argTypes){for(var i=1;i0?\", \":\"\")+argsListWired}invokerFnBody+=(returns||isAsync?\"var rv = \":\"\")+\"invoker(fn\"+(argsListWired.length>0?\", \":\"\")+argsListWired+\");\\n\";if(needsDestructorStack){invokerFnBody+=\"runDestructors(destructors);\\n\"}else{for(var i=isClassMethodFunc?1:2;i{var rawArgTypes=heap32VectorToArray(argCount,rawArgTypesAddr);invoker=embind__requireFunction(invokerSignature,invoker);whenDependentTypesAreResolved([],[rawClassType],classType=>{classType=classType[0];var humanName=`constructor ${classType.name}`;if(undefined===classType.registeredClass.constructor_body){classType.registeredClass.constructor_body=[]}if(undefined!==classType.registeredClass.constructor_body[argCount-1]){throw new BindingError(`Cannot register multiple constructors with identical number of parameters (${argCount-1}) for class '${classType.name}'! Overload resolution is currently only performed using the parameter count, not actual type info!`)}classType.registeredClass.constructor_body[argCount-1]=()=>{throwUnboundTypeError(`Cannot construct ${classType.name} due to unbound types`,rawArgTypes)};whenDependentTypesAreResolved([],rawArgTypes,argTypes=>{argTypes.splice(1,0,null);classType.registeredClass.constructor_body[argCount-1]=craftInvokerFunction(humanName,argTypes,null,invoker,rawConstructor);return[]});return[]})};var getFunctionName=signature=>{signature=signature.trim();const argsIndex=signature.indexOf(\"(\");if(argsIndex!==-1){return signature.substr(0,argsIndex)}else{return signature}};var __embind_register_class_function=(rawClassType,methodName,argCount,rawArgTypesAddr,invokerSignature,rawInvoker,context,isPureVirtual,isAsync)=>{var rawArgTypes=heap32VectorToArray(argCount,rawArgTypesAddr);methodName=readLatin1String(methodName);methodName=getFunctionName(methodName);rawInvoker=embind__requireFunction(invokerSignature,rawInvoker);whenDependentTypesAreResolved([],[rawClassType],classType=>{classType=classType[0];var humanName=`${classType.name}.${methodName}`;if(methodName.startsWith(\"@@\")){methodName=Symbol[methodName.substring(2)]}if(isPureVirtual){classType.registeredClass.pureVirtualFunctions.push(methodName)}function unboundTypesHandler(){throwUnboundTypeError(`Cannot call ${humanName} due to unbound types`,rawArgTypes)}var proto=classType.registeredClass.instancePrototype;var method=proto[methodName];if(undefined===method||undefined===method.overloadTable&&method.className!==classType.name&&method.argCount===argCount-2){unboundTypesHandler.argCount=argCount-2;unboundTypesHandler.className=classType.name;proto[methodName]=unboundTypesHandler}else{ensureOverloadTable(proto,methodName,humanName);proto[methodName].overloadTable[argCount-2]=unboundTypesHandler}whenDependentTypesAreResolved([],rawArgTypes,argTypes=>{var memberFunction=craftInvokerFunction(humanName,argTypes,classType,rawInvoker,context,isAsync);if(undefined===proto[methodName].overloadTable){memberFunction.argCount=argCount-2;proto[methodName]=memberFunction}else{proto[methodName].overloadTable[argCount-2]=memberFunction}return[]});return[]})};var emval_freelist=[];var emval_handles=[];var __emval_decref=handle=>{if(handle>9&&0===--emval_handles[handle+1]){emval_handles[handle]=undefined;emval_freelist.push(handle)}};var count_emval_handles=()=>emval_handles.length/2-5-emval_freelist.length;var init_emval=()=>{emval_handles.push(0,1,undefined,1,null,1,true,1,false,1);Module[\"count_emval_handles\"]=count_emval_handles};var Emval={toValue:handle=>{if(!handle){throwBindingError(\"Cannot use deleted val. handle = \"+handle)}return emval_handles[handle]},toHandle:value=>{switch(value){case undefined:return 2;case null:return 4;case true:return 6;case false:return 8;default:{const handle=emval_freelist.pop()||emval_handles.length;emval_handles[handle]=value;emval_handles[handle+1]=1;return handle}}}};var EmValType={name:\"emscripten::val\",fromWireType:handle=>{var rv=Emval.toValue(handle);__emval_decref(handle);return rv},toWireType:(destructors,value)=>Emval.toHandle(value),argPackAdvance:GenericWireTypeSize,readValueFromPointer:readPointer,destructorFunction:null};var __embind_register_emval=rawType=>registerType(rawType,EmValType);var embindRepr=v=>{if(v===null){return\"null\"}var t=typeof v;if(t===\"object\"||t===\"array\"||t===\"function\"){return v.toString()}else{return\"\"+v}};var floatReadValueFromPointer=(name,width)=>{switch(width){case 4:return function(pointer){return this[\"fromWireType\"](HEAPF32[pointer>>2])};case 8:return function(pointer){return this[\"fromWireType\"](HEAPF64[pointer>>3])};default:throw new TypeError(`invalid float width (${width}): ${name}`)}};var __embind_register_float=(rawType,name,size)=>{name=readLatin1String(name);registerType(rawType,{name:name,fromWireType:value=>value,toWireType:(destructors,value)=>value,argPackAdvance:GenericWireTypeSize,readValueFromPointer:floatReadValueFromPointer(name,size),destructorFunction:null})};var integerReadValueFromPointer=(name,width,signed)=>{switch(width){case 1:return signed?pointer=>HEAP8[pointer]:pointer=>HEAPU8[pointer];case 2:return signed?pointer=>HEAP16[pointer>>1]:pointer=>HEAPU16[pointer>>1];case 4:return signed?pointer=>HEAP32[pointer>>2]:pointer=>HEAPU32[pointer>>2];default:throw new TypeError(`invalid integer width (${width}): ${name}`)}};var __embind_register_integer=(primitiveType,name,size,minRange,maxRange)=>{name=readLatin1String(name);if(maxRange===-1){maxRange=4294967295}var fromWireType=value=>value;if(minRange===0){var bitshift=32-8*size;fromWireType=value=>value<>>bitshift}var isUnsignedType=name.includes(\"unsigned\");var checkAssertions=(value,toTypeName)=>{};var toWireType;if(isUnsignedType){toWireType=function(destructors,value){checkAssertions(value,this.name);return value>>>0}}else{toWireType=function(destructors,value){checkAssertions(value,this.name);return value}}registerType(primitiveType,{name:name,fromWireType:fromWireType,toWireType:toWireType,argPackAdvance:GenericWireTypeSize,readValueFromPointer:integerReadValueFromPointer(name,size,minRange!==0),destructorFunction:null})};var __embind_register_memory_view=(rawType,dataTypeIndex,name)=>{var typeMapping=[Int8Array,Uint8Array,Int16Array,Uint16Array,Int32Array,Uint32Array,Float32Array,Float64Array];var TA=typeMapping[dataTypeIndex];function decodeMemoryView(handle){var size=HEAPU32[handle>>2];var data=HEAPU32[handle+4>>2];return new TA(HEAP8.buffer,data,size)}name=readLatin1String(name);registerType(rawType,{name:name,fromWireType:decodeMemoryView,argPackAdvance:GenericWireTypeSize,readValueFromPointer:decodeMemoryView},{ignoreDuplicateRegistrations:true})};var stringToUTF8Array=(str,heap,outIdx,maxBytesToWrite)=>{if(!(maxBytesToWrite>0))return 0;var startIdx=outIdx;var endIdx=outIdx+maxBytesToWrite-1;for(var i=0;i=55296&&u<=57343){var u1=str.charCodeAt(++i);u=65536+((u&1023)<<10)|u1&1023}if(u<=127){if(outIdx>=endIdx)break;heap[outIdx++]=u}else if(u<=2047){if(outIdx+1>=endIdx)break;heap[outIdx++]=192|u>>6;heap[outIdx++]=128|u&63}else if(u<=65535){if(outIdx+2>=endIdx)break;heap[outIdx++]=224|u>>12;heap[outIdx++]=128|u>>6&63;heap[outIdx++]=128|u&63}else{if(outIdx+3>=endIdx)break;heap[outIdx++]=240|u>>18;heap[outIdx++]=128|u>>12&63;heap[outIdx++]=128|u>>6&63;heap[outIdx++]=128|u&63}}heap[outIdx]=0;return outIdx-startIdx};var stringToUTF8=(str,outPtr,maxBytesToWrite)=>stringToUTF8Array(str,HEAPU8,outPtr,maxBytesToWrite);var lengthBytesUTF8=str=>{var len=0;for(var i=0;i=55296&&c<=57343){len+=4;++i}else{len+=3}}return len};var UTF8Decoder=typeof TextDecoder!=\"undefined\"?new TextDecoder(\"utf8\"):undefined;var UTF8ArrayToString=(heapOrArray,idx,maxBytesToRead)=>{var endIdx=idx+maxBytesToRead;var endPtr=idx;while(heapOrArray[endPtr]&&!(endPtr>=endIdx))++endPtr;if(endPtr-idx>16&&heapOrArray.buffer&&UTF8Decoder){return UTF8Decoder.decode(heapOrArray.subarray(idx,endPtr))}var str=\"\";while(idx>10,56320|ch&1023)}}return str};var UTF8ToString=(ptr,maxBytesToRead)=>ptr?UTF8ArrayToString(HEAPU8,ptr,maxBytesToRead):\"\";var __embind_register_std_string=(rawType,name)=>{name=readLatin1String(name);var stdStringIsUTF8=name===\"std::string\";registerType(rawType,{name:name,fromWireType(value){var length=HEAPU32[value>>2];var payload=value+4;var str;if(stdStringIsUTF8){var decodeStartPtr=payload;for(var i=0;i<=length;++i){var currentBytePtr=payload+i;if(i==length||HEAPU8[currentBytePtr]==0){var maxRead=currentBytePtr-decodeStartPtr;var stringSegment=UTF8ToString(decodeStartPtr,maxRead);if(str===undefined){str=stringSegment}else{str+=String.fromCharCode(0);str+=stringSegment}decodeStartPtr=currentBytePtr+1}}}else{var a=new Array(length);for(var i=0;i>2]=length;if(stdStringIsUTF8&&valueIsOfTypeString){stringToUTF8(value,ptr,length+1)}else{if(valueIsOfTypeString){for(var i=0;i255){_free(ptr);throwBindingError(\"String has UTF-16 code units that do not fit in 8 bits\")}HEAPU8[ptr+i]=charCode}}else{for(var i=0;i{var endPtr=ptr;var idx=endPtr>>1;var maxIdx=idx+maxBytesToRead/2;while(!(idx>=maxIdx)&&HEAPU16[idx])++idx;endPtr=idx<<1;if(endPtr-ptr>32&&UTF16Decoder)return UTF16Decoder.decode(HEAPU8.subarray(ptr,endPtr));var str=\"\";for(var i=0;!(i>=maxBytesToRead/2);++i){var codeUnit=HEAP16[ptr+i*2>>1];if(codeUnit==0)break;str+=String.fromCharCode(codeUnit)}return str};var stringToUTF16=(str,outPtr,maxBytesToWrite)=>{maxBytesToWrite??=2147483647;if(maxBytesToWrite<2)return 0;maxBytesToWrite-=2;var startPtr=outPtr;var numCharsToWrite=maxBytesToWrite>1]=codeUnit;outPtr+=2}HEAP16[outPtr>>1]=0;return outPtr-startPtr};var lengthBytesUTF16=str=>str.length*2;var UTF32ToString=(ptr,maxBytesToRead)=>{var i=0;var str=\"\";while(!(i>=maxBytesToRead/4)){var utf32=HEAP32[ptr+i*4>>2];if(utf32==0)break;++i;if(utf32>=65536){var ch=utf32-65536;str+=String.fromCharCode(55296|ch>>10,56320|ch&1023)}else{str+=String.fromCharCode(utf32)}}return str};var stringToUTF32=(str,outPtr,maxBytesToWrite)=>{maxBytesToWrite??=2147483647;if(maxBytesToWrite<4)return 0;var startPtr=outPtr;var endPtr=startPtr+maxBytesToWrite-4;for(var i=0;i=55296&&codeUnit<=57343){var trailSurrogate=str.charCodeAt(++i);codeUnit=65536+((codeUnit&1023)<<10)|trailSurrogate&1023}HEAP32[outPtr>>2]=codeUnit;outPtr+=4;if(outPtr+4>endPtr)break}HEAP32[outPtr>>2]=0;return outPtr-startPtr};var lengthBytesUTF32=str=>{var len=0;for(var i=0;i=55296&&codeUnit<=57343)++i;len+=4}return len};var __embind_register_std_wstring=(rawType,charSize,name)=>{name=readLatin1String(name);var decodeString,encodeString,readCharAt,lengthBytesUTF;if(charSize===2){decodeString=UTF16ToString;encodeString=stringToUTF16;lengthBytesUTF=lengthBytesUTF16;readCharAt=pointer=>HEAPU16[pointer>>1]}else if(charSize===4){decodeString=UTF32ToString;encodeString=stringToUTF32;lengthBytesUTF=lengthBytesUTF32;readCharAt=pointer=>HEAPU32[pointer>>2]}registerType(rawType,{name:name,fromWireType:value=>{var length=HEAPU32[value>>2];var str;var decodeStartPtr=value+4;for(var i=0;i<=length;++i){var currentBytePtr=value+4+i*charSize;if(i==length||readCharAt(currentBytePtr)==0){var maxReadBytes=currentBytePtr-decodeStartPtr;var stringSegment=decodeString(decodeStartPtr,maxReadBytes);if(str===undefined){str=stringSegment}else{str+=String.fromCharCode(0);str+=stringSegment}decodeStartPtr=currentBytePtr+charSize}}_free(value);return str},toWireType:(destructors,value)=>{if(!(typeof value==\"string\")){throwBindingError(`Cannot pass non-string to C++ string type ${name}`)}var length=lengthBytesUTF(value);var ptr=_malloc(4+length+charSize);HEAPU32[ptr>>2]=length/charSize;encodeString(value,ptr+4,length+charSize);if(destructors!==null){destructors.push(_free,ptr)}return ptr},argPackAdvance:GenericWireTypeSize,readValueFromPointer:readPointer,destructorFunction(ptr){_free(ptr)}})};var __embind_register_void=(rawType,name)=>{name=readLatin1String(name);registerType(rawType,{isVoid:true,name:name,argPackAdvance:0,fromWireType:()=>undefined,toWireType:(destructors,o)=>undefined})};var __emscripten_memcpy_js=(dest,src,num)=>HEAPU8.copyWithin(dest,src,src+num);var requireRegisteredType=(rawType,humanName)=>{var impl=registeredTypes[rawType];if(undefined===impl){throwBindingError(`${humanName} has unknown type ${getTypeName(rawType)}`)}return impl};var emval_returnValue=(returnType,destructorsRef,handle)=>{var destructors=[];var result=returnType[\"toWireType\"](destructors,handle);if(destructors.length){HEAPU32[destructorsRef>>2]=Emval.toHandle(destructors)}return result};var __emval_as=(handle,returnType,destructorsRef)=>{handle=Emval.toValue(handle);returnType=requireRegisteredType(returnType,\"emval::as\");return emval_returnValue(returnType,destructorsRef,handle)};var emval_symbols={};var getStringOrSymbol=address=>{var symbol=emval_symbols[address];if(symbol===undefined){return readLatin1String(address)}return symbol};var emval_methodCallers=[];var __emval_call_method=(caller,objHandle,methodName,destructorsRef,args)=>{caller=emval_methodCallers[caller];objHandle=Emval.toValue(objHandle);methodName=getStringOrSymbol(methodName);return caller(objHandle,objHandle[methodName],destructorsRef,args)};var emval_addMethodCaller=caller=>{var id=emval_methodCallers.length;emval_methodCallers.push(caller);return id};var emval_lookupTypes=(argCount,argTypes)=>{var a=new Array(argCount);for(var i=0;i>2],\"parameter \"+i)}return a};var reflectConstruct=Reflect.construct;var __emval_get_method_caller=(argCount,argTypes,kind)=>{var types=emval_lookupTypes(argCount,argTypes);var retType=types.shift();argCount--;var functionBody=`return function (obj, func, destructorsRef, args) {\\n`;var offset=0;var argsList=[];if(kind===0){argsList.push(\"obj\")}var params=[\"retType\"];var args=[retType];for(var i=0;it.name).join(\", \")}) => ${retType.name}>`;return emval_addMethodCaller(createNamedFunction(functionName,invokerFunction))};var __emval_get_property=(handle,key)=>{handle=Emval.toValue(handle);key=Emval.toValue(key);return Emval.toHandle(handle[key])};var __emval_incref=handle=>{if(handle>9){emval_handles[handle+1]+=1}};var __emval_new_cstring=v=>Emval.toHandle(getStringOrSymbol(v));var __emval_run_destructors=handle=>{var destructors=Emval.toValue(handle);runDestructors(destructors);__emval_decref(handle)};var __emval_take_value=(type,arg)=>{type=requireRegisteredType(type,\"_emval_take_value\");var v=type[\"readValueFromPointer\"](arg);return Emval.toHandle(v)};var getHeapMax=()=>1073741824;var growMemory=size=>{var b=wasmMemory.buffer;var pages=(size-b.byteLength+65535)/65536;try{wasmMemory.grow(pages);updateMemoryViews();return 1}catch(e){}};var _emscripten_resize_heap=requestedSize=>{var oldSize=HEAPU8.length;requestedSize>>>=0;var maxHeapSize=getHeapMax();if(requestedSize>maxHeapSize){return false}var alignUp=(x,multiple)=>x+(multiple-x%multiple)%multiple;for(var cutDown=1;cutDown<=4;cutDown*=2){var overGrownHeapSize=oldSize*(1+.2/cutDown);overGrownHeapSize=Math.min(overGrownHeapSize,requestedSize+100663296);var newSize=Math.min(maxHeapSize,alignUp(Math.max(requestedSize,overGrownHeapSize),65536));var replacement=growMemory(newSize);if(replacement){return true}}return false};embind_init_charCodes();BindingError=Module[\"BindingError\"]=class BindingError extends Error{constructor(message){super(message);this.name=\"BindingError\"}};InternalError=Module[\"InternalError\"]=class InternalError extends Error{constructor(message){super(message);this.name=\"InternalError\"}};init_ClassHandle();init_embind();init_RegisteredPointer();UnboundTypeError=Module[\"UnboundTypeError\"]=extendError(Error,\"UnboundTypeError\");init_emval();var wasmImports={n:___cxa_throw,t:__abort_js,s:__embind_register_bigint,x:__embind_register_bool,r:__embind_register_class,q:__embind_register_class_constructor,k:__embind_register_class_function,w:__embind_register_emval,m:__embind_register_float,c:__embind_register_integer,a:__embind_register_memory_view,l:__embind_register_std_string,f:__embind_register_std_wstring,y:__embind_register_void,v:__emscripten_memcpy_js,h:__emval_as,o:__emval_call_method,b:__emval_decref,p:__emval_get_method_caller,i:__emval_get_property,g:__emval_incref,j:__emval_new_cstring,d:__emval_run_destructors,e:__emval_take_value,u:_emscripten_resize_heap};var wasmExports=createWasm();var ___wasm_call_ctors=()=>(___wasm_call_ctors=wasmExports[\"A\"])();var ___getTypeName=a0=>(___getTypeName=wasmExports[\"B\"])(a0);var _malloc=a0=>(_malloc=wasmExports[\"D\"])(a0);var _free=a0=>(_free=wasmExports[\"E\"])(a0);var ___cxa_is_pointer_type=a0=>(___cxa_is_pointer_type=wasmExports[\"F\"])(a0);Module[\"addOnPostRun\"]=addOnPostRun;var calledRun;dependenciesFulfilled=function runCaller(){if(!calledRun)run();if(!calledRun)dependenciesFulfilled=runCaller};function run(){if(runDependencies>0){return}preRun();if(runDependencies>0){return}function doRun(){if(calledRun)return;calledRun=true;Module[\"calledRun\"]=true;if(ABORT)return;initRuntime();readyPromiseResolve(Module);if(Module[\"onRuntimeInitialized\"])Module[\"onRuntimeInitialized\"]();postRun()}if(Module[\"setStatus\"]){Module[\"setStatus\"](\"Running...\");setTimeout(function(){setTimeout(function(){Module[\"setStatus\"](\"\")},1);doRun()},1)}else{doRun()}}if(Module[\"preInit\"]){if(typeof Module[\"preInit\"]==\"function\")Module[\"preInit\"]=[Module[\"preInit\"]];while(Module[\"preInit\"].length>0){Module[\"preInit\"].pop()()}}run();moduleRtn=readyPromise;\n\n\n return moduleRtn;\n}\n);\n})();\nexport default Module;\n" }, { "path": "viser/src/viser/client/src/Splatting/WasmSorter/build.sh", "content": "#!/usr/bin/env bash\n\nemcc --bind -O3 sorter.cpp -o Sorter.mjs -s WASM=1 -s NO_EXIT_RUNTIME=1 -s \"EXPORTED_RUNTIME_METHODS=['addOnPostRun']\" -s ALLOW_MEMORY_GROWTH=1 -s MAXIMUM_MEMORY=1GB -s STACK_SIZE=2097152 -msimd128;\n" }, { "path": "viser/src/viser/client/src/Splatting/WasmSorter/sorter.cpp", "content": "#include \n#include \n#include \n\n#include \n#include \n#include \n#include \n\n/** SIMD dot product between two 4D vectors. */\n__attribute__((always_inline)) inline float\ndot_f32x4(const v128_t &a, const v128_t &b) {\n v128_t product = wasm_f32x4_mul(a, b);\n v128_t temp = wasm_f32x4_add(\n product, wasm_i32x4_shuffle(product, product, 1, 0, 3, 2)\n );\n v128_t tmp =\n wasm_f32x4_add(temp, wasm_i32x4_shuffle(temp, temp, 2, 3, 0, 1));\n return wasm_f32x4_extract_lane(tmp, 0);\n}\n\n// Function to find the minimum value across a v128_t i32x4 vector.\n__attribute__((always_inline)) inline int32_t min_i32x4(v128_t vector) {\n int32_t elem0 = wasm_i32x4_extract_lane(vector, 0);\n int32_t elem1 = wasm_i32x4_extract_lane(vector, 1);\n int32_t elem2 = wasm_i32x4_extract_lane(vector, 2);\n int32_t elem3 = wasm_i32x4_extract_lane(vector, 3);\n return std::min({elem0, elem1, elem2, elem3});\n}\n\n// Function to find the maximum value across a v128_t i32x4 vector and return it\n// as a float.\n__attribute__((always_inline)) inline int32_t max_i32x4(v128_t vector) {\n int32_t elem0 = wasm_i32x4_extract_lane(vector, 0);\n int32_t elem1 = wasm_i32x4_extract_lane(vector, 1);\n int32_t elem2 = wasm_i32x4_extract_lane(vector, 2);\n int32_t elem3 = wasm_i32x4_extract_lane(vector, 3);\n return std::max({elem0, elem1, elem2, elem3});\n}\n\nclass Sorter {\n std::vector centers_homog; // Centers as homogeneous coordinates.\n std::vector group_indices;\n std::vector sorted_indices;\n\n public:\n Sorter(\n const emscripten::val &buffer, const emscripten::val &group_indices_val\n ) {\n const std::vector bufferVec =\n emscripten::convertJSArrayToNumberVector(buffer);\n const float *floatBuffer =\n reinterpret_cast(bufferVec.data());\n const int32_t num_gaussians = bufferVec.size() / 8;\n sorted_indices.resize(num_gaussians);\n centers_homog.resize(num_gaussians);\n for (int32_t i = 0; i < num_gaussians; i++) {\n centers_homog[i] = wasm_f32x4_make(\n floatBuffer[i * 8 + 0],\n floatBuffer[i * 8 + 1],\n floatBuffer[i * 8 + 2],\n 1.0\n );\n }\n group_indices =\n emscripten::convertJSArrayToNumberVector(group_indices_val\n );\n };\n\n // Run sorting using the newest view projection matrix. Mutates internal\n // buffers.\n emscripten::val sort(const emscripten::val &Tz_cam_groups_val) {\n const auto Tz_cam_groups_buffer =\n emscripten::convertJSArrayToNumberVector(Tz_cam_groups_val);\n const int32_t num_gaussians = centers_homog.size();\n\n // We do a 16-bit counting sort. This is mostly translated from Kevin\n // Kwok's Javascript implementation:\n // https://github.com/antimatter15/splat/blob/main/main.js\n //\n // Note: we want to sort from minimum Z (high depth) to maximum Z (low\n // depth).\n const int32_t padded_length = std::ceil(num_gaussians / 4.0);\n std::vector gaussian_zs(padded_length);\n std::array counts0({0});\n std::array starts0({0});\n\n const int32_t num_groups = Tz_cam_groups_buffer.size() / 4;\n std::vector Tz_cam_groups(num_groups);\n\n const v128_t row3 = wasm_f32x4_make(0.0, 0.0, 0.0, 1.0);\n for (int32_t i = 0; i < num_groups; i++) {\n Tz_cam_groups[i] = wasm_v128_load(&Tz_cam_groups_buffer[i * 4]);\n }\n\n v128_t min_z_i32x4;\n v128_t max_z_i32x4;\n const v128_t splat4096 = wasm_f32x4_splat(4096.0);\n for (int32_t i = 0; i < padded_length; i++) {\n // This should get inlined.\n int32_t gaussianIndex = i * 4;\n const float z0 = dot_f32x4(\n Tz_cam_groups[group_indices[gaussianIndex]],\n centers_homog[gaussianIndex]\n );\n gaussianIndex++;\n const float z1 = dot_f32x4(\n Tz_cam_groups[group_indices[gaussianIndex]],\n centers_homog[gaussianIndex]\n );\n gaussianIndex++;\n const float z2 = dot_f32x4(\n Tz_cam_groups[group_indices[gaussianIndex]],\n centers_homog[gaussianIndex]\n );\n gaussianIndex++;\n const float z3 = dot_f32x4(\n Tz_cam_groups[group_indices[gaussianIndex]],\n centers_homog[gaussianIndex]\n );\n const v128_t cam_z = wasm_f32x4_make(z0, z1, z2, z3);\n\n // OpenGL camera convention: -Z is forward.\n const v128_t depth = wasm_f32x4_neg(cam_z);\n const v128_t z_int =\n wasm_i32x4_trunc_sat_f32x4(wasm_f32x4_mul(cam_z, splat4096));\n gaussian_zs[i] = z_int;\n\n if (i == 0) {\n min_z_i32x4 = z_int;\n max_z_i32x4 = z_int;\n } else {\n // Currently, we incorrectly include padding elements in the\n // min/max.\n min_z_i32x4 = wasm_i32x4_min(min_z_i32x4, z_int);\n max_z_i32x4 = wasm_i32x4_max(max_z_i32x4, z_int);\n }\n }\n min_z_i32x4 = wasm_i32x4_splat(min_i32x4(min_z_i32x4));\n max_z_i32x4 = wasm_i32x4_splat(max_i32x4(max_z_i32x4));\n const v128_t z_inv = wasm_f32x4_div(\n wasm_f32x4_splat(256 * 256 - 1),\n wasm_f32x4_add(\n wasm_f32x4_convert_i32x4(\n wasm_i32x4_sub(max_z_i32x4, min_z_i32x4)\n ),\n wasm_f32x4_splat(1e-5f)\n )\n );\n for (int32_t i = 0; i < padded_length; i++) {\n const v128_t z_bin = wasm_i32x4_trunc_sat_f32x4(wasm_f32x4_mul(\n wasm_f32x4_convert_i32x4(\n wasm_i32x4_sub(gaussian_zs[i], min_z_i32x4)\n ),\n z_inv\n ));\n gaussian_zs[i] = z_bin;\n counts0[wasm_i32x4_extract_lane(z_bin, 0)]++;\n if (i == padded_length - 1) {\n if (i * 4 + 1 < num_gaussians)\n counts0[wasm_i32x4_extract_lane(z_bin, 1)]++;\n if (i * 4 + 2 < num_gaussians)\n counts0[wasm_i32x4_extract_lane(z_bin, 2)]++;\n if (i * 4 + 3 < num_gaussians)\n counts0[wasm_i32x4_extract_lane(z_bin, 3)]++;\n } else {\n counts0[wasm_i32x4_extract_lane(z_bin, 1)]++;\n counts0[wasm_i32x4_extract_lane(z_bin, 2)]++;\n counts0[wasm_i32x4_extract_lane(z_bin, 3)]++;\n }\n }\n for (int32_t i = 1; i < 256 * 256; i++) {\n starts0[i] = starts0[i - 1] + counts0[i - 1];\n }\n\n // Update and return sorted indices.\n for (int32_t i = 0; i < num_gaussians; i++)\n sorted_indices[starts0[((int32_t *)&gaussian_zs[0])[i]]++] = i;\n return emscripten::val(emscripten::typed_memory_view(\n sorted_indices.size(), &(sorted_indices[0])\n ));\n }\n};\n\nEMSCRIPTEN_BINDINGS(c) {\n emscripten::class_(\"Sorter\")\n .constructor()\n .function(\"sort\", &Sorter::sort, emscripten::allow_raw_pointers());\n};\n" }, { "path": "viser/src/viser/client/src/ThreeAssets.tsx", "content": "import { Instance, Instances, shaderMaterial } from \"@react-three/drei\";\nimport { createPortal, useFrame, useThree } from \"@react-three/fiber\";\nimport { Outlines } from \"./Outlines\";\nimport React from \"react\";\nimport * as THREE from \"three\";\nimport { GLTF, GLTFLoader } from \"three/examples/jsm/loaders/GLTFLoader\";\nimport {\n MeshBasicMaterial,\n MeshDepthMaterial,\n MeshDistanceMaterial,\n MeshLambertMaterial,\n MeshMatcapMaterial,\n MeshNormalMaterial,\n MeshPhongMaterial,\n MeshPhysicalMaterial,\n MeshStandardMaterial,\n MeshToonMaterial,\n ShadowMaterial,\n SpriteMaterial,\n RawShaderMaterial,\n ShaderMaterial,\n PointsMaterial,\n LineBasicMaterial,\n LineDashedMaterial,\n} from \"three\";\nimport { DRACOLoader } from \"three/examples/jsm/loaders/DRACOLoader\";\n\ntype AllPossibleThreeJSMaterials =\n | MeshBasicMaterial\n | MeshDepthMaterial\n | MeshDistanceMaterial\n | MeshLambertMaterial\n | MeshMatcapMaterial\n | MeshNormalMaterial\n | MeshPhongMaterial\n | MeshPhysicalMaterial\n | MeshStandardMaterial\n | MeshToonMaterial\n | ShadowMaterial\n | SpriteMaterial\n | RawShaderMaterial\n | ShaderMaterial\n | PointsMaterial\n | LineBasicMaterial\n | LineDashedMaterial;\n\nconst originGeom = new THREE.SphereGeometry(1.0);\nconst originMaterial = new THREE.MeshBasicMaterial({ color: 0xecec00 });\n\nconst PointCloudMaterial = /* @__PURE__ */ shaderMaterial(\n { scale: 1.0, point_ball_norm: 0.0 },\n `\n varying vec3 vPosition;\n varying vec3 vColor; // in the vertex shader\n uniform float scale;\n\n void main() {\n vPosition = position;\n vColor = color;\n vec4 world_pos = modelViewMatrix * vec4(position, 1.0);\n gl_Position = projectionMatrix * world_pos;\n gl_PointSize = (scale / -world_pos.z);\n }\n `,\n `varying vec3 vPosition;\n varying vec3 vColor;\n uniform float point_ball_norm;\n\n void main() {\n if (point_ball_norm < 1000.0) {\n float r = pow(\n pow(abs(gl_PointCoord.x - 0.5), point_ball_norm)\n + pow(abs(gl_PointCoord.y - 0.5), point_ball_norm),\n 1.0 / point_ball_norm);\n if (r > 0.5) discard;\n }\n gl_FragColor = vec4(vColor, 1.0);\n }\n `,\n);\n\nexport const PointCloud = React.forwardRef<\n THREE.Points,\n {\n pointSize: number;\n /** We visualize each point as a 2D ball, which is defined by some norm. */\n pointBallNorm: number;\n points: Float32Array;\n colors: Float32Array;\n }\n>(function PointCloud(props, ref) {\n const getThreeState = useThree((state) => state.get);\n\n const geometry = new THREE.BufferGeometry();\n geometry.setAttribute(\n \"position\",\n new THREE.Float32BufferAttribute(props.points, 3),\n );\n geometry.computeBoundingSphere();\n geometry.setAttribute(\n \"color\",\n new THREE.Float32BufferAttribute(props.colors, 3),\n );\n\n const [material] = React.useState(\n () => new PointCloudMaterial({ vertexColors: true }),\n );\n material.uniforms.scale.value = 10.0;\n material.uniforms.point_ball_norm.value = props.pointBallNorm;\n\n React.useEffect(() => {\n return () => {\n material.dispose();\n geometry.dispose();\n };\n });\n\n const rendererSize = new THREE.Vector2();\n useFrame(() => {\n // Match point scale to behavior of THREE.PointsMaterial().\n if (material === undefined) return;\n // point px height / actual height = point meters height / frustum meters height\n // frustum meters height = math.tan(fov / 2.0) * z\n // point px height = (point meters height / math.tan(fov / 2.0) * actual height) / z\n material.uniforms.scale.value =\n (props.pointSize /\n Math.tan(\n (((getThreeState().camera as THREE.PerspectiveCamera).fov / 180.0) *\n Math.PI) /\n 2.0,\n )) *\n getThreeState().gl.getSize(rendererSize).height *\n getThreeState().gl.getPixelRatio();\n });\n return ;\n});\n\n/** Component for rendering the contents of GLB files. */\nexport const GlbAsset = React.forwardRef<\n THREE.Group,\n { glb_data: Uint8Array; scale: number }\n>(function GlbAsset({ glb_data, scale }, ref) {\n // We track both the GLTF asset itself and all meshes within it. Meshes are\n // used for hover effects.\n const [gltf, setGltf] = React.useState();\n const [meshes, setMeshes] = React.useState([]);\n\n // glTF/GLB files support animations.\n const mixerRef = React.useRef(null);\n\n React.useEffect(() => {\n const loader = new GLTFLoader();\n\n // We use a CDN for Draco. We could move this locally if we want to use Viser offline.\n const dracoLoader = new DRACOLoader();\n dracoLoader.setDecoderPath(\"https://www.gstatic.com/draco/v1/decoders/\");\n loader.setDRACOLoader(dracoLoader);\n\n loader.parse(\n glb_data.buffer,\n \"\",\n (gltf) => {\n if (gltf.animations && gltf.animations.length) {\n mixerRef.current = new THREE.AnimationMixer(gltf.scene);\n gltf.animations.forEach((clip) => {\n mixerRef.current!.clipAction(clip).play();\n });\n }\n const meshes: THREE.Mesh[] = [];\n gltf?.scene.traverse((obj) => {\n if (obj instanceof THREE.Mesh) meshes.push(obj);\n });\n setMeshes(meshes);\n setGltf(gltf);\n },\n (error) => {\n console.log(\"Error loading GLB!\");\n console.log(error);\n },\n );\n\n return () => {\n if (mixerRef.current) mixerRef.current.stopAllAction();\n\n function disposeNode(node: any) {\n if (node instanceof THREE.Mesh) {\n if (node.geometry) {\n node.geometry.dispose();\n }\n if (node.material) {\n if (Array.isArray(node.material)) {\n node.material.forEach((material) => {\n disposeMaterial(material);\n });\n } else {\n disposeMaterial(node.material);\n }\n }\n }\n }\n function disposeMaterial(material: AllPossibleThreeJSMaterials) {\n if (\"map\" in material) material.map?.dispose();\n if (\"lightMap\" in material) material.lightMap?.dispose();\n if (\"bumpMap\" in material) material.bumpMap?.dispose();\n if (\"normalMap\" in material) material.normalMap?.dispose();\n if (\"specularMap\" in material) material.specularMap?.dispose();\n if (\"envMap\" in material) material.envMap?.dispose();\n if (\"alphaMap\" in material) material.alphaMap?.dispose();\n if (\"aoMap\" in material) material.aoMap?.dispose();\n if (\"displacementMap\" in material) material.displacementMap?.dispose();\n if (\"emissiveMap\" in material) material.emissiveMap?.dispose();\n if (\"gradientMap\" in material) material.gradientMap?.dispose();\n if (\"metalnessMap\" in material) material.metalnessMap?.dispose();\n if (\"roughnessMap\" in material) material.roughnessMap?.dispose();\n material.dispose(); // disposes any programs associated with the material\n }\n\n // Attempt to free resources.\n gltf?.scene.traverse(disposeNode);\n };\n }, [glb_data]);\n\n useFrame((_, delta) => {\n if (mixerRef.current) {\n mixerRef.current.update(delta);\n }\n });\n\n return (\n \n {gltf === undefined ? null : (\n <>\n \n {meshes.map((mesh) =>\n createPortal(, mesh),\n )}\n \n )}\n \n );\n});\n\n/** Helper for adding coordinate frames as scene nodes. */\nexport const CoordinateFrame = React.forwardRef<\n THREE.Group,\n {\n showAxes?: boolean;\n axesLength?: number;\n axesRadius?: number;\n originRadius?: number;\n }\n>(function CoordinateFrame(\n {\n showAxes = true,\n axesLength = 0.5,\n axesRadius = 0.0125,\n originRadius = undefined,\n },\n ref,\n) {\n originRadius = originRadius ?? axesRadius * 2;\n return (\n \n {showAxes && (\n <>\n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n \n )}\n \n );\n});\n\n/** Helper for adding batched/instanced coordinate frames as scene nodes. */\nexport const InstancedAxes = React.forwardRef<\n THREE.Group,\n {\n wxyzsBatched: Float32Array;\n positionsBatched: Float32Array;\n axes_length?: number;\n axes_radius?: number;\n }\n>(function InstancedAxes(\n {\n wxyzsBatched: instance_wxyzs,\n positionsBatched: instance_positions,\n axes_length = 0.5,\n axes_radius = 0.0125,\n },\n ref,\n) {\n const axesRef = React.useRef(null);\n\n const cylinderGeom = new THREE.CylinderGeometry(\n axes_radius,\n axes_radius,\n axes_length,\n 16,\n );\n const material = new MeshBasicMaterial();\n\n // Dispose when done.\n React.useEffect(() => {\n return () => {\n cylinderGeom.dispose();\n material.dispose();\n };\n });\n\n // Update instance matrices and colors.\n React.useEffect(() => {\n // Pre-allocate to avoid garbage collector from running during loop.\n const T_world_frame = new THREE.Matrix4();\n const T_world_framex = new THREE.Matrix4();\n const T_world_framey = new THREE.Matrix4();\n const T_world_framez = new THREE.Matrix4();\n\n const T_frame_framex = new THREE.Matrix4()\n .makeRotationFromEuler(new THREE.Euler(0.0, 0.0, (3.0 * Math.PI) / 2.0))\n .setPosition(0.5 * axes_length, 0.0, 0.0);\n const T_frame_framey = new THREE.Matrix4()\n .makeRotationFromEuler(new THREE.Euler(0.0, 0.0, 0.0))\n .setPosition(0.0, 0.5 * axes_length, 0.0);\n const T_frame_framez = new THREE.Matrix4()\n .makeRotationFromEuler(new THREE.Euler(Math.PI / 2.0, 0.0, 0.0))\n .setPosition(0.0, 0.0, 0.5 * axes_length);\n\n const tmpQuat = new THREE.Quaternion();\n\n const red = new THREE.Color(0xcc0000);\n const green = new THREE.Color(0x00cc00);\n const blue = new THREE.Color(0x0000cc);\n\n for (let i = 0; i < instance_wxyzs.length / 4; i++) {\n T_world_frame.makeRotationFromQuaternion(\n tmpQuat.set(\n instance_wxyzs[i * 4 + 1],\n instance_wxyzs[i * 4 + 2],\n instance_wxyzs[i * 4 + 3],\n instance_wxyzs[i * 4 + 0],\n ),\n ).setPosition(\n instance_positions[i * 3 + 0],\n instance_positions[i * 3 + 1],\n instance_positions[i * 3 + 2],\n );\n T_world_framex.copy(T_world_frame).multiply(T_frame_framex);\n T_world_framey.copy(T_world_frame).multiply(T_frame_framey);\n T_world_framez.copy(T_world_frame).multiply(T_frame_framez);\n\n axesRef.current!.setMatrixAt(i * 3 + 0, T_world_framex);\n axesRef.current!.setMatrixAt(i * 3 + 1, T_world_framey);\n axesRef.current!.setMatrixAt(i * 3 + 2, T_world_framez);\n\n axesRef.current!.setColorAt(i * 3 + 0, red);\n axesRef.current!.setColorAt(i * 3 + 1, green);\n axesRef.current!.setColorAt(i * 3 + 2, blue);\n }\n axesRef.current!.instanceMatrix.needsUpdate = true;\n axesRef.current!.instanceColor!.needsUpdate = true;\n }, [instance_wxyzs, instance_positions]);\n\n return (\n \n \n \n \n \n );\n});\n/** Helper for visualizing camera frustums. */\nexport const CameraFrustum = React.forwardRef<\n THREE.Group,\n {\n fov: number;\n aspect: number;\n scale: number;\n color: number;\n thickness?: number; // Added thickness property\n image?: THREE.Texture;\n }\n>(function CameraFrustum(props, ref) {\n let y = Math.tan(props.fov / 2.0);\n let x = y * props.aspect;\n let z = 1.0;\n\n const volumeScale = Math.cbrt((x * y * z) / 3.0);\n x /= volumeScale;\n y /= volumeScale;\n z /= volumeScale;\n\n function scaledLineSegments(points: [number, number, number][], thickness = 1.0) {\n points = points.map((xyz) => [xyz[0] * x, xyz[1] * y, xyz[2] * z]);\n return [...Array(points.length - 1).keys()].map((i) => (\n \n ));\n }\n\n const lineThickness = props.thickness || 1.0; // Default to 1.0 if not provided\n\n return (\n \n \n \n \n {scaledLineSegments(\n [\n // Rectangle.\n [-1, -1, 1],\n [1, -1, 1],\n [1, 1, 1],\n [-1, 1, 1],\n [-1, -1, 1],\n ],\n lineThickness // Pass thickness to scaledLineSegments\n )}\n {scaledLineSegments(\n [\n // Lines to origin.\n [-1, -1, 1],\n [0, 0, 0],\n [1, -1, 1],\n ],\n lineThickness // Pass thickness to scaledLineSegments\n )}\n {scaledLineSegments(\n [\n // Lines to origin.\n [-1, 1, 1],\n [0, 0, 0],\n [1, 1, 1],\n ],\n lineThickness // Pass thickness to scaledLineSegments\n )}\n {scaledLineSegments(\n [\n // Up direction.\n [0.0, -1.2, 1.0],\n [0.0, -0.9, 1.0],\n ],\n lineThickness // Pass thickness to scaledLineSegments\n )}\n \n {props.image && (\n \n \n \n \n )}\n \n );\n});\n\n\nfunction LineSegmentInstance(props: {\n start: THREE.Vector3;\n end: THREE.Vector3;\n color: number;\n thickness?: number; // Optional thickness property\n}) {\n const desiredDirection = new THREE.Vector3()\n .subVectors(props.end, props.start)\n .normalize();\n const canonicalDirection = new THREE.Vector3(0.0, 1.0, 0.0);\n const orientation = new THREE.Quaternion().setFromUnitVectors(\n canonicalDirection,\n desiredDirection,\n );\n\n const length = props.start.distanceTo(props.end);\n const midpoint = new THREE.Vector3()\n .addVectors(props.start, props.end)\n .divideScalar(2.0);\n\n const thickness = props.thickness || 1.0; // Default to 1.0 if not provided\n\n return (\n \n \n \n );\n}\n\nexport const HoverableContext =\n React.createContext | null>(null);\n\n/** Outlines object, which should be placed as a child of all meshes that might\n * be clickable. */\nexport function OutlinesIfHovered(\n props: { alwaysMounted?: boolean; creaseAngle?: number } = {\n // Can be set to true for objects like meshes which may be slow to mount.\n // It seems better to set to False for instanced meshes, there may be some\n // drei or fiber-related race conditions...\n alwaysMounted: false,\n // Some thing just look better with no creasing, like camera frustum objects.\n creaseAngle: Math.PI,\n },\n) {\n const groupRef = React.useRef(null);\n const hoveredRef = React.useContext(HoverableContext);\n const [mounted, setMounted] = React.useState(true);\n\n useFrame(() => {\n if (hoveredRef === null) return;\n if (props.alwaysMounted) {\n if (groupRef.current === null) return;\n groupRef.current.visible = hoveredRef.current;\n } else if (hoveredRef.current != mounted) {\n setMounted(hoveredRef.current);\n }\n });\n return hoveredRef === null || !mounted ? null : (\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/Titlebar.tsx", "content": "import { ViewerContext } from \"./App\";\nimport { ThemeConfigurationMessage } from \"./WebsocketMessages\";\nimport {\n Burger,\n Button,\n Container,\n Group,\n Paper,\n Box,\n useMantineColorScheme,\n Portal,\n} from \"@mantine/core\";\nimport {\n IconBrandGithub,\n IconFileDescription,\n IconKeyboard,\n} from \"@tabler/icons-react\";\nimport { useDisclosure } from \"@mantine/hooks\";\nimport { useContext } from \"react\";\n\n// Type helpers.\ntype ArrayElement =\n ArrayType extends readonly (infer ElementType)[] ? ElementType : never;\ntype TitlebarContent = NonNullable<\n ThemeConfigurationMessage[\"titlebar_content\"]\n>;\nfunction assertUnreachable(x: never): never {\n throw new Error(\"Didn't expect to get here\", x);\n}\n\nfunction getIcon(\n icon: ArrayElement>[\"icon\"],\n) {\n let Icon = null;\n switch (icon) {\n case null:\n break;\n case \"GitHub\":\n Icon = IconBrandGithub;\n break;\n case \"Description\":\n Icon = IconFileDescription;\n break;\n case \"Keyboard\":\n Icon = IconKeyboard;\n break;\n default:\n assertUnreachable(icon);\n }\n return Icon;\n}\n\n// We inherit props directly from message contents.\nexport function TitlebarButton(\n props: ArrayElement>,\n) {\n const Icon = getIcon(props.icon);\n return (\n }\n ml=\"xs\"\n color=\"gray\"\n >\n {props.text}\n \n );\n}\n\nexport function MobileTitlebarButton(\n props: ArrayElement>,\n) {\n const Icon = getIcon(props.icon);\n return (\n }\n ml=\"sm\"\n color=\"gray\"\n >\n {props.text}\n \n );\n}\n\nexport function TitlebarImage(\n props: NonNullable,\n colorScheme: string,\n) {\n let imageSource: string;\n if (props.image_url_dark == null || colorScheme === \"light\") {\n imageSource = props.image_url_light;\n } else {\n imageSource = props.image_url_dark;\n }\n const image = (\n \n );\n\n if (props.href == null) {\n return image;\n }\n return (\n \n {image}\n \n );\n}\n\nexport function Titlebar() {\n const viewer = useContext(ViewerContext)!;\n const content = viewer.useGui((state) => state.theme.titlebar_content);\n const colorScheme = useMantineColorScheme().colorScheme;\n\n const [burgerOpen, burgerHandlers] = useDisclosure(false);\n\n if (content == null) {\n return null;\n }\n\n const buttons = content.buttons;\n const imageData = content.image;\n\n return (\n \n \n ({\n height: \"100%\",\n display: \"flex\",\n alignItems: \"center\",\n })}\n >\n ({ marginRight: \"auto\" })}>\n {imageData !== null ? TitlebarImage(imageData, colorScheme) : null}\n \n ({\n flexWrap: \"nowrap\",\n overflowX: \"scroll\",\n msOverflowStyle: \"none\",\n scrollbarWidth: \"none\",\n })}\n >\n {buttons?.map((btn, index) => (\n \n ))}\n \n \n \n \n \n {buttons?.map((btn, index) => (\n \n ))}\n \n \n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/Utils.ts", "content": "// Drag Utils\nexport interface DragEvents {\n move: \"touchmove\" | \"mousemove\";\n end: \"touchend\" | \"mouseup\";\n}\nexport const touchEvents: DragEvents = { move: \"touchmove\", end: \"touchend\" };\nexport const mouseEvents: DragEvents = { move: \"mousemove\", end: \"mouseup\" };\n\nexport function isTouchEvent(\n event: TouchEvent | MouseEvent,\n): event is TouchEvent {\n return event.type === \"touchmove\";\n}\nexport function isMouseEvent(\n event: TouchEvent | MouseEvent,\n): event is MouseEvent {\n return event.type === \"mousemove\";\n}\n" }, { "path": "viser/src/viser/client/src/WebsocketFunctions.tsx", "content": "import React from \"react\";\nimport * as THREE from \"three\";\nimport { Message } from \"./WebsocketMessages\";\nimport { ViewerContext, ViewerContextContents } from \"./App\";\n\n/** Easier, hook version of makeThrottledMessageSender. */\nexport function useThrottledMessageSender(throttleMilliseconds: number) {\n const viewer = React.useContext(ViewerContext)!;\n return makeThrottledMessageSender(viewer, throttleMilliseconds);\n}\n\n/** Returns a function for sending messages, with automatic throttling. */\nexport function makeThrottledMessageSender(\n viewer: ViewerContextContents,\n throttleMilliseconds: number,\n) {\n let readyToSend = true;\n let stale = false;\n let latestMessage: Message | null = null;\n\n function send(message: Message) {\n if (viewer.sendMessageRef.current === null) return;\n latestMessage = message;\n if (readyToSend) {\n viewer.sendMessageRef.current(message);\n stale = false;\n readyToSend = false;\n\n setTimeout(() => {\n readyToSend = true;\n if (!stale) return;\n latestMessage && send(latestMessage);\n }, throttleMilliseconds);\n } else {\n stale = true;\n }\n }\n return send;\n}\n\n/** Type guard for threejs textures. Meant to be used with `scene.background`. */\nexport function isTexture(\n background:\n | THREE.Color\n | THREE.Texture\n | THREE.CubeTexture\n | null\n | undefined,\n): background is THREE.Texture {\n return (\n background !== null &&\n background !== undefined &&\n (background as THREE.Texture).isTexture !== undefined\n );\n}\n" }, { "path": "viser/src/viser/client/src/WebsocketInterface.tsx", "content": "import WebsocketServerWorker from \"./WebsocketServerWorker?worker\";\nimport React, { useContext } from \"react\";\n\nimport { ViewerContext } from \"./App\";\nimport { syncSearchParamServer } from \"./SearchParamsUtils\";\nimport { WsWorkerIncoming, WsWorkerOutgoing } from \"./WebsocketServerWorker\";\n\n/** Component for handling websocket connections. */\nexport function WebsocketMessageProducer() {\n const messageQueueRef = useContext(ViewerContext)!.messageQueueRef;\n const viewer = useContext(ViewerContext)!;\n const server = viewer.useGui((state) => state.server);\n const resetGui = viewer.useGui((state) => state.resetGui);\n\n syncSearchParamServer(server);\n\n React.useEffect(() => {\n const worker = new WebsocketServerWorker();\n\n worker.onmessage = (event) => {\n const data: WsWorkerOutgoing = event.data;\n if (data.type === \"connected\") {\n resetGui();\n viewer.useGui.setState({ websocketConnected: true });\n viewer.sendMessageRef.current = (message) => {\n postToWorker({ type: \"send\", message: message });\n };\n } else if (data.type === \"closed\") {\n resetGui();\n viewer.useGui.setState({ websocketConnected: false });\n viewer.sendMessageRef.current = (message) => {\n console.log(\n `Tried to send ${message.type} but websocket is not connected!`,\n );\n };\n } else if (data.type === \"message_batch\") {\n messageQueueRef.current.push(...data.messages);\n }\n };\n function postToWorker(data: WsWorkerIncoming) {\n worker.postMessage(data);\n }\n postToWorker({ type: \"set_server\", server: server });\n return () => {\n postToWorker({ type: \"close\" });\n viewer.sendMessageRef.current = (message) =>\n console.log(\n `Tried to send ${message.type} but websocket is not connected!`,\n );\n viewer.useGui.setState({ websocketConnected: false });\n };\n }, [server, resetGui]);\n\n return <>;\n}\n" }, { "path": "viser/src/viser/client/src/WebsocketMessages.tsx", "content": "// AUTOMATICALLY GENERATED message interfaces, from Python dataclass definitions.\n// This file should not be manually modified.\n/** Message for running some arbitrary Javascript on the client.\n * We use this to set up the Plotly.js package, via the plotly.min.js source\n * code.\n *\n * (automatically generated)\n */\nexport interface RunJavascriptMessage {\n type: \"RunJavascriptMessage\";\n source: string;\n}\n/** Notification message.\n *\n * (automatically generated)\n */\nexport interface NotificationMessage {\n type: \"NotificationMessage\";\n mode: \"show\" | \"update\";\n id: string;\n title: string;\n body: string;\n loading: boolean;\n with_close_button: boolean;\n auto_close: number | false;\n color:\n | \"dark\"\n | \"gray\"\n | \"red\"\n | \"pink\"\n | \"grape\"\n | \"violet\"\n | \"indigo\"\n | \"blue\"\n | \"cyan\"\n | \"green\"\n | \"lime\"\n | \"yellow\"\n | \"orange\"\n | \"teal\"\n | null;\n}\n/** Remove a specific notification.\n *\n * (automatically generated)\n */\nexport interface RemoveNotificationMessage {\n type: \"RemoveNotificationMessage\";\n id: string;\n}\n/** Message for a posed viewer camera.\n * Pose is in the form T_world_camera, OpenCV convention, +Z forward.\n *\n * (automatically generated)\n */\nexport interface ViewerCameraMessage {\n type: \"ViewerCameraMessage\";\n wxyz: [number, number, number, number];\n position: [number, number, number];\n fov: number;\n aspect: number;\n look_at: [number, number, number];\n up_direction: [number, number, number];\n}\n/** Message for a raycast-like pointer in the scene.\n * origin is the viewing camera position, in world coordinates.\n * direction is the vector if a ray is projected from the camera through the clicked pixel,\n *\n *\n * (automatically generated)\n */\nexport interface ScenePointerMessage {\n type: \"ScenePointerMessage\";\n event_type: \"click\" | \"rect-select\";\n ray_origin: [number, number, number] | null;\n ray_direction: [number, number, number] | null;\n screen_pos: [number, number][];\n}\n/** Message to enable/disable scene click events.\n *\n * (automatically generated)\n */\nexport interface ScenePointerEnableMessage {\n type: \"ScenePointerEnableMessage\";\n enable: boolean;\n event_type: \"click\" | \"rect-select\";\n}\n/** Variant of CameraMessage used for visualizing camera frustums.\n *\n * OpenCV convention, +Z forward.\n *\n * (automatically generated)\n */\nexport interface CameraFrustumMessage {\n type: \"CameraFrustumMessage\";\n name: string;\n fov: number;\n aspect: number;\n scale: number;\n color: number;\n thickness: number;\n image_media_type: \"image/jpeg\" | \"image/png\" | null;\n image_binary: Uint8Array | null;\n}\n/** GlTF message.\n *\n * (automatically generated)\n */\nexport interface GlbMessage {\n type: \"GlbMessage\";\n name: string;\n glb_data: Uint8Array;\n scale: number;\n}\n/** Coordinate frame message.\n *\n * (automatically generated)\n */\nexport interface FrameMessage {\n type: \"FrameMessage\";\n name: string;\n show_axes: boolean;\n axes_length: number;\n axes_radius: number;\n origin_radius: number;\n}\n/** Batched axes message.\n *\n * Positions and orientations should follow a `T_parent_local` convention, which\n * corresponds to the R matrix and t vector in `p_parent = [R | t] p_local`.\n *\n * (automatically generated)\n */\nexport interface BatchedAxesMessage {\n type: \"BatchedAxesMessage\";\n name: string;\n wxyzs_batched: Uint8Array;\n positions_batched: Uint8Array;\n axes_length: number;\n axes_radius: number;\n}\n/** Grid message. Helpful for visualizing things like ground planes.\n *\n * (automatically generated)\n */\nexport interface GridMessage {\n type: \"GridMessage\";\n name: string;\n width: number;\n height: number;\n width_segments: number;\n height_segments: number;\n plane: \"xz\" | \"xy\" | \"yx\" | \"yz\" | \"zx\" | \"zy\";\n cell_color: number;\n cell_thickness: number;\n cell_size: number;\n section_color: number;\n section_thickness: number;\n section_size: number;\n}\n/** Add a 2D label to the scene.\n *\n * (automatically generated)\n */\nexport interface LabelMessage {\n type: \"LabelMessage\";\n name: string;\n text: string;\n}\n/** Add a 3D gui element to the scene.\n *\n * (automatically generated)\n */\nexport interface Gui3DMessage {\n type: \"Gui3DMessage\";\n order: number;\n name: string;\n container_id: string;\n}\n/** Point cloud message.\n *\n * Positions are internally canonicalized to float32, colors to uint8.\n *\n * Float color inputs should be in the range [0,1], int color inputs should be in the\n * range [0,255].\n *\n * (automatically generated)\n */\nexport interface PointCloudMessage {\n type: \"PointCloudMessage\";\n name: string;\n points: Uint8Array;\n colors: Uint8Array;\n point_size: number;\n point_ball_norm: number;\n}\n/** Message for a bone of a skinned mesh.\n *\n * (automatically generated)\n */\nexport interface MeshBoneMessage {\n type: \"MeshBoneMessage\";\n name: string;\n}\n/** Mesh message.\n *\n * Vertices are internally canonicalized to float32, faces to uint32.\n *\n * (automatically generated)\n */\nexport interface MeshMessage {\n type: \"MeshMessage\";\n name: string;\n vertices: Uint8Array;\n faces: Uint8Array;\n color: number | null;\n vertex_colors: Uint8Array | null;\n wireframe: boolean;\n opacity: number | null;\n flat_shading: boolean;\n side: \"front\" | \"back\" | \"double\";\n material: \"standard\" | \"toon3\" | \"toon5\";\n}\n/** Mesh message.\n *\n * Vertices are internally canonicalized to float32, faces to uint32.\n *\n * (automatically generated)\n */\nexport interface SkinnedMeshMessage {\n type: \"SkinnedMeshMessage\";\n name: string;\n vertices: Uint8Array;\n faces: Uint8Array;\n color: number | null;\n vertex_colors: Uint8Array | null;\n wireframe: boolean;\n opacity: number | null;\n flat_shading: boolean;\n side: \"front\" | \"back\" | \"double\";\n material: \"standard\" | \"toon3\" | \"toon5\";\n bone_wxyzs: [number, number, number, number][];\n bone_positions: [number, number, number][];\n skin_indices: Uint8Array;\n skin_weights: Uint8Array;\n}\n/** Server -> client message to set a skinned mesh bone's orientation.\n *\n * As with all other messages, transforms take the `T_parent_local` convention.\n *\n * (automatically generated)\n */\nexport interface SetBoneOrientationMessage {\n type: \"SetBoneOrientationMessage\";\n name: string;\n bone_index: number;\n wxyz: [number, number, number, number];\n}\n/** Server -> client message to set a skinned mesh bone's position.\n *\n * As with all other messages, transforms take the `T_parent_local` convention.\n *\n * (automatically generated)\n */\nexport interface SetBonePositionMessage {\n type: \"SetBonePositionMessage\";\n name: string;\n bone_index: number;\n position: [number, number, number];\n}\n/** Message for transform gizmos.\n *\n * (automatically generated)\n */\nexport interface TransformControlsMessage {\n type: \"TransformControlsMessage\";\n name: string;\n scale: number;\n line_width: number;\n fixed: boolean;\n auto_transform: boolean;\n active_axes: [boolean, boolean, boolean];\n disable_axes: boolean;\n disable_sliders: boolean;\n disable_rotations: boolean;\n translation_limits: [[number, number], [number, number], [number, number]];\n rotation_limits: [[number, number], [number, number], [number, number]];\n depth_test: boolean;\n opacity: number;\n}\n/** Server -> client message to set the camera's position.\n *\n * (automatically generated)\n */\nexport interface SetCameraPositionMessage {\n type: \"SetCameraPositionMessage\";\n position: [number, number, number];\n}\n/** Server -> client message to set the camera's up direction.\n *\n * (automatically generated)\n */\nexport interface SetCameraUpDirectionMessage {\n type: \"SetCameraUpDirectionMessage\";\n position: [number, number, number];\n}\n/** Server -> client message to set the camera's look-at point.\n *\n * (automatically generated)\n */\nexport interface SetCameraLookAtMessage {\n type: \"SetCameraLookAtMessage\";\n look_at: [number, number, number];\n}\n/** Server -> client message to set the camera's field of view.\n *\n * (automatically generated)\n */\nexport interface SetCameraFovMessage {\n type: \"SetCameraFovMessage\";\n fov: number;\n}\n/** Server -> client message to set a scene node's orientation.\n *\n * As with all other messages, transforms take the `T_parent_local` convention.\n *\n * (automatically generated)\n */\nexport interface SetOrientationMessage {\n type: \"SetOrientationMessage\";\n name: string;\n wxyz: [number, number, number, number];\n}\n/** Server -> client message to set a scene node's position.\n *\n * As with all other messages, transforms take the `T_parent_local` convention.\n *\n * (automatically generated)\n */\nexport interface SetPositionMessage {\n type: \"SetPositionMessage\";\n name: string;\n position: [number, number, number];\n}\n/** Client -> server message when a transform control is updated.\n *\n * As with all other messages, transforms take the `T_parent_local` convention.\n *\n * (automatically generated)\n */\nexport interface TransformControlsUpdateMessage {\n type: \"TransformControlsUpdateMessage\";\n name: string;\n wxyz: [number, number, number, number];\n position: [number, number, number];\n}\n/** Message for rendering a background image.\n *\n * (automatically generated)\n */\nexport interface BackgroundImageMessage {\n type: \"BackgroundImageMessage\";\n media_type: \"image/jpeg\" | \"image/png\";\n rgb_bytes: Uint8Array;\n depth_bytes: Uint8Array | null;\n}\n/** Message for rendering 2D images.\n *\n * (automatically generated)\n */\nexport interface ImageMessage {\n type: \"ImageMessage\";\n name: string;\n media_type: \"image/jpeg\" | \"image/png\";\n data: Uint8Array;\n render_width: number;\n render_height: number;\n}\n/** Remove a particular node from the scene.\n *\n * (automatically generated)\n */\nexport interface RemoveSceneNodeMessage {\n type: \"RemoveSceneNodeMessage\";\n name: string;\n}\n/** Set the visibility of a particular node in the scene.\n *\n * (automatically generated)\n */\nexport interface SetSceneNodeVisibilityMessage {\n type: \"SetSceneNodeVisibilityMessage\";\n name: string;\n visible: boolean;\n}\n/** Set the clickability of a particular node in the scene.\n *\n * (automatically generated)\n */\nexport interface SetSceneNodeClickableMessage {\n type: \"SetSceneNodeClickableMessage\";\n name: string;\n clickable: boolean;\n}\n/** Message for clicked objects.\n *\n * (automatically generated)\n */\nexport interface SceneNodeClickMessage {\n type: \"SceneNodeClickMessage\";\n name: string;\n instance_index: number | null;\n ray_origin: [number, number, number];\n ray_direction: [number, number, number];\n screen_pos: [number, number];\n}\n/** Reset scene.\n *\n * (automatically generated)\n */\nexport interface ResetSceneMessage {\n type: \"ResetSceneMessage\";\n}\n/** Reset GUI.\n *\n * (automatically generated)\n */\nexport interface ResetGuiMessage {\n type: \"ResetGuiMessage\";\n}\n/** GuiAddFolderMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', expand_by_default: 'bool', visible: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddFolderMessage {\n type: \"GuiAddFolderMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n expand_by_default: boolean;\n visible: boolean;\n}\n/** GuiAddMarkdownMessage(order: 'float', id: 'str', markdown: 'str', container_id: 'str', visible: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddMarkdownMessage {\n type: \"GuiAddMarkdownMessage\";\n order: number;\n id: string;\n markdown: string;\n container_id: string;\n visible: boolean;\n}\n/** GuiAddProgressBarMessage(order: 'float', id: 'str', value: 'float', animated: 'bool', color: 'Optional[Color]', container_id: 'str', visible: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddProgressBarMessage {\n type: \"GuiAddProgressBarMessage\";\n order: number;\n id: string;\n value: number;\n animated: boolean;\n color:\n | \"dark\"\n | \"gray\"\n | \"red\"\n | \"pink\"\n | \"grape\"\n | \"violet\"\n | \"indigo\"\n | \"blue\"\n | \"cyan\"\n | \"green\"\n | \"lime\"\n | \"yellow\"\n | \"orange\"\n | \"teal\"\n | null;\n container_id: string;\n visible: boolean;\n}\n/** GuiAddPlotlyMessage(order: 'float', id: 'str', plotly_json_str: 'str', aspect: 'float', container_id: 'str', visible: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddPlotlyMessage {\n type: \"GuiAddPlotlyMessage\";\n order: number;\n id: string;\n plotly_json_str: string;\n aspect: number;\n container_id: string;\n visible: boolean;\n}\n/** GuiAddTabGroupMessage(order: 'float', id: 'str', container_id: 'str', tab_labels: 'Tuple[str, ...]', tab_icons_html: 'Tuple[Union[str, None], ...]', tab_container_ids: 'Tuple[str, ...]', visible: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddTabGroupMessage {\n type: \"GuiAddTabGroupMessage\";\n order: number;\n id: string;\n container_id: string;\n tab_labels: string[];\n tab_icons_html: (string | null)[];\n tab_container_ids: string[];\n visible: boolean;\n}\n/** Base message type containing fields commonly used by GUI inputs.\n *\n * (automatically generated)\n */\nexport interface _GuiAddInputBase {\n type: \"_GuiAddInputBase\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: any;\n visible: boolean;\n disabled: boolean;\n}\n/** GuiAddButtonMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'bool', visible: 'bool', disabled: 'bool', color: 'Optional[Color]', icon_html: 'Optional[str]')\n *\n * (automatically generated)\n */\nexport interface GuiAddButtonMessage {\n type: \"GuiAddButtonMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: boolean;\n visible: boolean;\n disabled: boolean;\n color:\n | \"dark\"\n | \"gray\"\n | \"red\"\n | \"pink\"\n | \"grape\"\n | \"violet\"\n | \"indigo\"\n | \"blue\"\n | \"cyan\"\n | \"green\"\n | \"lime\"\n | \"yellow\"\n | \"orange\"\n | \"teal\"\n | null;\n icon_html: string | null;\n}\n/** GuiAddUploadButtonMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'Any', visible: 'bool', disabled: 'bool', color: 'Optional[Color]', icon_html: 'Optional[str]', mime_type: 'str')\n *\n * (automatically generated)\n */\nexport interface GuiAddUploadButtonMessage {\n type: \"GuiAddUploadButtonMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: any;\n visible: boolean;\n disabled: boolean;\n color:\n | \"dark\"\n | \"gray\"\n | \"red\"\n | \"pink\"\n | \"grape\"\n | \"violet\"\n | \"indigo\"\n | \"blue\"\n | \"cyan\"\n | \"green\"\n | \"lime\"\n | \"yellow\"\n | \"orange\"\n | \"teal\"\n | null;\n icon_html: string | null;\n mime_type: string;\n}\n/** GuiAddSliderMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'float', visible: 'bool', disabled: 'bool', min: 'float', max: 'float', step: 'Optional[float]', precision: 'int', marks: 'Optional[Tuple[GuiSliderMark, ...]]' = None)\n *\n * (automatically generated)\n */\nexport interface GuiAddSliderMessage {\n type: \"GuiAddSliderMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: number;\n visible: boolean;\n disabled: boolean;\n min: number;\n max: number;\n step: number | null;\n precision: number;\n marks: { value: number; label?: string }[] | null;\n}\n/** GuiAddMultiSliderMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'Any', visible: 'bool', disabled: 'bool', min: 'float', max: 'float', step: 'Optional[float]', min_range: 'Optional[float]', precision: 'int', fixed_endpoints: 'bool' = False, marks: 'Optional[Tuple[GuiSliderMark, ...]]' = None)\n *\n * (automatically generated)\n */\nexport interface GuiAddMultiSliderMessage {\n type: \"GuiAddMultiSliderMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: any;\n visible: boolean;\n disabled: boolean;\n min: number;\n max: number;\n step: number | null;\n min_range: number | null;\n precision: number;\n fixed_endpoints: boolean;\n marks: { value: number; label?: string }[] | null;\n}\n/** GuiAddNumberMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'float', visible: 'bool', disabled: 'bool', precision: 'int', step: 'float', min: 'Optional[float]', max: 'Optional[float]')\n *\n * (automatically generated)\n */\nexport interface GuiAddNumberMessage {\n type: \"GuiAddNumberMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: number;\n visible: boolean;\n disabled: boolean;\n precision: number;\n step: number;\n min: number | null;\n max: number | null;\n}\n/** GuiAddRgbMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'Tuple[int, int, int]', visible: 'bool', disabled: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddRgbMessage {\n type: \"GuiAddRgbMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: [number, number, number];\n visible: boolean;\n disabled: boolean;\n}\n/** GuiAddRgbaMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'Tuple[int, int, int, int]', visible: 'bool', disabled: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddRgbaMessage {\n type: \"GuiAddRgbaMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: [number, number, number, number];\n visible: boolean;\n disabled: boolean;\n}\n/** GuiAddCheckboxMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'bool', visible: 'bool', disabled: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddCheckboxMessage {\n type: \"GuiAddCheckboxMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: boolean;\n visible: boolean;\n disabled: boolean;\n}\n/** GuiAddVector2Message(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'Tuple[float, float]', visible: 'bool', disabled: 'bool', min: 'Optional[Tuple[float, float]]', max: 'Optional[Tuple[float, float]]', step: 'float', precision: 'int')\n *\n * (automatically generated)\n */\nexport interface GuiAddVector2Message {\n type: \"GuiAddVector2Message\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: [number, number];\n visible: boolean;\n disabled: boolean;\n min: [number, number] | null;\n max: [number, number] | null;\n step: number;\n precision: number;\n}\n/** GuiAddVector3Message(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'Tuple[float, float, float]', visible: 'bool', disabled: 'bool', min: 'Optional[Tuple[float, float, float]]', max: 'Optional[Tuple[float, float, float]]', step: 'float', precision: 'int')\n *\n * (automatically generated)\n */\nexport interface GuiAddVector3Message {\n type: \"GuiAddVector3Message\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: [number, number, number];\n visible: boolean;\n disabled: boolean;\n min: [number, number, number] | null;\n max: [number, number, number] | null;\n step: number;\n precision: number;\n}\n/** GuiAddTextMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'str', visible: 'bool', disabled: 'bool')\n *\n * (automatically generated)\n */\nexport interface GuiAddTextMessage {\n type: \"GuiAddTextMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: string;\n visible: boolean;\n disabled: boolean;\n}\n/** GuiAddDropdownMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'str', visible: 'bool', disabled: 'bool', options: 'Tuple[str, ...]')\n *\n * (automatically generated)\n */\nexport interface GuiAddDropdownMessage {\n type: \"GuiAddDropdownMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: string;\n visible: boolean;\n disabled: boolean;\n options: string[];\n}\n/** GuiAddButtonGroupMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', value: 'str', visible: 'bool', disabled: 'bool', options: 'Tuple[str, ...]')\n *\n * (automatically generated)\n */\nexport interface GuiAddButtonGroupMessage {\n type: \"GuiAddButtonGroupMessage\";\n order: number;\n id: string;\n label: string;\n container_id: string;\n hint: string | null;\n value: string;\n visible: boolean;\n disabled: boolean;\n options: string[];\n}\n/** GuiModalMessage(order: 'float', id: 'str', title: 'str')\n *\n * (automatically generated)\n */\nexport interface GuiModalMessage {\n type: \"GuiModalMessage\";\n order: number;\n id: string;\n title: string;\n}\n/** GuiCloseModalMessage(id: 'str')\n *\n * (automatically generated)\n */\nexport interface GuiCloseModalMessage {\n type: \"GuiCloseModalMessage\";\n id: string;\n}\n/** Sent server->client to remove a GUI element.\n *\n * (automatically generated)\n */\nexport interface GuiRemoveMessage {\n type: \"GuiRemoveMessage\";\n id: string;\n}\n/** Sent client<->server when any property of a GUI component is changed.\n *\n * (automatically generated)\n */\nexport interface GuiUpdateMessage {\n type: \"GuiUpdateMessage\";\n id: string;\n updates: Partial;\n}\n/** Message from server->client to configure parts of the GUI.\n *\n * (automatically generated)\n */\nexport interface ThemeConfigurationMessage {\n type: \"ThemeConfigurationMessage\";\n titlebar_content: {\n buttons:\n | {\n text: string | null;\n icon: \"GitHub\" | \"Description\" | \"Keyboard\" | null;\n href: string | null;\n }[]\n | null;\n image: {\n image_url_light: string;\n image_url_dark: string | null;\n image_alt: string;\n href: string | null;\n } | null;\n } | null;\n control_layout: \"floating\" | \"collapsible\" | \"fixed\";\n control_width: \"small\" | \"medium\" | \"large\";\n show_logo: boolean;\n show_share_button: boolean;\n dark_mode: boolean;\n colors:\n | [\n string,\n string,\n string,\n string,\n string,\n string,\n string,\n string,\n string,\n string,\n ]\n | null;\n}\n/** Message from server->client carrying Catmull-Rom spline information.\n *\n * (automatically generated)\n */\nexport interface CatmullRomSplineMessage {\n type: \"CatmullRomSplineMessage\";\n name: string;\n positions: [number, number, number][];\n curve_type: \"centripetal\" | \"chordal\" | \"catmullrom\";\n tension: number;\n closed: boolean;\n line_width: number;\n color: number;\n segments: number | null;\n}\n/** Message from server->client carrying Cubic Bezier spline information.\n *\n * (automatically generated)\n */\nexport interface CubicBezierSplineMessage {\n type: \"CubicBezierSplineMessage\";\n name: string;\n positions: [number, number, number][];\n control_points: [number, number, number][];\n line_width: number;\n color: number;\n segments: number | null;\n}\n/** Message from server->client carrying splattable Gaussians.\n *\n * (automatically generated)\n */\nexport interface GaussianSplatsMessage {\n type: \"GaussianSplatsMessage\";\n name: string;\n buffer: Uint8Array;\n}\n/** Message from server->client requesting a render of the current viewport.\n *\n * (automatically generated)\n */\nexport interface GetRenderRequestMessage {\n type: \"GetRenderRequestMessage\";\n format: \"image/jpeg\" | \"image/png\";\n height: number;\n width: number;\n quality: number;\n}\n/** Message from client->server carrying a render.\n *\n * (automatically generated)\n */\nexport interface GetRenderResponseMessage {\n type: \"GetRenderResponseMessage\";\n payload: Uint8Array;\n}\n/** Signal that a file is about to be sent.\n *\n * (automatically generated)\n */\nexport interface FileTransferStart {\n type: \"FileTransferStart\";\n source_component_id: string | null;\n transfer_uuid: string;\n filename: string;\n mime_type: string;\n part_count: number;\n size_bytes: number;\n}\n/** Send a file for clients to download or upload files from client.\n *\n * (automatically generated)\n */\nexport interface FileTransferPart {\n type: \"FileTransferPart\";\n source_component_id: string | null;\n transfer_uuid: string;\n part: number;\n content: Uint8Array;\n}\n/** Send a file for clients to download or upload files from client.\n *\n * (automatically generated)\n */\nexport interface FileTransferPartAck {\n type: \"FileTransferPartAck\";\n source_component_id: string | null;\n transfer_uuid: string;\n transferred_bytes: number;\n total_bytes: number;\n}\n/** Message from client->server to connect to the share URL server.\n *\n * (automatically generated)\n */\nexport interface ShareUrlRequest {\n type: \"ShareUrlRequest\";\n}\n/** Message from server->client to indicate that the share URL has been updated.\n *\n * (automatically generated)\n */\nexport interface ShareUrlUpdated {\n type: \"ShareUrlUpdated\";\n share_url: string | null;\n}\n/** Message from client->server to disconnect from the share URL server.\n *\n * (automatically generated)\n */\nexport interface ShareUrlDisconnect {\n type: \"ShareUrlDisconnect\";\n}\n/** Message from server->client to set the label of the GUI panel.\n *\n * (automatically generated)\n */\nexport interface SetGuiPanelLabelMessage {\n type: \"SetGuiPanelLabelMessage\";\n label: string | null;\n}\n\nexport type Message =\n | RunJavascriptMessage\n | NotificationMessage\n | RemoveNotificationMessage\n | ViewerCameraMessage\n | ScenePointerMessage\n | ScenePointerEnableMessage\n | CameraFrustumMessage\n | GlbMessage\n | FrameMessage\n | BatchedAxesMessage\n | GridMessage\n | LabelMessage\n | Gui3DMessage\n | PointCloudMessage\n | MeshBoneMessage\n | MeshMessage\n | SkinnedMeshMessage\n | SetBoneOrientationMessage\n | SetBonePositionMessage\n | TransformControlsMessage\n | SetCameraPositionMessage\n | SetCameraUpDirectionMessage\n | SetCameraLookAtMessage\n | SetCameraFovMessage\n | SetOrientationMessage\n | SetPositionMessage\n | TransformControlsUpdateMessage\n | BackgroundImageMessage\n | ImageMessage\n | RemoveSceneNodeMessage\n | SetSceneNodeVisibilityMessage\n | SetSceneNodeClickableMessage\n | SceneNodeClickMessage\n | ResetSceneMessage\n | ResetGuiMessage\n | GuiAddFolderMessage\n | GuiAddMarkdownMessage\n | GuiAddProgressBarMessage\n | GuiAddPlotlyMessage\n | GuiAddTabGroupMessage\n | _GuiAddInputBase\n | GuiAddButtonMessage\n | GuiAddUploadButtonMessage\n | GuiAddSliderMessage\n | GuiAddMultiSliderMessage\n | GuiAddNumberMessage\n | GuiAddRgbMessage\n | GuiAddRgbaMessage\n | GuiAddCheckboxMessage\n | GuiAddVector2Message\n | GuiAddVector3Message\n | GuiAddTextMessage\n | GuiAddDropdownMessage\n | GuiAddButtonGroupMessage\n | GuiModalMessage\n | GuiCloseModalMessage\n | GuiRemoveMessage\n | GuiUpdateMessage\n | ThemeConfigurationMessage\n | CatmullRomSplineMessage\n | CubicBezierSplineMessage\n | GaussianSplatsMessage\n | GetRenderRequestMessage\n | GetRenderResponseMessage\n | FileTransferStart\n | FileTransferPart\n | FileTransferPartAck\n | ShareUrlRequest\n | ShareUrlUpdated\n | ShareUrlDisconnect\n | SetGuiPanelLabelMessage;\nexport type GuiAddComponentMessage =\n | GuiAddFolderMessage\n | GuiAddMarkdownMessage\n | GuiAddProgressBarMessage\n | GuiAddPlotlyMessage\n | GuiAddTabGroupMessage\n | GuiAddButtonMessage\n | GuiAddUploadButtonMessage\n | GuiAddSliderMessage\n | GuiAddMultiSliderMessage\n | GuiAddNumberMessage\n | GuiAddRgbMessage\n | GuiAddRgbaMessage\n | GuiAddCheckboxMessage\n | GuiAddVector2Message\n | GuiAddVector3Message\n | GuiAddTextMessage\n | GuiAddDropdownMessage\n | GuiAddButtonGroupMessage;\n" }, { "path": "viser/src/viser/client/src/WebsocketServerWorker.ts", "content": "import { encode, decode } from \"@msgpack/msgpack\";\nimport { Message } from \"./WebsocketMessages\";\nimport AwaitLock from \"await-lock\";\n\nexport type WsWorkerIncoming =\n | { type: \"send\"; message: Message }\n | { type: \"set_server\"; server: string }\n | { type: \"close\" };\n\nexport type WsWorkerOutgoing =\n | { type: \"connected\" }\n | { type: \"closed\" }\n | { type: \"message_batch\"; messages: Message[] };\n\n// Helper function to collect all ArrayBuffer objects. This is used for postMessage() move semantics.\nfunction collectArrayBuffers(obj: any, buffers: Set) {\n if (obj instanceof ArrayBuffer) {\n buffers.add(obj);\n } else if (obj instanceof Uint8Array) {\n buffers.add(obj.buffer);\n } else if (obj && typeof obj === \"object\") {\n for (const key in obj) {\n if (Object.prototype.hasOwnProperty.call(obj, key)) {\n collectArrayBuffers(obj[key], buffers);\n }\n }\n }\n return buffers;\n}\n{\n let server: string | null = null;\n let ws: WebSocket | null = null;\n const orderLock = new AwaitLock();\n\n const postOutgoing = (\n data: WsWorkerOutgoing,\n transferable?: Transferable[],\n ) => {\n // @ts-ignore\n self.postMessage(data, transferable);\n };\n\n const tryConnect = () => {\n if (ws !== null) ws.close();\n ws = new WebSocket(server!);\n\n // Timeout is necessary when we're connecting to an SSH/tunneled port.\n const retryTimeout = setTimeout(() => {\n ws?.close();\n }, 5000);\n\n ws.onopen = () => {\n postOutgoing({ type: \"connected\" });\n clearTimeout(retryTimeout);\n console.log(`Connected! ${server}`);\n };\n\n ws.onclose = (event) => {\n postOutgoing({ type: \"closed\" });\n console.log(`Disconnected! ${server} code=${event.code}`);\n clearTimeout(retryTimeout);\n\n // Try to reconnect.\n if (server !== null) setTimeout(tryConnect, 1000);\n };\n\n ws.onmessage = async (event) => {\n // Reduce websocket backpressure.\n const messagePromise = new Promise((resolve) => {\n (event.data.arrayBuffer() as Promise).then((buffer) => {\n resolve(decode(new Uint8Array(buffer)) as Message[]);\n });\n });\n\n // Try our best to handle messages in order. If this takes more than 1 second, we give up. :)\n await orderLock.acquireAsync({ timeout: 1000 }).catch(() => {\n console.log(\"Order lock timed out.\");\n orderLock.release();\n });\n try {\n const messages = await messagePromise;\n const arrayBuffers = collectArrayBuffers(messages, new Set());\n postOutgoing(\n { type: \"message_batch\", messages: messages },\n Array.from(arrayBuffers),\n );\n } finally {\n orderLock.acquired && orderLock.release();\n }\n };\n };\n\n self.onmessage = (e) => {\n const data: WsWorkerIncoming = e.data;\n\n if (data.type === \"send\") {\n ws!.send(encode(data.message));\n } else if (data.type === \"set_server\") {\n server = data.server;\n tryConnect();\n } else if (data.type == \"close\") {\n server = null;\n ws !== null && ws.close();\n self.close();\n } else {\n console.log(\n `WebSocket worker: got ${data}, not sure what to do with it!`,\n );\n }\n };\n}\n" }, { "path": "viser/src/viser/client/src/WorldTransformUtils.ts", "content": "import { ViewerContextContents } from \"./App\";\nimport * as THREE from \"three\";\n\n/** Helper for computing the transformation between the three.js world and the\n * Python-exposed world frames. This is useful for things like switching\n * between +Y and +Z up directions for the world frame. */\nexport function computeT_threeworld_world(viewer: ViewerContextContents) {\n const wxyz = viewer.nodeAttributesFromName.current[\"\"]!.wxyz!;\n const position = viewer.nodeAttributesFromName.current[\"\"]!.position ?? [\n 0, 0, 0,\n ];\n return new THREE.Matrix4()\n .makeRotationFromQuaternion(\n new THREE.Quaternion(wxyz[1], wxyz[2], wxyz[3], wxyz[0]),\n )\n .setPosition(position[0], position[1], position[2]);\n}\n\n/** Helper for converting a ray from the three.js world frame to the Python\n * world frame. Applies the transformation from computeT_threeworld_world.\n */\nexport function rayToViserCoords(\n viewer: ViewerContextContents,\n ray: THREE.Ray,\n): THREE.Ray {\n const T_world_threeworld = computeT_threeworld_world(viewer).invert();\n\n const origin = ray.origin.clone().applyMatrix4(T_world_threeworld);\n\n // Compute just the rotation term without new memory allocation; this\n // will mutate T_world_threeworld!\n const R_world_threeworld = T_world_threeworld.setPosition(0.0, 0.0, 0);\n const direction = ray.direction.clone().applyMatrix4(R_world_threeworld);\n\n return new THREE.Ray(origin, direction);\n}\n" }, { "path": "viser/src/viser/client/src/components/Button.tsx", "content": "import { GuiAddButtonMessage } from \"../WebsocketMessages\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { Box } from \"@mantine/core\";\n\nimport { Button } from \"@mantine/core\";\nimport React from \"react\";\nimport { htmlIconWrapper } from \"./ComponentStyles.css\";\n\nexport default function ButtonComponent({\n id,\n visible,\n disabled,\n label,\n ...otherProps\n}: GuiAddButtonMessage) {\n const { messageSender } = React.useContext(GuiComponentContext)!;\n const { color, icon_html } = otherProps;\n if (!(visible ?? true)) return <>;\n\n return (\n \n \n messageSender({\n type: \"GuiUpdateMessage\",\n id: id,\n updates: { value: true },\n })\n }\n style={{\n height: \"2.125em\",\n }}\n disabled={disabled ?? false}\n size=\"sm\"\n leftSection={\n icon_html === null ? undefined : (\n \n )\n }\n >\n {label}\n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/ButtonGroup.tsx", "content": "import * as React from \"react\";\nimport { Button, Flex } from \"@mantine/core\";\nimport { ViserInputComponent } from \"./common\";\nimport { GuiAddButtonGroupMessage } from \"../WebsocketMessages\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\n\nexport default function ButtonGroupComponent({\n id,\n hint,\n label,\n visible,\n disabled,\n options,\n}: GuiAddButtonGroupMessage) {\n const { messageSender } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n \n {options.map((option, index) => (\n \n messageSender({\n type: \"GuiUpdateMessage\",\n id: id,\n updates: { value: option },\n })\n }\n style={{ flexGrow: 1, width: 0 }}\n disabled={disabled}\n size=\"compact-xs\"\n variant=\"outline\"\n >\n {option}\n \n ))}\n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/Checkbox.tsx", "content": "import * as React from \"react\";\nimport { ViserInputComponent } from \"./common\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { GuiAddCheckboxMessage } from \"../WebsocketMessages\";\nimport { Box, Checkbox, Tooltip } from \"@mantine/core\";\n\nexport default function CheckboxComponent({\n id,\n disabled,\n visible,\n hint,\n label,\n value,\n}: GuiAddCheckboxMessage) {\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n let input = (\n {\n setValue(id, value.target.checked);\n }}\n disabled={disabled}\n />\n );\n if (hint !== null && hint !== undefined) {\n // For checkboxes, we want to make sure that the wrapper\n // doesn't expand to the full width of the parent. This will\n // de-center the tooltip.\n input = (\n \n {input}\n \n );\n }\n return {input};\n}\n" }, { "path": "viser/src/viser/client/src/components/ComponentStyles.css.ts", "content": "import { globalStyle, style } from \"@vanilla-extract/css\";\n\nexport const htmlIconWrapper = style({\n height: \"1em\",\n width: \"1em\",\n position: \"relative\",\n});\n\nglobalStyle(`${htmlIconWrapper} svg`, {\n height: \"auto\",\n width: \"1em\",\n position: \"absolute\",\n top: \"50%\",\n transform: \"translateY(-50%)\",\n});\n\n// Class for sliders with default min/max marks. We use this for aestheticn\n// its; global styles are used to shift the min/max mark labels to stay closer\n// within the bounds of the slider.\nexport const sliderDefaultMarks = style({});\n\nglobalStyle(\n `${sliderDefaultMarks} .mantine-Slider-markWrapper:first-of-type div:nth-of-type(2)`,\n {\n transform: \"translate(-0.1rem, 0.03rem) !important\",\n },\n);\n\nglobalStyle(\n `${sliderDefaultMarks} .mantine-Slider-markWrapper:last-of-type div:nth-of-type(2)`,\n {\n transform: \"translate(-85%, 0.03rem) !important\",\n },\n);\n" }, { "path": "viser/src/viser/client/src/components/Dropdown.tsx", "content": "import * as React from \"react\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { ViserInputComponent } from \"./common\";\nimport { GuiAddDropdownMessage } from \"../WebsocketMessages\";\nimport { Select } from \"@mantine/core\";\n\nexport default function DropdownComponent({\n id,\n hint,\n label,\n value,\n disabled,\n visible,\n options,\n}: GuiAddDropdownMessage) {\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n value !== null && setValue(id, value)}\n disabled={disabled}\n searchable\n maxDropdownHeight={400}\n size=\"xs\"\n rightSectionWidth=\"1.2em\"\n styles={{\n input: {\n padding: \"0.5em\",\n letterSpacing: \"-0.5px\",\n minHeight: \"1.625rem\",\n height: \"1.625rem\",\n },\n }}\n // zIndex of dropdown should be >modal zIndex.\n // On edge cases: it seems like existing dropdowns are always closed when a new modal is opened.\n comboboxProps={{ zIndex: 1000 }}\n />\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/Folder.css.ts", "content": "import { style } from \"@vanilla-extract/css\";\nimport { vars } from \"../AppTheme\";\n\nexport const folderWrapper = style({\n position: \"relative\",\n marginTop: vars.spacing.xs,\n marginLeft: vars.spacing.xs,\n marginRight: vars.spacing.xs,\n marginBottom: vars.spacing.md,\n \":last-child\": {\n marginBottom: vars.spacing.xs,\n },\n});\n\nexport const folderLabel = style({\n fontSize: \"0.875em\",\n position: \"absolute\",\n padding: \"0 0.375em 0 0.375em\",\n top: 0,\n left: \"0.375em\",\n transform: \"translateY(-50%)\",\n userSelect: \"none\",\n fontWeight: 500,\n});\n\nexport const folderToggleIcon = style({\n width: \"0.9em\",\n height: \"0.9em\",\n strokeWidth: 3,\n top: \"0.1em\",\n position: \"relative\",\n marginLeft: \"0.25em\",\n marginRight: \"-0.1em\",\n opacity: 0.5,\n});\n" }, { "path": "viser/src/viser/client/src/components/Folder.tsx", "content": "import * as React from \"react\";\nimport { useDisclosure } from \"@mantine/hooks\";\nimport { GuiAddFolderMessage } from \"../WebsocketMessages\";\nimport { IconChevronDown, IconChevronUp } from \"@tabler/icons-react\";\nimport { Box, Collapse, Paper } from \"@mantine/core\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { ViewerContext } from \"../App\";\nimport { folderLabel, folderToggleIcon, folderWrapper } from \"./Folder.css\";\n\nexport default function FolderComponent({\n id,\n label,\n visible,\n expand_by_default,\n}: GuiAddFolderMessage) {\n const viewer = React.useContext(ViewerContext)!;\n const [opened, { toggle }] = useDisclosure(expand_by_default);\n const guiIdSet = viewer.useGui((state) => state.guiIdSetFromContainerId[id]);\n const guiContext = React.useContext(GuiComponentContext)!;\n const isEmpty = guiIdSet === undefined || Object.keys(guiIdSet).length === 0;\n\n const ToggleIcon = opened ? IconChevronUp : IconChevronDown;\n if (!visible) return <>;\n return (\n \n \n {label}\n \n \n \n \n \n \n \n \n \n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/Markdown.tsx", "content": "import { Box, Text } from \"@mantine/core\";\nimport Markdown from \"../Markdown\";\nimport { ErrorBoundary } from \"react-error-boundary\";\nimport { GuiAddMarkdownMessage } from \"../WebsocketMessages\";\n\nexport default function MarkdownComponent({\n visible,\n markdown,\n}: GuiAddMarkdownMessage) {\n if (!visible) return <>;\n return (\n \n Markdown Failed to Render}\n >\n {markdown}\n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/MultiSlider.tsx", "content": "import React from \"react\";\nimport { GuiAddMultiSliderMessage } from \"../WebsocketMessages\";\nimport { Box, useMantineColorScheme } from \"@mantine/core\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { ViserInputComponent } from \"./common\";\nimport { MultiSlider } from \"./MultiSliderPrimitive\";\nimport { sliderDefaultMarks } from \"./ComponentStyles.css\";\n\nexport default function MultiSliderComponent({\n id,\n label,\n hint,\n visible,\n disabled,\n value,\n ...otherProps\n}: GuiAddMultiSliderMessage) {\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n const updateValue = (value: number[]) => setValue(id, value);\n const { min, max, precision, step, marks, fixed_endpoints, min_range } =\n otherProps;\n const colorScheme = useMantineColorScheme().colorScheme;\n const input = (\n \n ({\n thumb: {\n height: \"0.75rem\",\n width: \"0.5rem\",\n },\n trackContainer: {\n zIndex: 3,\n position: \"relative\",\n },\n markLabel: {\n transform: \"translate(-50%, 0.03rem)\",\n fontSize: \"0.6rem\",\n textAlign: \"center\",\n },\n mark: {\n transform: \"scale(1.85)\",\n },\n markFilled: {\n background: disabled\n ? colorScheme === \"dark\"\n ? theme.colors.dark[3]\n : theme.colors.gray[4]\n : theme.primaryColor,\n },\n })}\n pt=\"0.2em\"\n pb=\"0.4em\"\n min={min}\n max={max}\n step={step ?? undefined}\n fixedEndpoints={fixed_endpoints}\n precision={precision}\n minRange={min_range ?? undefined}\n marks={\n marks === null\n ? [\n {\n value: min,\n label: `${parseInt(min.toFixed(6))}`,\n },\n {\n value: max,\n label: `${parseInt(max.toFixed(6))}`,\n },\n ]\n : marks\n }\n value={value}\n onChange={updateValue}\n />\n \n );\n\n return (\n {input}\n );\n}\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/LICENSE", "content": "This component is modified from Mantine v7.\n\nhttps://github.com/mantinedev/mantine/tree/e3e3bb834de1f2f75a27dbc757dc0a2fc6a6cba8/packages/%40mantine/core/src/components/Slider\n\n--\n\nMIT License\n\nCopyright (c) 2021 Vitaly Rtishchev\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/Marks/Marks.tsx", "content": "import React from \"react\";\nimport { Box } from \"@mantine/core\";\nimport { useSliderContext } from \"../Slider.context\";\nimport { getPosition } from \"../utils/get-position/get-position\";\n\nexport interface MarksProps {\n marks: { value: number; label?: React.ReactNode }[] | undefined;\n min: number;\n max: number;\n value: number;\n offset: number | undefined;\n disabled: boolean | undefined;\n inverted: boolean | undefined;\n}\n\nexport function Marks({\n marks,\n min,\n max,\n disabled,\n value, // eslint-disable-line\n offset, // eslint-disable-line\n inverted, // eslint-disable-line\n}: MarksProps) {\n const { getStyles } = useSliderContext();\n\n if (!marks) {\n return null;\n }\n\n const items = marks.map((mark, index) => (\n \n \n {mark.label &&
{mark.label}
}\n \n ));\n\n return
{items}
;\n}\n\nMarks.displayName = \"@mantine/core/SliderMarks\";\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/MultiSlider/MultiSlider.tsx", "content": "import React, { useCallback, useRef, useState } from \"react\";\nimport { clamp, useMergedRef, useMove, useUncontrolled } from \"@mantine/hooks\";\nimport {\n BoxProps,\n createVarsResolver,\n ElementProps,\n factory,\n Factory,\n getRadius,\n getSize,\n getThemeColor,\n MantineColor,\n MantineRadius,\n MantineSize,\n rem,\n StylesApiProps,\n useDirection,\n useProps,\n useStyles,\n TransitionOverride,\n} from \"@mantine/core\";\nimport {\n SliderCssVariables,\n SliderProvider,\n SliderStylesNames,\n} from \"../Slider.context\";\nimport { SliderRoot } from \"../SliderRoot/SliderRoot\";\nimport { Thumb } from \"../Thumb/Thumb\";\nimport { Track } from \"../Track/Track\";\nimport { getChangeValue } from \"../utils/get-change-value/get-change-value\";\nimport { getPosition } from \"../utils/get-position/get-position\";\nimport { getPrecision } from \"../utils/get-precision/get-precision\";\nimport classes from \"../Slider.module.css\";\n\nexport interface SliderProps\n extends BoxProps,\n StylesApiProps,\n ElementProps<\"div\", \"onChange\" | \"defaultValue\"> {\n /** Key of `theme.colors` or any valid CSS color, controls color of track and thumb, `theme.primaryColor` by default */\n color?: MantineColor;\n\n /** Key of `theme.radius` or any valid CSS value to set `border-radius`, numbers are converted to rem, `'xl'` by default */\n radius?: MantineRadius;\n\n /** Controls size of the track, `'md'` by default */\n size?: MantineSize | (string & NonNullable) | number;\n\n /** Minimal possible value, `0` by default */\n min?: number;\n\n /** Maximum possible value, `100` by default */\n max?: number;\n\n /** Number by which value will be incremented/decremented with thumb drag and arrows, `1` by default */\n step?: number;\n\n /** Number of significant digits after the decimal point */\n precision?: number;\n\n fixedEndpoints: boolean;\n minRange?: number;\n\n /** Controlled component value */\n value?: number[];\n\n /** Uncontrolled component default value */\n defaultValue?: number[];\n\n /** Called when value changes */\n onChange?: (value: number[]) => void;\n\n /** Called when user stops dragging slider or changes value with arrows */\n onChangeEnd?: (value: number[]) => void;\n\n /** Hidden input name, use with uncontrolled component */\n name?: string;\n\n /** Marks displayed on the track */\n marks?: { value: number; label?: React.ReactNode }[];\n\n /** Function to generate label or any react node to render instead, set to null to disable label */\n label?: React.ReactNode | ((value: number) => React.ReactNode);\n\n /** Props passed down to the `Transition` component, `{ transition: 'fade', duration: 0 }` by default */\n labelTransitionProps?: TransitionOverride;\n\n /** Determines whether the label should be visible when the slider is not being dragged or hovered, `false` by default */\n labelAlwaysOn?: boolean;\n\n /** Thumb `aria-label` */\n thumbLabel?: string;\n\n /** Determines whether the label should be displayed when the slider is hovered, `true` by default */\n showLabelOnHover?: boolean;\n\n /** Content rendered inside thumb */\n thumbChildren?: React.ReactNode;\n\n /** Disables slider */\n disabled?: boolean;\n\n /** Thumb `width` and `height`, by default value is computed based on `size` prop */\n thumbSize?: number | string;\n\n /** A transformation function to change the scale of the slider */\n scale?: (value: number) => number;\n\n /** Determines whether track value representation should be inverted, `false` by default */\n inverted?: boolean;\n\n /** Props passed down to the hidden input */\n hiddenInputProps?: React.ComponentPropsWithoutRef<\"input\">;\n}\n\nexport type SliderFactory = Factory<{\n props: SliderProps;\n ref: HTMLDivElement;\n stylesNames: SliderStylesNames;\n vars: SliderCssVariables;\n}>;\n\nconst defaultProps: Partial = {\n radius: \"xl\",\n min: 0,\n max: 100,\n step: 1,\n fixedEndpoints: true,\n marks: [],\n label: (f) => f,\n labelTransitionProps: { transition: \"fade\", duration: 0 },\n labelAlwaysOn: false,\n thumbLabel: \"\",\n showLabelOnHover: true,\n disabled: false,\n scale: (v) => v,\n};\n\nconst varsResolver = createVarsResolver(\n (theme, { size, color, thumbSize, radius }) => ({\n root: {\n \"--slider-size\": getSize(size, \"slider-size\"),\n \"--slider-color\": color ? getThemeColor(color, theme) : undefined,\n \"--slider-radius\": radius === undefined ? undefined : getRadius(radius),\n \"--slider-thumb-size\":\n thumbSize !== undefined\n ? rem(thumbSize)\n : \"calc(var(--slider-size) * 2)\",\n },\n }),\n);\n\nexport const MultiSlider = factory((_props, ref) => {\n const props = useProps(\"MultiSlider\", defaultProps, _props);\n const {\n classNames,\n styles,\n value,\n onChange,\n onChangeEnd,\n size,\n min,\n max,\n step,\n fixedEndpoints,\n minRange,\n precision: _precision,\n defaultValue,\n name,\n marks,\n label,\n labelTransitionProps,\n labelAlwaysOn,\n thumbLabel,\n showLabelOnHover,\n thumbChildren,\n disabled,\n unstyled,\n scale,\n inverted,\n className,\n style,\n vars,\n // hiddenInputProps,\n ...others\n } = props;\n\n const getStyles = useStyles({\n name: \"MultiSlider\",\n props,\n classes,\n classNames,\n className,\n styles,\n style,\n vars,\n varsResolver,\n unstyled,\n });\n\n const { dir } = useDirection();\n const [hovered, setHovered] = useState(false);\n const [_value, setValue] = useUncontrolled({\n value: value === undefined ? value : value.map((x) => clamp(x, min!, max!)),\n defaultValue:\n defaultValue === undefined\n ? defaultValue\n : defaultValue.map((x) => clamp(x, min!, max!)),\n finalValue: [clamp(0, min!, max!)],\n onChange,\n });\n\n const valueRef = useRef(_value);\n const root = useRef();\n const thumbs = useRef<(HTMLDivElement | null)[]>([]);\n const thumbIndex = useRef(-1);\n const positions = _value.map((x) =>\n getPosition({ value: x, min: min!, max: max! }),\n );\n const precision = _precision ?? getPrecision(step!);\n\n valueRef.current = _value;\n\n const setRangedValue = (\n val: number,\n thumbIndex: number,\n triggerChangeEnd: boolean,\n ) => {\n const clone: number[] = [...valueRef.current];\n clone[thumbIndex] = val;\n\n const _minRange = minRange || step!;\n if (thumbIndex < clone.length - 1) {\n if (val > clone[thumbIndex + 1] - (_minRange - 0.000000001)) {\n clone[thumbIndex] = Math.max(min!, clone[thumbIndex + 1] - _minRange);\n }\n\n if (val > (max! - (_minRange - 0.000000001) || min!)) {\n clone[thumbIndex] = valueRef.current[thumbIndex];\n }\n }\n\n if (thumbIndex > 0) {\n if (val < clone[thumbIndex - 1] + _minRange) {\n clone[thumbIndex] = Math.min(max!, clone[thumbIndex - 1] + _minRange);\n }\n }\n\n // const fixedEndpoints = false;\n if (\n fixedEndpoints &&\n (thumbIndex === 0 || thumbIndex == clone.length - 1)\n ) {\n clone[thumbIndex] = valueRef.current[thumbIndex];\n }\n\n setValue(clone);\n valueRef.current = clone;\n\n if (triggerChangeEnd) {\n onChangeEnd?.(valueRef.current);\n }\n };\n\n const handleChange = useCallback(\n ({ x }: { x: number }) => {\n if (!disabled) {\n const nextValue = getChangeValue({\n value: x,\n min: min!,\n max: max!,\n step: step!,\n precision,\n });\n setRangedValue(nextValue, thumbIndex.current, false);\n }\n },\n [disabled, min, max, step, precision, setValue],\n );\n\n const { ref: container, active } = useMove(\n handleChange,\n { onScrubEnd: () => onChangeEnd?.(valueRef.current) },\n dir,\n );\n\n function getClientPosition(event: any) {\n if (\"TouchEvent\" in window && event instanceof window.TouchEvent) {\n const touch = event.touches[0];\n return touch.clientX;\n }\n\n return event.clientX;\n }\n const handleTrackMouseDownCapture = (\n event: React.MouseEvent | React.TouchEvent,\n ) => {\n container.current!.focus();\n const rect = container.current!.getBoundingClientRect();\n const changePosition = getClientPosition(event.nativeEvent);\n const changeValue = getChangeValue({\n value: changePosition - rect.left,\n max: max!,\n min: min!,\n step: step!,\n containerWidth: rect.width,\n });\n\n const _nearestHandle = _value\n .map((v) => Math.abs(v - changeValue))\n .indexOf(Math.min(..._value.map((v) => Math.abs(v - changeValue))));\n\n thumbIndex.current = _nearestHandle;\n };\n\n const handleTrackKeydownCapture = (\n event: React.KeyboardEvent,\n ) => {\n if (!disabled) {\n const focusedIndex = thumbIndex.current;\n switch (event.key) {\n case \"ArrowUp\": {\n event.preventDefault();\n thumbs.current[focusedIndex]!.focus();\n setRangedValue(\n Math.min(\n Math.max(valueRef.current[focusedIndex] + step!, min!),\n max!,\n ),\n focusedIndex,\n true,\n );\n break;\n }\n case \"ArrowRight\": {\n event.preventDefault();\n thumbs.current[focusedIndex]?.focus();\n setRangedValue(\n Math.min(\n Math.max(\n dir === \"rtl\"\n ? valueRef.current[focusedIndex] - step!\n : valueRef.current[focusedIndex] + step!,\n min!,\n ),\n max!,\n ),\n focusedIndex,\n true,\n );\n break;\n }\n\n case \"ArrowDown\": {\n event.preventDefault();\n thumbs.current[focusedIndex]?.focus();\n setRangedValue(\n Math.min(\n Math.max(valueRef.current[focusedIndex] - step!, min!),\n max!,\n ),\n focusedIndex,\n true,\n );\n break;\n }\n case \"ArrowLeft\": {\n event.preventDefault();\n thumbs.current[focusedIndex]?.focus();\n setRangedValue(\n Math.min(\n Math.max(\n dir === \"rtl\"\n ? valueRef.current[focusedIndex] + step!\n : valueRef.current[focusedIndex] - step!,\n min!,\n ),\n max!,\n ),\n focusedIndex,\n true,\n );\n break;\n }\n\n case \"Home\": {\n event.preventDefault();\n thumbs.current[focusedIndex]?.focus();\n setRangedValue(min!, focusedIndex, true);\n break;\n }\n\n case \"End\": {\n event.preventDefault();\n thumbs.current[focusedIndex]?.focus();\n setRangedValue(max!, focusedIndex, true);\n break;\n }\n\n default: {\n break;\n }\n }\n }\n };\n\n return (\n \n root.current?.focus()}\n size={size!}\n disabled={disabled}\n >\n setHovered(true) : undefined,\n onMouseLeave: showLabelOnHover\n ? () => setHovered(false)\n : undefined,\n onTouchStartCapture: handleTrackMouseDownCapture,\n onTouchEndCapture: () => {\n thumbIndex.current = -1;\n },\n onMouseDownCapture: handleTrackMouseDownCapture,\n onMouseUpCapture: () => {\n thumbIndex.current = -1;\n },\n }}\n >\n {_value.map((value, index) => (\n {\n thumbs.current[index] = node;\n }}\n labelTransitionProps={labelTransitionProps}\n labelAlwaysOn={labelAlwaysOn}\n thumbLabel={thumbLabel}\n showLabelOnHover={showLabelOnHover}\n isHovered={hovered}\n disabled={disabled}\n >\n {thumbChildren}\n \n ))}\n \n {_value.map((value, index) => (\n \n ))}\n \n \n );\n});\n\nMultiSlider.classes = classes;\nMultiSlider.displayName = \"MultiSlider\";\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/Slider.context.ts", "content": "import { createSafeContext, GetStylesApi } from \"@mantine/core\";\n\nexport type SliderStylesNames =\n | \"root\"\n | \"label\"\n | \"thumb\"\n | \"trackContainer\"\n | \"track\"\n | \"bar\"\n | \"markWrapper\"\n | \"mark\"\n | \"markLabel\";\n\nexport type SliderCssVariables = {\n root:\n | \"--slider-size\"\n | \"--slider-color\"\n | \"--slider-thumb-size\"\n | \"--slider-radius\";\n};\n\ninterface SliderContextValue {\n getStyles: GetStylesApi<{\n stylesNames: SliderStylesNames;\n props: any;\n ref: any;\n vars: any;\n variant: any;\n }>;\n}\n\nexport const [SliderProvider, useSliderContext] =\n createSafeContext(\"SliderProvider was not found in tree\");\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/Slider.module.css", "content": ".root {\n --slider-size-xs: rem(4px);\n --slider-size-sm: rem(6px);\n --slider-size-md: rem(8px);\n --slider-size-lg: rem(10px);\n --slider-size-xl: rem(12px);\n\n --slider-size: var(--slider-size-md);\n --slider-radius: rem(1000px);\n --slider-color: var(--mantine-primary-color-filled);\n\n -webkit-tap-highlight-color: transparent;\n outline: none;\n height: calc(var(--slider-size) * 2);\n padding-inline: var(--slider-size);\n display: flex;\n flex-direction: column;\n align-items: center;\n touch-action: none;\n position: relative;\n\n @mixin light {\n --slider-track-bg: var(--mantine-color-gray-2);\n --slider-track-disabled-bg: var(--mantine-color-gray-4);\n }\n\n @mixin dark {\n --slider-track-bg: var(--mantine-color-dark-4);\n --slider-track-disabled-bg: var(--mantine-color-dark-3);\n }\n}\n\n.label {\n position: absolute;\n top: rem(-36px);\n font-size: var(--mantine-font-size-xs);\n color: var(--mantine-color-white);\n padding: calc(var(--mantine-spacing-xs) / 2);\n border-radius: var(--mantine-radius-sm);\n white-space: nowrap;\n pointer-events: none;\n user-select: none;\n touch-action: none;\n\n @mixin where-light {\n background-color: var(--mantine-color-gray-9);\n }\n\n @mixin where-dark {\n background-color: var(--mantine-color-dark-4);\n }\n}\n\n.thumb {\n position: absolute;\n display: flex;\n height: var(--slider-thumb-size);\n width: var(--slider-thumb-size);\n border: rem(4px) solid;\n transform: translate(-50%, -50%);\n color: var(--slider-color);\n top: 50%;\n cursor: pointer;\n border-radius: var(--slider-radius);\n align-items: center;\n justify-content: center;\n transition:\n box-shadow 100ms ease,\n transform 100ms ease;\n z-index: 3;\n user-select: none;\n touch-action: none;\n outline-offset: rem(2px);\n left: var(--slider-thumb-offset);\n\n @mixin where-rtl {\n left: auto;\n right: calc(var(--slider-thumb-offset) - var(--slider-thumb-size));\n }\n\n fieldset:disabled &,\n &:where([data-disabled]) {\n display: none;\n }\n\n &:where([data-dragging]) {\n transform: translate(-50%, -50%) scale(1.05);\n box-shadow: var(--mantine-shadow-sm);\n }\n\n @mixin where-light {\n border-color: var(--slider-color);\n background-color: var(--mantine-color-white);\n }\n\n @mixin where-dark {\n border-color: var(--mantine-color-white);\n background-color: var(--slider-color);\n }\n}\n\n.trackContainer {\n display: flex;\n align-items: center;\n width: 100%;\n height: calc(var(--slider-size) * 2);\n cursor: pointer;\n\n fieldset:disabled &,\n &:where([data-disabled]) {\n cursor: not-allowed;\n }\n}\n\n.track {\n position: relative;\n width: 100%;\n height: var(--slider-size);\n\n &:where([data-inverted]:not([data-disabled])) {\n --track-bg: var(--slider-color);\n }\n\n fieldset:disabled &:where([data-inverted]),\n &:where([data-inverted][data-disabled]) {\n --track-bg: var(--slider-track-disabled-bg);\n }\n\n &::before {\n content: \"\";\n position: absolute;\n top: 0;\n bottom: 0;\n border-radius: var(--slider-radius);\n inset-inline: calc(var(--slider-size) * -1);\n background-color: var(--track-bg, var(--slider-track-bg));\n z-index: 0;\n }\n}\n\n.bar {\n position: absolute;\n z-index: 1;\n top: 0;\n bottom: 0;\n background-color: var(--slider-color);\n border-radius: var(--slider-radius);\n width: var(--slider-bar-width);\n inset-inline-start: var(--slider-bar-offset);\n\n &:where([data-inverted]) {\n background-color: var(--slider-track-bg);\n }\n\n fieldset:disabled &:where(:not([data-inverted])),\n &:where([data-disabled]:not([data-inverted])) {\n @mixin where-light {\n background-color: var(--mantine-color-gray-4);\n }\n\n @mixin where-dark {\n background-color: var(--mantine-color-dark-3);\n }\n }\n}\n\n.markWrapper {\n position: absolute;\n inset-inline-start: calc(var(--mark-offset) - var(--slider-size) / 2);\n top: 0;\n z-index: 2;\n height: 0;\n pointer-events: none;\n}\n\n.mark {\n border: rem(2px) solid;\n height: var(--slider-size);\n width: var(--slider-size);\n border-radius: rem(1000px);\n transform: translateX((calc(var(--slider-size) / -2)));\n background-color: var(--mantine-color-white);\n pointer-events: none;\n\n @mixin where-light {\n border-color: var(--mantine-color-gray-2);\n }\n\n @mixin where-dark {\n border-color: var(--mantine-color-dark-4);\n }\n\n &:where([data-filled]) {\n border-color: var(--slider-color);\n\n &:where([data-disabled]) {\n @mixin where-light {\n border-color: var(--mantine-color-gray-4);\n }\n\n @mixin where-dark {\n border-color: var(--mantine-color-dark-3);\n }\n }\n }\n}\n\n.markLabel {\n transform: translate(\n calc(-50% + var(--slider-size) / 2),\n calc(var(--mantine-spacing-xs) / 2)\n );\n font-size: var(--mantine-font-size-sm);\n white-space: nowrap;\n cursor: pointer;\n user-select: none;\n\n @mixin where-light {\n color: var(--mantine-color-gray-6);\n }\n\n @mixin where-dark {\n color: var(--mantine-color-dark-2);\n }\n}\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/SliderRoot/SliderRoot.tsx", "content": "import React, { forwardRef } from \"react\";\nimport {\n Box,\n BoxProps,\n ElementProps,\n MantineColor,\n MantineRadius,\n MantineSize,\n} from \"@mantine/core\";\nimport { useSliderContext } from \"../Slider.context\";\n\nexport interface SliderRootProps extends BoxProps, ElementProps<\"div\"> {\n size: MantineSize | (string & NonNullable) | number;\n children: React.ReactNode;\n color: MantineColor | undefined;\n disabled: boolean | undefined;\n variant?: string;\n thumbSize: string | number | undefined;\n radius: MantineRadius | undefined;\n}\n\nexport const SliderRoot = forwardRef(\n ({ size, variant, ...others }: SliderRootProps, ref) => {\n const { getStyles } = useSliderContext();\n\n return (\n \n );\n },\n);\n\nSliderRoot.displayName = \"@mantine/core/SliderRoot\";\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/Thumb/Thumb.tsx", "content": "import React, { forwardRef, useState } from \"react\";\nimport { Box } from \"@mantine/core\";\nimport { Transition, TransitionOverride } from \"@mantine/core\";\nimport { useSliderContext } from \"../Slider.context\";\n\nexport interface ThumbProps {\n max: number;\n min: number;\n value: number;\n position: number;\n dragging: boolean;\n draggingThisThumb: boolean;\n label: React.ReactNode;\n onKeyDownCapture?: (event: React.KeyboardEvent) => void;\n onMouseDown?: (\n event: React.MouseEvent | React.TouchEvent,\n ) => void;\n labelTransitionProps: TransitionOverride | undefined;\n labelAlwaysOn: boolean | undefined;\n thumbLabel: string | undefined;\n onFocus?: () => void;\n onBlur?: () => void;\n showLabelOnHover: boolean | undefined;\n isHovered?: boolean;\n children?: React.ReactNode;\n disabled: boolean | undefined;\n className?: string;\n style?: React.CSSProperties;\n}\n\nexport const Thumb = forwardRef(\n (\n {\n max,\n min,\n value,\n position,\n label,\n dragging,\n draggingThisThumb,\n onMouseDown,\n onKeyDownCapture,\n labelTransitionProps,\n labelAlwaysOn,\n thumbLabel,\n onFocus,\n onBlur,\n showLabelOnHover,\n isHovered,\n children = null,\n disabled,\n }: ThumbProps,\n ref,\n ) => {\n const { getStyles } = useSliderContext();\n\n const [focused, setFocused] = useState(false);\n\n const isVisible =\n labelAlwaysOn || dragging || focused || (showLabelOnHover && isHovered);\n\n return (\n \n tabIndex={0}\n role=\"slider\"\n aria-label={thumbLabel}\n aria-valuemax={max}\n aria-valuemin={min}\n aria-valuenow={value}\n ref={ref}\n __vars={{ \"--slider-thumb-offset\": `${position}%` }}\n {...getStyles(\"thumb\", {\n focusable: true,\n style: {\n /* Put active thumb + its label in front of others. */\n ...(draggingThisThumb ? { zIndex: 1000 } : {}),\n },\n })}\n mod={{ dragging, disabled }}\n onFocus={() => {\n setFocused(true);\n typeof onFocus === \"function\" && onFocus();\n }}\n onBlur={() => {\n setFocused(false);\n typeof onBlur === \"function\" && onBlur();\n }}\n onTouchStart={onMouseDown}\n onMouseDown={onMouseDown}\n onKeyDownCapture={onKeyDownCapture}\n onClick={(event) => event.stopPropagation()}\n >\n {children}\n \n {(transitionStyles) => (\n \n {label}\n \n )}\n \n \n );\n },\n);\n\nThumb.displayName = \"@mantine/core/SliderThumb\";\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/Track/Track.tsx", "content": "import React from \"react\";\nimport { Box } from \"@mantine/core\";\nimport { Marks } from \"../Marks/Marks\";\nimport { useSliderContext } from \"../Slider.context\";\n\nexport interface TrackProps {\n filled: number;\n offset?: number;\n marksOffset?: number;\n marks: { value: number; label?: React.ReactNode }[] | undefined;\n min: number;\n max: number;\n value: number;\n children: React.ReactNode;\n disabled: boolean | undefined;\n inverted: boolean | undefined;\n containerProps?: React.PropsWithRef>;\n}\n\nexport function Track({\n children,\n disabled,\n marksOffset,\n inverted,\n containerProps,\n ...others\n}: TrackProps) {\n const { getStyles } = useSliderContext();\n\n return (\n <>\n \n \n {children}\n\n \n \n \n \n );\n}\n\nTrack.displayName = \"@mantine/core/SliderTrack\";\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/index.ts", "content": "export { MultiSlider } from \"./MultiSlider/MultiSlider\";\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/utils/get-change-value/get-change-value.ts", "content": "interface GetChangeValue {\n value: number;\n containerWidth?: number;\n min: number;\n max: number;\n step: number;\n precision?: number;\n}\n\nexport function getChangeValue({\n value,\n containerWidth,\n min,\n max,\n step,\n precision,\n}: GetChangeValue) {\n const left = !containerWidth\n ? value\n : Math.min(Math.max(value, 0), containerWidth) / containerWidth;\n const dx = left * (max - min);\n const nextValue = (dx !== 0 ? Math.round(dx / step) * step : 0) + min;\n\n const nextValueWithinStep = Math.max(nextValue, min);\n\n if (precision !== undefined) {\n return Number(nextValueWithinStep.toFixed(precision));\n }\n\n return nextValueWithinStep;\n}\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/utils/get-client-position/get-client-position.ts", "content": "export function getClientPosition(event: any) {\n if (\"TouchEvent\" in window && event instanceof window.TouchEvent) {\n const touch = event.touches[0];\n return touch.clientX;\n }\n\n return event.clientX;\n}\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/utils/get-floating-value/get-gloating-value.ts", "content": "export function getFloatingValue(value: number, precision: number) {\n return parseFloat(value.toFixed(precision));\n}\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/utils/get-position/get-position.ts", "content": "interface GetPosition {\n value: number;\n min: number;\n max: number;\n}\n\nexport function getPosition({ value, min, max }: GetPosition) {\n const position = ((value - min) / (max - min)) * 100;\n return Math.min(Math.max(position, 0), 100);\n}\n" }, { "path": "viser/src/viser/client/src/components/MultiSliderPrimitive/utils/get-precision/get-precision.ts", "content": "export function getPrecision(step: number) {\n if (!step) return 0;\n const split = step.toString().split(\".\");\n return split.length > 1 ? split[1].length : 0;\n}\n" }, { "path": "viser/src/viser/client/src/components/NumberInput.tsx", "content": "import * as React from \"react\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { GuiAddNumberMessage } from \"../WebsocketMessages\";\nimport { ViserInputComponent } from \"./common\";\nimport { NumberInput } from \"@mantine/core\";\n\nexport default function NumberInputComponent({\n visible,\n id,\n label,\n hint,\n value,\n disabled,\n ...otherProps\n}: GuiAddNumberMessage) {\n const { setValue } = React.useContext(GuiComponentContext)!;\n const { precision, min, max, step } = otherProps;\n if (!visible) return <>;\n return (\n \n {\n // Ignore empty values.\n newValue !== \"\" && setValue(id, newValue);\n }}\n styles={{\n input: {\n minHeight: \"1.625rem\",\n height: \"1.625rem\",\n },\n controls: {\n height: \"1.625em\",\n width: \"0.825em\",\n },\n }}\n disabled={disabled}\n stepHoldDelay={500}\n stepHoldInterval={(t) => Math.max(1000 / t ** 2, 25)}\n />\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/PlotlyComponent.tsx", "content": "import React from \"react\";\nimport { GuiAddPlotlyMessage } from \"../WebsocketMessages\";\nimport { useDisclosure } from \"@mantine/hooks\";\nimport { Modal, Box, Paper, Tooltip } from \"@mantine/core\";\nimport { useElementSize } from \"@mantine/hooks\";\n\n// When drawing border around the plot, it should be aligned with the folder's.\nimport { folderWrapper } from \"./Folder.css\";\n\nconst PlotWithAspect = React.memo(function PlotWithAspect({\n jsonStr,\n aspectRatio,\n staticPlot,\n}: {\n jsonStr: string;\n aspectRatio: number;\n staticPlot: boolean;\n}) {\n // Catch if the jsonStr is empty; if so, render an empty div.\n if (jsonStr === \"\") return
;\n\n // Parse json string, to construct plotly object.\n // Note that only the JSON string is kept as state, not the json object.\n const plotJson = JSON.parse(jsonStr);\n\n // This keeps the zoom-in state, etc, see https://plotly.com/javascript/uirevision/.\n plotJson.layout.uirevision = \"true\";\n\n // Box size change -> width value change -> plot rerender trigger.\n const { ref, width } = useElementSize();\n plotJson.layout.width = width;\n plotJson.layout.height = width * aspectRatio;\n\n // Make the plot non-interactable, if specified.\n // Ideally, we would use `staticplot`, but this has a known bug with 3D plots:\n // - https://github.com/plotly/plotly.js/issues/457\n // In the meantime, we choose to disable all interactions.\n if (staticPlot) {\n if (plotJson.config === undefined) plotJson.config = {};\n plotJson.config.displayModeBar = false;\n plotJson.layout.dragmode = false;\n plotJson.layout.hovermode = false;\n plotJson.layout.clickmode = \"none\";\n }\n\n // Use React hooks to update the plotly object, when the plot data changes.\n // based on https://github.com/plotly/react-plotly.js/issues/242.\n const plotRef = React.useRef(null);\n React.useEffect(() => {\n // @ts-ignore - Plotly.js is dynamically imported with an eval() call.\n Plotly.react(\n plotRef.current!,\n plotJson.data,\n plotJson.layout,\n plotJson.config,\n );\n }, [plotJson]);\n\n return (\n \n
\n {/* Add a div on top of the plot, to prevent interaction + cursor changes. */}\n {staticPlot ? (\n \n ) : null}\n \n );\n});\n\nexport default function PlotlyComponent({\n visible,\n plotly_json_str,\n aspect,\n}: GuiAddPlotlyMessage) {\n if (!visible) return <>;\n\n // Create a modal with the plot, and a button to open it.\n const [opened, { open, close }] = useDisclosure(false);\n return (\n \n {/* Draw static plot in the controlpanel, which can be clicked. */}\n \n \n \n \n \n\n {/* Modal contents. keepMounted makes state changes (eg zoom) to the plot\n persistent. */}\n \n \n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/ProgressBar.tsx", "content": "import { Box, Progress } from \"@mantine/core\";\nimport { GuiAddProgressBarMessage } from \"../WebsocketMessages\";\n\nexport default function ProgressBarComponent({\n visible,\n color,\n value,\n animated,\n}: GuiAddProgressBarMessage) {\n if (!visible) return <>;\n return (\n \n \n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/Rgb.tsx", "content": "import * as React from \"react\";\nimport { ColorInput } from \"@mantine/core\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { rgbToHex, hexToRgb } from \"./utils\";\nimport { ViserInputComponent } from \"./common\";\nimport { GuiAddRgbMessage } from \"../WebsocketMessages\";\n\nexport default function RgbComponent({\n id,\n label,\n hint,\n value,\n disabled,\n visible,\n}: GuiAddRgbMessage) {\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n setValue(id, hexToRgb(v))}\n format=\"hex\"\n // zIndex of dropdown should be >modal zIndex.\n // On edge cases: it seems like existing dropdowns are always closed when a new modal is opened.\n popoverProps={{ zIndex: 1000 }}\n styles={{\n input: { height: \"1.625rem\", minHeight: \"1.625rem\" },\n // icon: { transform: \"scale(0.8)\" },\n }}\n />\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/Rgba.tsx", "content": "import * as React from \"react\";\nimport { ColorInput } from \"@mantine/core\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { rgbaToHex, hexToRgba } from \"./utils\";\nimport { ViserInputComponent } from \"./common\";\nimport { GuiAddRgbaMessage } from \"../WebsocketMessages\";\n\nexport default function RgbaComponent({\n id,\n label,\n hint,\n value,\n disabled,\n visible,\n}: GuiAddRgbaMessage) {\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n setValue(id, hexToRgba(v))}\n format=\"hexa\"\n // zIndex of dropdown should be >modal zIndex.\n // On edge cases: it seems like existing dropdowns are always closed when a new modal is opened.\n popoverProps={{ zIndex: 1000 }}\n styles={{\n input: { height: \"1.625rem\", minHeight: \"1.625rem\" },\n }}\n />\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/Slider.tsx", "content": "import React from \"react\";\nimport { GuiAddSliderMessage } from \"../WebsocketMessages\";\nimport {\n Slider,\n Flex,\n NumberInput,\n useMantineColorScheme,\n} from \"@mantine/core\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { ViserInputComponent } from \"./common\";\nimport { sliderDefaultMarks } from \"./ComponentStyles.css\";\n\nexport default function SliderComponent({\n id,\n label,\n hint,\n visible,\n disabled,\n value,\n ...otherProps\n}: GuiAddSliderMessage) {\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n const updateValue = (value: number) => setValue(id, value);\n const { min, max, precision, step, marks } = otherProps;\n const colorScheme = useMantineColorScheme().colorScheme;\n const input = (\n \n ({\n thumb: {\n height: \"0.75rem\",\n width: \"0.5rem\",\n },\n trackContainer: {\n zIndex: 3,\n position: \"relative\",\n },\n markLabel: {\n transform: \"translate(-50%, 0.03rem)\",\n fontSize: \"0.6rem\",\n textAlign: \"center\",\n },\n mark: {\n transform: \"scale(1.95)\",\n },\n markFilled: {\n background: disabled\n ? colorScheme === \"dark\"\n ? theme.colors.dark[3]\n : theme.colors.gray[4]\n : theme.primaryColor,\n },\n })}\n pt=\"0.3em\"\n pb=\"0.2em\"\n showLabelOnHover={false}\n min={min}\n max={max}\n step={step ?? undefined}\n precision={precision}\n value={value}\n onChange={updateValue}\n marks={\n marks === null\n ? [\n {\n value: min,\n label: `${parseInt(min.toFixed(6))}`,\n },\n {\n value: max,\n label: `${parseInt(max.toFixed(6))}`,\n },\n ]\n : marks\n }\n disabled={disabled}\n />\n {\n // Ignore empty values.\n newValue !== \"\" && updateValue(Number(newValue));\n }}\n size=\"xs\"\n min={min}\n max={max}\n hideControls\n step={step ?? undefined}\n // precision={precision}\n style={{ width: \"3rem\" }}\n styles={{\n input: {\n padding: \"0.375em\",\n letterSpacing: \"-0.5px\",\n minHeight: \"1.875em\",\n height: \"1.875em\",\n },\n }}\n ml=\"xs\"\n />\n \n );\n\n return (\n {input}\n );\n}\n" }, { "path": "viser/src/viser/client/src/components/TabGroup.tsx", "content": "import * as React from \"react\";\nimport { GuiAddTabGroupMessage } from \"../WebsocketMessages\";\nimport { Tabs } from \"@mantine/core\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { htmlIconWrapper } from \"./ComponentStyles.css\";\n\nexport default function TabGroupComponent({\n tab_labels,\n tab_icons_html,\n tab_container_ids,\n visible,\n}: GuiAddTabGroupMessage) {\n const { GuiContainer } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n \n {tab_labels.map((label, index) => (\n \n )\n }\n >\n {label}\n \n ))}\n \n {tab_container_ids.map((containerId, index) => (\n \n \n \n ))}\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/TextInput.tsx", "content": "import * as React from \"react\";\nimport { TextInput } from \"@mantine/core\";\nimport { ViserInputComponent } from \"./common\";\nimport { GuiAddTextMessage } from \"../WebsocketMessages\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\n\nexport default function TextInputComponent(props: GuiAddTextMessage) {\n const { id, hint, label, value, disabled, visible } = props;\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n {\n setValue(id, value.target.value);\n }}\n styles={{\n input: {\n minHeight: \"1.625rem\",\n height: \"1.625rem\",\n padding: \"0 0.5em\",\n },\n }}\n disabled={disabled}\n />\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/UploadButton.tsx", "content": "import { GuiAddUploadButtonMessage } from \"../WebsocketMessages\";\nimport { v4 as uuid } from \"uuid\";\nimport { Box, Progress } from \"@mantine/core\";\n\nimport { Button } from \"@mantine/core\";\nimport React, { useContext } from \"react\";\nimport { ViewerContext, ViewerContextContents } from \"../App\";\nimport { IconCheck } from \"@tabler/icons-react\";\nimport { notifications } from \"@mantine/notifications\";\nimport { htmlIconWrapper } from \"./ComponentStyles.css\";\n\nexport default function UploadButtonComponent(conf: GuiAddUploadButtonMessage) {\n // Handle GUI input types.\n const viewer = useContext(ViewerContext)!;\n const fileUploadRef = React.useRef(null);\n const { isUploading, upload } = useFileUpload({\n viewer,\n componentId: conf.id,\n });\n\n const disabled = conf.disabled || isUploading;\n return (\n \n {\n const input = e.target as HTMLInputElement;\n if (!input.files) return;\n upload(input.files[0]);\n }}\n />\n {\n if (fileUploadRef.current === null) return;\n fileUploadRef.current.value = fileUploadRef.current.defaultValue;\n fileUploadRef.current.click();\n }}\n style={{ height: \"2.125em\" }}\n disabled={disabled}\n size=\"sm\"\n leftSection={\n conf.icon_html === null ? undefined : (\n \n )\n }\n >\n {conf.label}\n \n \n );\n}\n\nfunction useFileUpload({\n viewer,\n componentId,\n}: {\n componentId: string;\n viewer: ViewerContextContents;\n}) {\n const updateUploadState = viewer.useGui((state) => state.updateUploadState);\n const uploadState = viewer.useGui(\n (state) => state.uploadsInProgress[componentId],\n );\n const totalBytes = uploadState?.totalBytes;\n\n // Cache total bytes string\n const totalBytesString = React.useMemo(() => {\n if (totalBytes === undefined) return \"\";\n let displaySize = totalBytes;\n const displayUnits = [\"B\", \"K\", \"M\", \"G\", \"T\", \"P\"];\n let displayUnitIndex = 0;\n while (displaySize >= 100 && displayUnitIndex < displayUnits.length - 1) {\n displaySize /= 1024;\n displayUnitIndex += 1;\n }\n return `${displaySize.toFixed(1)}${displayUnits[displayUnitIndex]}`;\n }, [totalBytes]);\n\n // Update notification status\n React.useEffect(() => {\n if (uploadState === undefined) return;\n const { notificationId, filename } = uploadState;\n if (uploadState.uploadedBytes === 0) {\n // Show notification.\n notifications.show({\n id: notificationId,\n title: \"Uploading \" + `${filename} (${totalBytesString})`,\n message: ,\n autoClose: false,\n withCloseButton: false,\n loading: true,\n });\n } else {\n // Update progress.\n const progressValue = uploadState.uploadedBytes / uploadState.totalBytes;\n const isDone = progressValue === 1.0;\n notifications.update({\n id: notificationId,\n title: \"Uploading \" + `${filename} (${totalBytesString})`,\n message: !isDone ? (\n \n ) : (\n \"File uploaded successfully.\"\n ),\n autoClose: isDone,\n withCloseButton: isDone,\n loading: !isDone,\n icon: isDone ? : undefined,\n });\n }\n }, [uploadState, totalBytesString]);\n\n const isUploading =\n uploadState !== undefined &&\n uploadState.uploadedBytes < uploadState.totalBytes;\n\n async function upload(file: File) {\n const chunkSize = 512 * 1024; // bytes\n const numChunks = Math.ceil(file.size / chunkSize);\n const transferUuid = uuid();\n const notificationId = \"upload-\" + transferUuid;\n\n // Begin upload by setting initial state\n updateUploadState({\n componentId: componentId,\n uploadedBytes: 0,\n totalBytes: file.size,\n filename: file.name,\n notificationId,\n });\n\n viewer.sendMessageRef.current({\n type: \"FileTransferStart\",\n source_component_id: componentId,\n transfer_uuid: transferUuid,\n filename: file.name,\n mime_type: file.type,\n size_bytes: file.size,\n part_count: numChunks,\n });\n\n for (let i = 0; i < numChunks; i++) {\n const start = i * chunkSize;\n const end = (i + 1) * chunkSize;\n const chunk = file.slice(start, end);\n const buffer = await chunk.arrayBuffer();\n\n viewer.sendMessageRef.current({\n type: \"FileTransferPart\",\n source_component_id: componentId,\n transfer_uuid: transferUuid,\n part: i,\n content: new Uint8Array(buffer),\n });\n }\n }\n\n return {\n isUploading,\n upload,\n };\n}\n" }, { "path": "viser/src/viser/client/src/components/Vector2.tsx", "content": "import * as React from \"react\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { GuiAddVector2Message } from \"../WebsocketMessages\";\nimport { VectorInput, ViserInputComponent } from \"./common\";\n\nexport default function Vector2Component({\n id,\n hint,\n label,\n visible,\n disabled,\n value,\n ...otherProps\n}: GuiAddVector2Message) {\n const { min, max, step, precision } = otherProps;\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n setValue(id, value)}\n min={min}\n max={max}\n step={step}\n precision={precision}\n disabled={disabled}\n />\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/Vector3.tsx", "content": "import * as React from \"react\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\nimport { GuiAddVector3Message } from \"../WebsocketMessages\";\nimport { VectorInput, ViserInputComponent } from \"./common\";\n\nexport default function Vector3Component({\n id,\n hint,\n label,\n visible,\n disabled,\n value,\n ...otherProps\n}: GuiAddVector3Message) {\n const { min, max, step, precision } = otherProps;\n const { setValue } = React.useContext(GuiComponentContext)!;\n if (!visible) return <>;\n return (\n \n setValue(id, value)}\n min={min}\n max={max}\n step={step}\n precision={precision}\n disabled={disabled}\n />\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/common.tsx", "content": "import * as React from \"react\";\nimport { Box, Flex, Text, NumberInput, Tooltip } from \"@mantine/core\";\nimport { GuiComponentContext } from \"../ControlPanel/GuiComponentContext\";\n\nexport function ViserInputComponent({\n id,\n label,\n hint,\n children,\n}: {\n id: string;\n children: React.ReactNode;\n label?: string;\n hint?: string | null;\n}) {\n const { folderDepth } = React.useContext(GuiComponentContext)!;\n if (hint !== undefined && hint !== null) {\n children = // We need to add for inputs that we can't assign refs to.\n (\n \n {children}\n \n );\n }\n\n if (label !== undefined)\n children = (\n \n );\n\n return (\n \n {children}\n \n );\n}\n\n/** GUI input with a label horizontally placed to the left of it. */\nfunction LabeledInput(props: {\n id: string;\n label: string;\n input: React.ReactNode;\n folderDepth: number;\n}) {\n return (\n \n \n \n \n \n \n {props.input}\n \n );\n}\n\nexport function VectorInput(\n props:\n | {\n id: string;\n n: 2;\n value: [number, number];\n min: [number, number] | null;\n max: [number, number] | null;\n step: number;\n precision: number;\n onChange: (value: number[]) => void;\n disabled: boolean;\n }\n | {\n id: string;\n n: 3;\n value: [number, number, number];\n min: [number, number, number] | null;\n max: [number, number, number] | null;\n step: number;\n precision: number;\n onChange: (value: number[]) => void;\n disabled: boolean;\n },\n) {\n return (\n \n {[...Array(props.n).keys()].map((i) => (\n {\n const updated = [...props.value];\n updated[i] = v === \"\" ? 0.0 : Number(v);\n props.onChange(updated);\n }}\n size=\"xs\"\n styles={{\n root: { flexGrow: 1, width: 0 },\n input: {\n paddingLeft: \"0.5em\",\n paddingRight: \"1.75em\",\n textAlign: \"right\",\n height: \"1.875em\",\n minHeight: \"1.875em\",\n },\n controls: {\n height: \"1.25em\",\n width: \"0.825em\",\n },\n }}\n rightSectionWidth=\"1em\"\n decimalScale={props.precision}\n step={props.step}\n min={props.min === null ? undefined : props.min[i]}\n max={props.max === null ? undefined : props.max[i]}\n stepHoldDelay={500}\n stepHoldInterval={(t) => Math.max(1000 / t ** 2, 25)}\n disabled={props.disabled}\n />\n ))}\n \n );\n}\n" }, { "path": "viser/src/viser/client/src/components/utils.tsx", "content": "// Color conversion helpers.\n\nexport function rgbToHex([r, g, b]: [number, number, number]): string {\n const hexR = r.toString(16).padStart(2, \"0\");\n const hexG = g.toString(16).padStart(2, \"0\");\n const hexB = b.toString(16).padStart(2, \"0\");\n return `#${hexR}${hexG}${hexB}`;\n}\n\nexport function hexToRgb(hexColor: string): [number, number, number] {\n const hex = hexColor.slice(1); // Remove the # in #ffffff.\n const r = parseInt(hex.substring(0, 2), 16);\n const g = parseInt(hex.substring(2, 4), 16);\n const b = parseInt(hex.substring(4, 6), 16);\n return [r, g, b];\n}\nexport function rgbaToHex([r, g, b, a]: [\n number,\n number,\n number,\n number,\n]): string {\n const hexR = r.toString(16).padStart(2, \"0\");\n const hexG = g.toString(16).padStart(2, \"0\");\n const hexB = b.toString(16).padStart(2, \"0\");\n const hexA = a.toString(16).padStart(2, \"0\");\n return `#${hexR}${hexG}${hexB}${hexA}`;\n}\n\nexport function hexToRgba(hexColor: string): [number, number, number, number] {\n const hex = hexColor.slice(1); // Remove the # in #ffffff.\n const r = parseInt(hex.substring(0, 2), 16);\n const g = parseInt(hex.substring(2, 4), 16);\n const b = parseInt(hex.substring(4, 6), 16);\n const a = parseInt(hex.substring(6, 8), 16);\n return [r, g, b, a];\n}\n" }, { "path": "viser/src/viser/client/src/index.css", "content": "@font-face {\n font-family: \"Inter\";\n src: url(\"/Inter-VariableFont_slnt,wght.ttf\") format(\"truetype\");\n font-weight: 1 100 200 300 400 500 600 700 800 900 1000;\n font-style: normal italic;\n}\n\nbody,\nhtml {\n width: 100%;\n height: 100%;\n margin: 0;\n padding: 0;\n overflow: hidden;\n}\n\nbody {\n font-family:\n \"Inter\",\n -apple-system,\n BlinkMacSystemFont,\n \"Segoe UI\",\n \"Roboto\",\n \"Oxygen\",\n \"Ubuntu\",\n \"Cantarell\",\n \"Fira Sans\",\n \"Droid Sans\",\n \"Helvetica Neue\",\n sans-serif;\n -webkit-font-smoothing: antialiased;\n -moz-osx-font-smoothing: grayscale;\n}\n\ncode {\n font-family: source-code-pro, Menlo, Monaco, Consolas, \"Courier New\",\n monospace;\n}\n\n#root {\n width: 100%;\n height: 100%;\n overflow: hidden;\n}\n\n/* Styling for threejs stats panel. Switches position: fixed to position:\n * absolute, to respect parent + for better multi-pane layout. */\n.stats-panel {\n position: absolute !important;\n}\n\n/* Styling for color chips in markdown */\n.gfm-color-chip {\n margin-left: 0.125rem;\n display: inline-block;\n height: 0.625rem;\n width: 0.625rem;\n border-radius: 9999px;\n border: 1px solid gray;\n transform: TranslateY(0.07em);\n}\n" }, { "path": "viser/src/viser/client/src/index.tsx", "content": "import ReactDOM from \"react-dom/client\";\nimport { Root } from \"./App\";\nimport { enableMapSet } from \"immer\";\n\nenableMapSet();\n\nReactDOM.createRoot(document.getElementById(\"root\") as HTMLElement).render(\n ,\n);\n" }, { "path": "viser/src/viser/client/src/react-app-env.d.ts", "content": "/// \n" }, { "path": "viser/src/viser/client/tsconfig.json", "content": "{\n \"compilerOptions\": {\n \"target\": \"ESNext\",\n \"lib\": [\"dom\", \"dom.iterable\", \"esnext\"],\n \"types\": [\n \"vite/client\",\n \"vite-plugin-svgr/client\",\n \"node\",\n \"@types/wicg-file-system-access\"\n ],\n \"allowJs\": true,\n \"skipLibCheck\": true,\n \"esModuleInterop\": true,\n \"allowSyntheticDefaultImports\": true,\n \"strict\": true,\n \"forceConsistentCasingInFileNames\": true,\n \"noFallthroughCasesInSwitch\": true,\n \"module\": \"esnext\",\n \"moduleResolution\": \"node\",\n \"resolveJsonModule\": true,\n \"isolatedModules\": true,\n \"noEmit\": true,\n \"jsx\": \"react-jsx\"\n },\n \"include\": [\"src\"]\n}\n" }, { "path": "viser/src/viser/client/vite-env.d.ts", "content": "/// \n" }, { "path": "viser/src/viser/client/vite.config.mts", "content": "import { defineConfig } from \"vite\";\nimport react from \"@vitejs/plugin-react\";\nimport { vanillaExtractPlugin } from \"@vanilla-extract/vite-plugin\";\n\nimport viteTsconfigPaths from \"vite-tsconfig-paths\";\nimport svgrPlugin from \"vite-plugin-svgr\";\nimport eslint from \"vite-plugin-eslint\";\nimport browserslistToEsbuild from \"browserslist-to-esbuild\";\n\n// https://vitejs.dev/config/\nexport default defineConfig({\n plugins: [\n react(),\n eslint({ failOnError: false, failOnWarning: false }),\n viteTsconfigPaths(),\n svgrPlugin(),\n vanillaExtractPlugin(),\n ],\n server: {\n port: 3000,\n hmr: { port: 1025 },\n },\n worker: {\n format: \"es\",\n },\n build: {\n outDir: \"build\",\n target: browserslistToEsbuild(),\n },\n});\n" }, { "path": "viser/src/viser/extras/__init__.py", "content": "\"\"\"Extra utilities. Used for example scripts.\"\"\"\n\nfrom ._record3d import Record3dFrame as Record3dFrame\nfrom ._record3d import Record3dLoader as Record3dLoader\nfrom ._record3d_customized import Record3dLoader_Customized as Record3dLoader_Customized\nfrom ._record3d_customized_megasam import Record3dLoader_Customized_Megasam as Record3dLoader_Customized_Megasam\nfrom ._urdf import ViserUrdf as ViserUrdf\n" }, { "path": "viser/src/viser/extras/_record3d.py", "content": "from __future__ import annotations\n\nimport dataclasses\nimport json\nfrom pathlib import Path\nfrom typing import Tuple, cast\n\nimport imageio.v3 as iio\nimport liblzfse\nimport numpy as np\nimport numpy as onp\nimport numpy.typing as onpt\nimport skimage.transform\nfrom scipy.spatial.transform import Rotation\n\n\nclass Record3dLoader:\n \"\"\"Helper for loading frames for Record3D captures.\"\"\"\n\n # NOTE(hangg): Consider moving this module into\n # `examples/7_record3d_visualizer.py` since it is usecase-specific.\n\n def __init__(self, data_dir: Path):\n metadata_path = data_dir / \"metadata\"\n\n # Read metadata.\n metadata = json.loads(metadata_path.read_text())\n\n K: onp.ndarray = np.array(metadata[\"K\"], np.float32).reshape(3, 3).T\n fps = metadata[\"fps\"]\n\n T_world_cameras: onp.ndarray = np.array(metadata[\"poses\"], np.float32)\n T_world_cameras = np.concatenate(\n [\n Rotation.from_quat(T_world_cameras[:, :4]).as_matrix(),\n T_world_cameras[:, 4:, None],\n ],\n -1,\n )\n T_world_cameras = (T_world_cameras @ np.diag([1, -1, -1, 1])).astype(np.float32)\n\n self.K = K\n self.fps = fps\n self.T_world_cameras = T_world_cameras\n\n rgbd_dir = data_dir / \"rgbd\"\n self.rgb_paths = sorted(rgbd_dir.glob(\"*.jpg\"), key=lambda p: int(p.stem))\n self.depth_paths = [\n rgb_path.with_suffix(\".depth\") for rgb_path in self.rgb_paths\n ]\n self.conf_paths = [rgb_path.with_suffix(\".conf\") for rgb_path in self.rgb_paths]\n\n def num_frames(self) -> int:\n return len(self.rgb_paths)\n\n def get_frame(self, index: int) -> Record3dFrame:\n # Read conf.\n conf: onp.ndarray = np.frombuffer(\n liblzfse.decompress(self.conf_paths[index].read_bytes()), dtype=np.uint8\n )\n if conf.shape[0] == 640 * 480:\n conf = conf.reshape((640, 480)) # For a FaceID camera 3D Video\n elif conf.shape[0] == 256 * 192:\n conf = conf.reshape((256, 192)) # For a LiDAR 3D Video\n else:\n assert False, f\"Unexpected conf shape {conf.shape}\"\n\n # Read depth.\n depth: onp.ndarray = np.frombuffer(\n liblzfse.decompress(self.depth_paths[index].read_bytes()), dtype=np.float32\n ).copy()\n if depth.shape[0] == 640 * 480:\n depth = depth.reshape((640, 480)) # For a FaceID camera 3D Video\n elif depth.shape[0] == 256 * 192:\n depth = depth.reshape((256, 192)) # For a LiDAR 3D Video\n else:\n assert False, f\"Unexpected depth shape {depth.shape}\"\n\n # Read RGB.\n rgb = iio.imread(self.rgb_paths[index])\n return Record3dFrame(\n K=self.K,\n rgb=rgb,\n depth=depth,\n mask=conf == 2,\n T_world_camera=self.T_world_cameras[index],\n )\n\n\n@dataclasses.dataclass\nclass Record3dFrame:\n \"\"\"A single frame from a Record3D capture.\"\"\"\n\n K: onpt.NDArray[onp.float32]\n rgb: onpt.NDArray[onp.uint8]\n depth: onpt.NDArray[onp.float32]\n mask: onpt.NDArray[onp.bool_]\n T_world_camera: onpt.NDArray[onp.float32]\n\n def get_point_cloud(\n self, downsample_factor: int = 1\n ) -> Tuple[onpt.NDArray[onp.float32], onpt.NDArray[onp.uint8]]:\n rgb = self.rgb[::downsample_factor, ::downsample_factor]\n depth = skimage.transform.resize(self.depth, rgb.shape[:2], order=0)\n mask = cast(\n onpt.NDArray[onp.bool_],\n skimage.transform.resize(self.mask, rgb.shape[:2], order=0),\n )\n assert depth.shape == rgb.shape[:2]\n\n K = self.K\n T_world_camera = self.T_world_camera\n\n img_wh = rgb.shape[:2][::-1]\n\n grid = (\n np.stack(np.meshgrid(np.arange(img_wh[0]), np.arange(img_wh[1])), 2) + 0.5\n )\n grid = grid * downsample_factor\n\n homo_grid = np.pad(grid[mask], np.array([[0, 0], [0, 1]]), constant_values=1)\n local_dirs = np.einsum(\"ij,bj->bi\", np.linalg.inv(K), homo_grid)\n dirs = np.einsum(\"ij,bj->bi\", T_world_camera[:3, :3], local_dirs)\n points = (T_world_camera[:, -1] + dirs * depth[mask, None]).astype(np.float32)\n point_colors = rgb[mask]\n\n return points, point_colors\n" }, { "path": "viser/src/viser/extras/_record3d_customized.py", "content": "from __future__ import annotations\n\nimport dataclasses\nimport os\nimport json\nfrom pathlib import Path\nfrom typing import Tuple, cast\n\nimport imageio.v3 as iio\nimport liblzfse\nimport numpy as np\nimport numpy as onp\nimport numpy.typing as onpt\nimport skimage.transform\nfrom scipy.spatial.transform import Rotation\nfrom scipy.spatial import cKDTree\n\nclass Record3dLoader_Customized:\n \"\"\"Helper for loading frames for Record3D captures.\"\"\"\n\n def __init__(self, data_dir: Path, conf_threshold: float = 1.0, foreground_conf_threshold: float = 0.1, no_mask: bool = False, xyzw=True, init_conf=False):\n\n # Read metadata.\n intrinsics_path = data_dir / \"pred_intrinsics.txt\"\n intrinsics = np.loadtxt(intrinsics_path)\n\n self.K: onp.ndarray = np.array(intrinsics, np.float32).reshape(-1, 3, 3)\n fps = 30\n\n self.init_conf = init_conf\n\n poses_path = data_dir / \"pred_traj.txt\"\n poses = np.loadtxt(poses_path)\n self.T_world_cameras: onp.ndarray = np.array(poses, np.float32)\n self.T_world_cameras = np.concatenate(\n [\n # Convert TUM pose to SE3 pose\n Rotation.from_quat(self.T_world_cameras[:, 4:]).as_matrix() if not xyzw\n else Rotation.from_quat(np.concatenate([self.T_world_cameras[:, 5:], self.T_world_cameras[:, 4:5]], -1)).as_matrix(),\n self.T_world_cameras[:, 1:4, None],\n ],\n -1,\n )\n self.T_world_cameras = self.T_world_cameras.astype(np.float32)\n\n # Convert to homogeneous transformation matrices (ensure shape is (N, 4, 4))\n num_frames = self.T_world_cameras.shape[0]\n ones = np.tile(np.array([0, 0, 0, 1], dtype=np.float32), (num_frames, 1, 1))\n self.T_world_cameras = np.concatenate([self.T_world_cameras, ones], axis=1)\n\n self.fps = fps\n self.conf_threshold = conf_threshold\n self.foreground_conf_threshold = foreground_conf_threshold\n self.no_mask = no_mask\n\n # Read frames.\n self.rgb_paths = sorted(data_dir.glob(\"frame_*.png\"), key=lambda p: int(p.stem.split(\"_\")[-1]))\n self.depth_paths = sorted(data_dir.glob(\"frame_*.npy\"), key=lambda p: int(p.stem.split(\"_\")[-1]))\n if init_conf:\n self.init_conf_paths = sorted(data_dir.glob(\"init_conf_*.npy\"), key=lambda p: int(p.stem.split(\"_\")[-1]))\n else:\n self.init_conf_paths = []\n self.conf_paths = sorted(data_dir.glob(\"conf_*.npy\"), key=lambda p: int(p.stem.split(\"_\")[-1]))\n self.mask_paths = sorted(data_dir.glob(\"enlarged_dynamic_mask_*.png\"), key=lambda p: int(p.stem.split(\"_\")[-1]))\n\n # Remove the last frame since it does not have a ground truth dynamic mask\n self.rgb_paths = self.rgb_paths[:-1]\n\n # Align all camera poses by the first frame\n T0 = self.T_world_cameras[len(self.T_world_cameras) // 2] # First camera pose (4x4 matrix)\n T0_inv = np.linalg.inv(T0) # Inverse of the first camera pose\n\n # Apply T0_inv to all camera poses\n self.T_world_cameras = np.matmul(T0_inv[np.newaxis, :, :], self.T_world_cameras)\n\n\n def num_frames(self) -> int:\n return len(self.rgb_paths)\n\n def get_frame(self, index: int) -> Record3dFrame:\n\n # Read depth.\n depth = np.load(self.depth_paths[index])\n depth: onp.NDArray[onp.float32] = depth\n \n # Check if conf file exists, otherwise initialize with ones\n if len(self.conf_paths) == 0:\n conf = np.ones_like(depth, dtype=onp.float32)\n else:\n conf_path = self.conf_paths[index]\n if os.path.exists(conf_path):\n conf = np.load(conf_path)\n conf: onpt.NDArray[onp.float32] = conf\n # Clip confidence to avoid negative values\n conf = np.clip(conf, 0.0001, 99999)\n else:\n conf = np.ones_like(depth, dtype=onp.float32)\n\n # Check if init conf file exists, otherwise initialize with ones\n if len(self.init_conf_paths) == 0: # If init conf is not available, use conf\n init_conf = conf\n else:\n init_conf_path = self.init_conf_paths[index]\n if os.path.exists(init_conf_path):\n init_conf = np.load(init_conf_path)\n init_conf: onpt.NDArray[onp.float32] = init_conf\n # Clip confidence to avoid negative values\n init_conf = np.clip(init_conf, 0.0001, 99999)\n else:\n init_conf = np.ones_like(depth, dtype=onp.float32)\n \n # Check if mask file exists, otherwise initialize with zeros\n if len(self.mask_paths) == 0:\n mask = np.ones_like(depth, dtype=onp.bool_)\n else:\n mask_path = self.mask_paths[index]\n if os.path.exists(mask_path):\n mask = iio.imread(mask_path) > 0\n mask: onpt.NDArray[onp.bool_] = mask\n else:\n mask = np.ones_like(depth, dtype=onp.bool_)\n\n if self.no_mask:\n mask = np.ones_like(mask).astype(np.bool_)\n\n # Read RGB.\n rgb = iio.imread(self.rgb_paths[index])\n # if 4 channels, remove the alpha channel\n if rgb.shape[-1] == 4:\n rgb = rgb[..., :3]\n\n return Record3dFrame(\n K=self.K[index],\n rgb=rgb,\n depth=depth,\n mask=mask,\n conf=conf,\n init_conf=init_conf,\n T_world_camera=self.T_world_cameras[index],\n conf_threshold=self.conf_threshold,\n foreground_conf_threshold=self.foreground_conf_threshold,\n )\n\n\n@dataclasses.dataclass\nclass Record3dFrame:\n \"\"\"A single frame from a Record3D capture.\"\"\"\n\n K: onpt.NDArray[onp.float32]\n rgb: onpt.NDArray[onp.uint8]\n depth: onpt.NDArray[onp.float32]\n mask: onpt.NDArray[onp.bool_]\n conf: onpt.NDArray[onp.float32]\n init_conf: onpt.NDArray[onp.float32]\n T_world_camera: onpt.NDArray[onp.float32]\n conf_threshold: float = 1.0\n foreground_conf_threshold: float = 0.1\n\n def get_point_cloud(\n self, downsample_factor: int = 1, bg_downsample_factor: int = 1,\n ) -> Tuple[onpt.NDArray[onp.float32], onpt.NDArray[onp.uint8], onpt.NDArray[onp.float32], onpt.NDArray[onp.uint8]]:\n rgb = self.rgb[::downsample_factor, ::downsample_factor]\n depth = skimage.transform.resize(self.depth, rgb.shape[:2], order=0)\n mask = cast(\n onpt.NDArray[onp.bool_],\n skimage.transform.resize(self.mask, rgb.shape[:2], order=0),\n )\n assert depth.shape == rgb.shape[:2]\n\n K = self.K\n T_world_camera = self.T_world_camera\n\n img_wh = rgb.shape[:2][::-1]\n\n grid = (\n np.stack(np.meshgrid(np.arange(img_wh[0]), np.arange(img_wh[1])), 2) + 0.5\n )\n grid = grid * downsample_factor\n conf_mask = self.conf > self.conf_threshold\n if self.init_conf is not None:\n fg_conf_mask = self.init_conf > self.foreground_conf_threshold\n else:\n fg_conf_mask = self.conf > self.foreground_conf_threshold\n # reshape the conf mask to the shape of the depth\n conf_mask = skimage.transform.resize(conf_mask, depth.shape, order=0)\n fg_conf_mask = skimage.transform.resize(fg_conf_mask, depth.shape, order=0)\n\n # Foreground points\n homo_grid = np.pad(grid[fg_conf_mask & mask], ((0, 0), (0, 1)), constant_values=1)\n local_dirs = np.einsum(\"ij,bj->bi\", np.linalg.inv(K), homo_grid)\n dirs = np.einsum(\"ij,bj->bi\", T_world_camera[:3, :3], local_dirs)\n points = (T_world_camera[:3, 3] + dirs * depth[fg_conf_mask & mask, None]).astype(np.float32)\n point_colors = rgb[fg_conf_mask & mask]\n\n # Background points\n bg_homo_grid = np.pad(grid[conf_mask & ~mask], ((0, 0), (0, 1)), constant_values=1)\n bg_local_dirs = np.einsum(\"ij,bj->bi\", np.linalg.inv(K), bg_homo_grid)\n bg_dirs = np.einsum(\"ij,bj->bi\", T_world_camera[:3, :3], bg_local_dirs)\n bg_points = (T_world_camera[:3, 3] + bg_dirs * depth[conf_mask & ~mask, None]).astype(np.float32)\n bg_point_colors = rgb[conf_mask & ~mask]\n\n if bg_downsample_factor > 1 and bg_points.shape[0] > 0:\n indices = np.random.choice(\n bg_points.shape[0],\n size=bg_points.shape[0] // bg_downsample_factor,\n replace=False\n )\n bg_points = bg_points[indices]\n bg_point_colors = bg_point_colors[indices]\n\n return points, point_colors, bg_points, bg_point_colors\n" }, { "path": "viser/src/viser/extras/_record3d_customized_megasam.py", "content": "from __future__ import annotations\n\nimport dataclasses\nimport os\nimport json\nfrom pathlib import Path\nfrom typing import Tuple, cast\n\nimport imageio.v3 as iio\nimport liblzfse\nimport numpy as np\nimport numpy as onp\nimport numpy.typing as onpt\nimport skimage.transform\nfrom scipy.spatial.transform import Rotation\nfrom scipy.spatial import cKDTree\n\nclass Record3dLoader_Customized_Megasam:\n \"\"\"Helper for loading frames for Record3D captures directly from a NPZ file.\"\"\"\n\n def __init__(self, npz_data: dict, conf_threshold: float = 1.0, foreground_conf_threshold: float = 0.1, no_mask: bool = False, xyzw=True, init_conf=False):\n # Assuming npz_data is a dictionary containing all the necessary arrays from the NPZ file\n self.K = np.expand_dims(npz_data['intrinsic'], 0) # (3,3) -> (1,3,3) Intrinsic matrix\n self.K = np.repeat(self.K, npz_data['images'].shape[0], axis=0) # (1,3,3) -> (N,3,3)\n self.T_world_cameras = npz_data['cam_c2w'] # (N,4,4) Camera poses (extrinsics)\n self.fps = 30 # Assuming a frame rate of 30\n self.conf_threshold = conf_threshold\n self.foreground_conf_threshold = foreground_conf_threshold\n self.no_mask = no_mask\n\n # Initialize the other parameters\n self.init_conf = init_conf\n \n # Read frames from the NPZ file\n self.images = npz_data['images'] # (N,H,W,3) RGB images\n self.depths = npz_data['depths'] # (N,H,W) Depth maps\n self.confidences = npz_data.get('conf', [])\n self.init_conf_data = npz_data.get('init_conf', [])\n self.masks = npz_data.get('enlarged_dynamic_mask', [])\n\n # Align all camera poses by the first frame\n T0 = self.T_world_cameras[len(self.T_world_cameras) // 2] # First camera pose (4x4 matrix)\n T0_inv = np.linalg.inv(T0) # Inverse of the first camera pose\n\n # Apply T0_inv to all camera poses\n self.T_world_cameras = np.matmul(T0_inv[np.newaxis, :, :], self.T_world_cameras)\n\n def num_frames(self) -> int:\n return len(self.images)\n\n def get_frame(self, index: int) -> Record3dFrame:\n # Read the depth for the given frame\n depth = self.depths[index]\n depth = depth.astype(np.float32)\n\n # Check if conf file exists, otherwise initialize with ones\n if len(self.confidences) == 0:\n conf = np.ones_like(depth, dtype=np.float32)\n else:\n conf = self.confidences[index]\n conf = np.clip(conf, 0.0001, 99999)\n\n # Check if init conf file exists, otherwise initialize with ones\n if len(self.init_conf_data) == 0:\n init_conf = conf\n else:\n init_conf = self.init_conf_data[index]\n init_conf = np.clip(init_conf, 0.0001, 99999)\n \n # Check if mask exists, otherwise initialize with zeros\n if len(self.masks) == 0:\n mask = np.ones_like(depth, dtype=bool)\n else:\n mask = self.masks[index] > 0 # Assuming mask is a binary image\n\n if self.no_mask:\n mask = np.ones_like(mask).astype(np.bool_)\n\n # Read RGB image\n rgb = self.images[index]\n if rgb.shape[-1] == 4:\n rgb = rgb[..., :3]\n\n return Record3dFrame(\n K=self.K[index],\n rgb=rgb,\n depth=depth,\n mask=mask,\n conf=conf,\n init_conf=init_conf,\n T_world_camera=self.T_world_cameras[index],\n conf_threshold=self.conf_threshold,\n foreground_conf_threshold=self.foreground_conf_threshold,\n )\n\n\n@dataclasses.dataclass\nclass Record3dFrame:\n \"\"\"A single frame from a Record3D capture.\"\"\"\n\n K: onpt.NDArray[onp.float32]\n rgb: onpt.NDArray[onp.uint8]\n depth: onpt.NDArray[onp.float32]\n mask: onpt.NDArray[onp.bool_]\n conf: onpt.NDArray[onp.float32]\n init_conf: onpt.NDArray[onp.float32]\n T_world_camera: onpt.NDArray[onp.float32]\n conf_threshold: float = 1.0\n foreground_conf_threshold: float = 0.1\n\n def get_point_cloud(\n self, downsample_factor: int = 1, bg_downsample_factor: int = 1,\n ) -> Tuple[onpt.NDArray[onp.float32], onpt.NDArray[onp.uint8], onpt.NDArray[onp.float32], onpt.NDArray[onp.uint8]]:\n rgb = self.rgb[::downsample_factor, ::downsample_factor]\n depth = skimage.transform.resize(self.depth, rgb.shape[:2], order=0)\n mask = cast(\n onpt.NDArray[onp.bool_],\n skimage.transform.resize(self.mask, rgb.shape[:2], order=0),\n )\n assert depth.shape == rgb.shape[:2]\n\n K = self.K\n T_world_camera = self.T_world_camera\n\n img_wh = rgb.shape[:2][::-1]\n\n grid = (\n np.stack(np.meshgrid(np.arange(img_wh[0]), np.arange(img_wh[1])), 2) + 0.5\n )\n grid = grid * downsample_factor\n conf_mask = self.conf > self.conf_threshold\n if self.init_conf is not None:\n fg_conf_mask = self.init_conf > self.foreground_conf_threshold\n else:\n fg_conf_mask = self.conf > self.foreground_conf_threshold\n # reshape the conf mask to the shape of the depth\n conf_mask = skimage.transform.resize(conf_mask, depth.shape, order=0)\n fg_conf_mask = skimage.transform.resize(fg_conf_mask, depth.shape, order=0)\n\n # Foreground points\n homo_grid = np.pad(grid[fg_conf_mask & mask], ((0, 0), (0, 1)), constant_values=1)\n local_dirs = np.einsum(\"ij,bj->bi\", np.linalg.inv(K), homo_grid)\n dirs = np.einsum(\"ij,bj->bi\", T_world_camera[:3, :3], local_dirs)\n points = (T_world_camera[:3, 3] + dirs * depth[fg_conf_mask & mask, None]).astype(np.float32)\n point_colors = rgb[fg_conf_mask & mask]\n\n # Background points\n bg_homo_grid = np.pad(grid[conf_mask & ~mask], ((0, 0), (0, 1)), constant_values=1)\n bg_local_dirs = np.einsum(\"ij,bj->bi\", np.linalg.inv(K), bg_homo_grid)\n bg_dirs = np.einsum(\"ij,bj->bi\", T_world_camera[:3, :3], bg_local_dirs)\n bg_points = (T_world_camera[:3, 3] + bg_dirs * depth[conf_mask & ~mask, None]).astype(np.float32)\n bg_point_colors = rgb[conf_mask & ~mask]\n\n if bg_downsample_factor > 1 and bg_points.shape[0] > 0:\n indices = np.random.choice(\n bg_points.shape[0],\n size=bg_points.shape[0] // bg_downsample_factor,\n replace=False\n )\n bg_points = bg_points[indices]\n bg_point_colors = bg_point_colors[indices]\n\n return points, point_colors, bg_points, bg_point_colors\n" }, { "path": "viser/src/viser/extras/_urdf.py", "content": "from __future__ import annotations\n\nfrom functools import partial\nfrom pathlib import Path\nfrom typing import List, Tuple\n\nimport numpy as onp\nimport trimesh\nimport yourdfpy\n\nimport viser\n\nfrom .. import transforms as tf\n\n\nclass ViserUrdf:\n \"\"\"Helper for rendering URDFs in Viser.\n\n Args:\n target: ViserServer or ClientHandle object to add URDF to.\n urdf_or_path: Either a path to a URDF file or a yourdfpy URDF object.\n scale: Scale factor to apply to resize the URDF.\n root_node_name: Viser scene tree name for the root of the URDF geometry.\n mesh_color_override: Optional color to override the URDF's mesh colors.\n \"\"\"\n\n def __init__(\n self,\n target: viser.ViserServer | viser.ClientHandle,\n urdf_or_path: yourdfpy.URDF | Path,\n scale: float = 1.0,\n root_node_name: str = \"/\",\n mesh_color_override: tuple[float, float, float] | None = None,\n ) -> None:\n assert root_node_name.startswith(\"/\")\n assert len(root_node_name) == 1 or not root_node_name.endswith(\"/\")\n\n if isinstance(urdf_or_path, Path):\n urdf = yourdfpy.URDF.load(\n urdf_or_path,\n filename_handler=partial(\n yourdfpy.filename_handler_magic, dir=urdf_or_path.parent\n ),\n )\n else:\n urdf = urdf_or_path\n assert isinstance(urdf, yourdfpy.URDF)\n\n self._target = target\n self._urdf = urdf\n self._scale = scale\n self._root_node_name = root_node_name\n\n # Add coordinate frame for each joint.\n self._joint_frames: List[viser.SceneNodeHandle] = []\n for joint in self._urdf.joint_map.values():\n assert isinstance(joint, yourdfpy.Joint)\n self._joint_frames.append(\n self._target.scene.add_frame(\n _viser_name_from_frame(\n self._urdf, joint.child, self._root_node_name\n ),\n show_axes=False,\n )\n )\n\n # Add the URDF's meshes/geometry to viser.\n self._meshes: List[viser.SceneNodeHandle] = []\n for link_name, mesh in urdf.scene.geometry.items():\n assert isinstance(mesh, trimesh.Trimesh)\n T_parent_child = urdf.get_transform(\n link_name, urdf.scene.graph.transforms.parents[link_name]\n )\n name = _viser_name_from_frame(urdf, link_name, root_node_name)\n\n # Scale + transform the mesh. (these will mutate it!)\n #\n # It's important that we use apply_transform() instead of unpacking\n # the rotation/translation terms, since the scene graph transform\n # can also contain scale and reflection terms.\n mesh = mesh.copy()\n mesh.apply_scale(self._scale)\n mesh.apply_transform(T_parent_child)\n\n if mesh_color_override is None:\n self._meshes.append(target.scene.add_mesh_trimesh(name, mesh))\n else:\n self._meshes.append(\n target.scene.add_mesh_simple(\n name,\n mesh.vertices,\n mesh.faces,\n color=mesh_color_override,\n )\n )\n\n def remove(self) -> None:\n \"\"\"Remove URDF from scene.\"\"\"\n # Some of this will be redundant, since children are removed when\n # parents are removed.\n for frame in self._joint_frames:\n frame.remove()\n for mesh in self._meshes:\n mesh.remove()\n\n def update_cfg(self, configuration: onp.ndarray) -> None:\n \"\"\"Update the joint angles of the visualized URDF.\"\"\"\n self._urdf.update_cfg(configuration)\n with self._target.atomic():\n for joint, frame_handle in zip(\n self._urdf.joint_map.values(), self._joint_frames\n ):\n assert isinstance(joint, yourdfpy.Joint)\n T_parent_child = self._urdf.get_transform(joint.child, joint.parent)\n frame_handle.wxyz = tf.SO3.from_matrix(T_parent_child[:3, :3]).wxyz\n frame_handle.position = T_parent_child[:3, 3] * self._scale\n\n def get_actuated_joint_limits(\n self,\n ) -> dict[str, tuple[float | None, float | None]]:\n \"\"\"Returns an ordered mapping from actuated joint names to position limits.\"\"\"\n out: dict[str, tuple[float | None, float | None]] = {}\n for joint_name, joint in zip(\n self._urdf.actuated_joint_names, self._urdf.actuated_joints\n ):\n assert isinstance(joint_name, str)\n assert isinstance(joint, yourdfpy.Joint)\n if joint.limit is None:\n out[joint_name] = (-onp.pi, onp.pi)\n else:\n out[joint_name] = (joint.limit.lower, joint.limit.upper)\n return out\n\n def get_actuated_joint_names(self) -> Tuple[str, ...]:\n \"\"\"Returns a tuple of actuated joint names, in order.\"\"\"\n return tuple(self._urdf.actuated_joint_names)\n\n\ndef _viser_name_from_frame(\n urdf: yourdfpy.URDF,\n frame_name: str,\n root_node_name: str = \"/\",\n) -> str:\n \"\"\"Given the (unique) name of a frame in our URDF's kinematic tree, return a\n scene node name for viser.\n\n For a robot manipulator with four frames, that looks like:\n\n\n ((shoulder)) == ((elbow))\n / / |X|\n / / ((wrist))\n ____/ /____ |X|\n [ ] [=======]\n [ base_link ] [] []\n [___________]\n\n\n this would map a name like \"elbow\" to \"base_link/shoulder/elbow\".\n \"\"\"\n assert root_node_name.startswith(\"/\")\n assert len(root_node_name) == 1 or not root_node_name.endswith(\"/\")\n\n frames = []\n while frame_name != urdf.scene.graph.base_frame:\n frames.append(frame_name)\n frame_name = urdf.scene.graph.transforms.parents[frame_name]\n if root_node_name != \"/\":\n frames.append(root_node_name)\n return \"/\".join(frames[::-1])\n" }, { "path": "viser/src/viser/extras/colmap/__init__.py", "content": "\"\"\"Colmap utilities.\"\"\"\n\nfrom ._colmap_utils import read_cameras_binary as read_cameras_binary\nfrom ._colmap_utils import read_cameras_text as read_cameras_text\nfrom ._colmap_utils import read_images_binary as read_images_binary\nfrom ._colmap_utils import read_images_text as read_images_text\nfrom ._colmap_utils import read_points3d_binary as read_points3d_binary\nfrom ._colmap_utils import read_points3D_text as read_points3D_text\n" }, { "path": "viser/src/viser/extras/colmap/_colmap_utils.py", "content": "# Copyright (c) 2018, 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 inf.ethz.ch)\n#\n# Modified 5/3/2023 to add type annotations.\n\nimport struct\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Dict, Union\n\nimport numpy as np\n\n\n@dataclass(frozen=True)\nclass CameraModel:\n model_id: int\n model_name: str\n num_params: int\n\n\n@dataclass(frozen=True)\nclass Camera:\n id: int\n model: str\n width: int\n height: int\n params: np.ndarray\n\n\n@dataclass(frozen=True)\nclass BaseImage:\n id: int\n qvec: np.ndarray\n tvec: np.ndarray\n camera_id: int\n name: str\n xys: np.ndarray\n point3D_ids: np.ndarray\n\n\n@dataclass(frozen=True)\nclass Point3D:\n id: int\n xyz: np.ndarray\n rgb: np.ndarray\n error: Union[float, np.ndarray]\n image_ids: np.ndarray\n point2D_idxs: np.ndarray\n\n\nclass Image(BaseImage):\n def qvec2rotmat(self):\n return qvec2rotmat(self.qvec)\n\n\nCAMERA_MODELS = {\n CameraModel(model_id=0, model_name=\"SIMPLE_PINHOLE\", num_params=3),\n CameraModel(model_id=1, model_name=\"PINHOLE\", num_params=4),\n CameraModel(model_id=2, model_name=\"SIMPLE_RADIAL\", num_params=4),\n CameraModel(model_id=3, model_name=\"RADIAL\", num_params=5),\n CameraModel(model_id=4, model_name=\"OPENCV\", num_params=8),\n CameraModel(model_id=5, model_name=\"OPENCV_FISHEYE\", num_params=8),\n CameraModel(model_id=6, model_name=\"FULL_OPENCV\", num_params=12),\n CameraModel(model_id=7, model_name=\"FOV\", num_params=5),\n CameraModel(model_id=8, model_name=\"SIMPLE_RADIAL_FISHEYE\", num_params=4),\n CameraModel(model_id=9, model_name=\"RADIAL_FISHEYE\", num_params=5),\n CameraModel(model_id=10, model_name=\"THIN_PRISM_FISHEYE\", num_params=12),\n}\nCAMERA_MODEL_IDS = dict(\n [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS]\n)\n\n\ndef read_next_bytes(fid, num_bytes, format_char_sequence, endian_character=\"<\"):\n \"\"\"Read and unpack the next bytes from a binary file.\n :param fid:\n :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.\n :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.\n :param endian_character: Any of {@, =, <, >, !}\n :return: 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: Union[str, Path]) -> Dict[int, Camera]:\n \"\"\"\n see: src/base/reconstruction.cc\n void Reconstruction::WriteCamerasText(const std::string& path)\n void Reconstruction::ReadCamerasText(const std::string& path)\n \"\"\"\n cameras = {}\n with open(path, \"r\") 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(\n id=camera_id, model=model, width=width, height=height, params=params\n )\n return cameras\n\n\ndef read_cameras_binary(path_to_model_file: Union[str, Path]) -> Dict[int, Camera]:\n \"\"\"\n see: src/base/reconstruction.cc\n void Reconstruction::WriteCamerasBinary(const std::string& path)\n void Reconstruction::ReadCamerasBinary(const std::string& path)\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 camera_line_index in range(num_cameras):\n camera_properties = read_next_bytes(\n fid, num_bytes=24, format_char_sequence=\"iiQQ\"\n )\n camera_id = camera_properties[0]\n model_id = camera_properties[1]\n model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name\n width = camera_properties[2]\n height = camera_properties[3]\n num_params = CAMERA_MODEL_IDS[model_id].num_params\n params = read_next_bytes(\n fid, num_bytes=8 * num_params, format_char_sequence=\"d\" * num_params\n )\n cameras[camera_id] = Camera(\n id=camera_id,\n model=model_name,\n width=width,\n height=height,\n params=np.array(params),\n )\n assert len(cameras) == num_cameras\n return cameras\n\n\ndef read_images_text(path: Union[str, Path]) -> Dict[int, Image]:\n \"\"\"\n see: src/base/reconstruction.cc\n void Reconstruction::ReadImagesText(const std::string& path)\n void Reconstruction::WriteImagesText(const std::string& path)\n \"\"\"\n images = {}\n with open(path, \"r\") 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(\n [tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))]\n )\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: Union[str, Path]) -> Dict[int, Image]:\n \"\"\"\n see: src/base/reconstruction.cc\n void Reconstruction::ReadImagesBinary(const std::string& path)\n void Reconstruction::WriteImagesBinary(const std::string& path)\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 image_index in range(num_reg_images):\n binary_image_properties = read_next_bytes(\n fid, num_bytes=64, format_char_sequence=\"idddddddi\"\n )\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\")[\n 0\n ]\n x_y_id_s = read_next_bytes(\n fid,\n num_bytes=24 * num_points2D,\n format_char_sequence=\"ddq\" * num_points2D,\n )\n xys = np.column_stack(\n [tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))]\n )\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: Union[str, Path]):\n \"\"\"\n see: src/base/reconstruction.cc\n void Reconstruction::ReadPoints3DText(const std::string& path)\n void Reconstruction::WritePoints3DText(const std::string& path)\n \"\"\"\n points3D = {}\n with open(path, \"r\") 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,\n xyz=xyz,\n rgb=rgb,\n error=error,\n image_ids=image_ids,\n point2D_idxs=point2D_idxs,\n )\n return points3D\n\n\ndef read_points3d_binary(path_to_model_file: Union[str, Path]) -> Dict[int, Point3D]:\n \"\"\"\n see: src/base/reconstruction.cc\n void Reconstruction::ReadPoints3DBinary(const std::string& path)\n void Reconstruction::WritePoints3DBinary(const std::string& path)\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 point_line_index in range(num_points):\n binary_point_line_properties = read_next_bytes(\n fid, num_bytes=43, format_char_sequence=\"QdddBBBd\"\n )\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\")[\n 0\n ]\n track_elems = read_next_bytes(\n fid,\n num_bytes=8 * track_length,\n format_char_sequence=\"ii\" * track_length,\n )\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,\n xyz=xyz,\n rgb=rgb,\n error=error,\n image_ids=image_ids,\n point2D_idxs=point2D_idxs,\n )\n return points3D\n\n\ndef qvec2rotmat(qvec):\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" }, { "path": "viser/src/viser/infra/__init__.py", "content": "\"\"\":mod:`viser.infra` provides WebSocket-based communication infrastructure.\n\nWe implement abstractions for:\n- Launching a WebSocket+HTTP server on a shared port.\n- Registering callbacks for connection events and incoming messages.\n- Asynchronous message sending, both broadcasted and to individual clients.\n- Defining dataclass-based message types.\n- Translating Python message types to TypeScript interfaces.\n\nThese are what `viser` runs on under-the-hood, and generally won't be useful unless\nyou're building a web-based application from scratch.\n\"\"\"\n\nfrom ._infra import ClientId as ClientId\nfrom ._infra import WebsockClientConnection as WebsockClientConnection\nfrom ._infra import WebsockMessageHandler as WebsockMessageHandler\nfrom ._infra import WebsockServer as WebsockServer\nfrom ._messages import Message as Message\nfrom ._typescript_interface_gen import (\n TypeScriptAnnotationOverride as TypeScriptAnnotationOverride,\n)\nfrom ._typescript_interface_gen import (\n generate_typescript_interfaces as generate_typescript_interfaces,\n)\n" }, { "path": "viser/src/viser/infra/_async_message_buffer.py", "content": "from __future__ import annotations\n\nimport asyncio\nimport dataclasses\nimport threading\nfrom asyncio.events import AbstractEventLoop\nfrom typing import AsyncGenerator, Dict, List, Sequence\n\nfrom ._messages import Message\n\n\n@dataclasses.dataclass\nclass AsyncMessageBuffer:\n \"\"\"Async iterable for keeping a persistent buffer of messages.\n\n Uses heuristics on message names to automatically cull out redundant messages.\"\"\"\n\n event_loop: AbstractEventLoop\n persistent_messages: bool\n message_event: asyncio.Event = dataclasses.field(default_factory=asyncio.Event)\n flush_event: asyncio.Event = dataclasses.field(default_factory=asyncio.Event)\n\n message_counter: int = 0\n message_from_id: Dict[int, Message] = dataclasses.field(default_factory=dict)\n id_from_redundancy_key: Dict[str, int] = dataclasses.field(default_factory=dict)\n\n buffer_lock: threading.Lock = dataclasses.field(default_factory=threading.Lock)\n \"\"\"Lock to prevent race conditions when pushing messages from different threads.\"\"\"\n\n max_window_size: int = 128\n window_duration_sec: float = 1.0 / 60.0\n done: bool = False\n\n def push(self, message: Message) -> None:\n \"\"\"Push a new message to our buffer, and remove old redundant ones.\"\"\"\n\n assert isinstance(message, Message)\n\n # Add message to buffer.\n redundancy_key = message.redundancy_key()\n with self.buffer_lock:\n new_message_id = self.message_counter\n self.message_from_id[new_message_id] = message\n self.message_counter += 1\n\n # If an existing message with the same key already exists in our buffer, we\n # don't need the old one anymore. :-)\n if (\n redundancy_key is not None\n and redundancy_key in self.id_from_redundancy_key\n ):\n old_message_id = self.id_from_redundancy_key.pop(redundancy_key)\n self.message_from_id.pop(old_message_id)\n self.id_from_redundancy_key[redundancy_key] = new_message_id\n\n # Pulse message event to notify consumers that a new message is available.\n self.event_loop.call_soon_threadsafe(self.message_event.set)\n\n def flush(self) -> None:\n \"\"\"Flush the message buffer; signals to yield a message window immediately.\"\"\"\n self.event_loop.call_soon_threadsafe(self.flush_event.set)\n\n def set_done(self) -> None:\n \"\"\"Set the done flag. Kills the generator.\"\"\"\n self.done = True\n\n # Pulse message event to make sure we aren't waiting for a new message.\n self.event_loop.call_soon_threadsafe(self.message_event.set)\n\n # Pulse flush event to skip any windowing delay.\n self.event_loop.call_soon_threadsafe(self.flush_event.set)\n\n async def window_generator(\n self, client_id: int\n ) -> AsyncGenerator[Sequence[Message], None]:\n \"\"\"Async iterator over messages. Loops infinitely, and waits when no messages\n are available.\"\"\"\n\n last_sent_id = -1\n flush_wait = asyncio.create_task(self.flush_event.wait())\n while not self.done:\n window: List[Message] = []\n most_recent_message_id = self.message_counter - 1\n while (\n last_sent_id < most_recent_message_id\n and len(window) < self.max_window_size\n ):\n last_sent_id += 1\n if self.persistent_messages:\n message = self.message_from_id.get(last_sent_id, None)\n else:\n # If we're not persisting messages, remove them from the buffer.\n with self.buffer_lock:\n message = self.message_from_id.pop(last_sent_id, None)\n if message is not None:\n redundancy_key = message.redundancy_key()\n self.id_from_redundancy_key.pop(redundancy_key, None)\n\n if message is not None and message.excluded_self_client != client_id:\n window.append(message)\n\n if len(window) > 0:\n # Yield a window!\n yield window\n else:\n # Wait for a new message to come in.\n await self.message_event.wait()\n self.message_event.clear()\n\n # Add a delay if either (a) we failed to yield or (b) there's currently no messages to send.\n most_recent_message_id = self.message_counter - 1\n if len(window) == 0 or most_recent_message_id == last_sent_id:\n done, pending = await asyncio.wait(\n [flush_wait], timeout=self.window_duration_sec\n )\n del pending\n if flush_wait in done and not self.done:\n self.flush_event.clear()\n flush_wait = asyncio.create_task(self.flush_event.wait())\n" }, { "path": "viser/src/viser/infra/_infra.py", "content": "from __future__ import annotations\n\nimport abc\nimport asyncio\nimport contextlib\nimport dataclasses\nimport gzip\nimport http\nimport mimetypes\nimport queue\nimport threading\nfrom asyncio.events import AbstractEventLoop\nfrom concurrent.futures import ThreadPoolExecutor\nfrom pathlib import Path\nfrom typing import Any, Callable, Generator, NewType, TypeVar\n\nimport msgspec\nimport rich\nimport websockets.connection\nimport websockets.datastructures\nimport websockets.exceptions\nimport websockets.server\nfrom typing_extensions import Literal, assert_never, override\nfrom websockets.legacy.server import WebSocketServerProtocol\n\nfrom ._async_message_buffer import AsyncMessageBuffer\nfrom ._messages import Message\n\n\n@dataclasses.dataclass\nclass _ClientHandleState:\n # Internal state for ClientConnection objects.\n # message_buffer: asyncio.Queue\n message_buffer: AsyncMessageBuffer\n event_loop: AbstractEventLoop\n\n\nClientId = NewType(\"ClientId\", int)\nTMessage = TypeVar(\"TMessage\", bound=Message)\n\n\nclass RecordHandle:\n \"\"\"**Experimental.**\n\n Handle for recording outgoing messages. Useful for logging + debugging.\"\"\"\n\n def __init__(\n self, handler: WebsockMessageHandler, filter: Callable[[Message], bool]\n ):\n self._handler = handler\n self._filter = filter\n self._loop_start_index: int | None = None\n self._time: float = 0.0\n self._messages: list[tuple[float, dict[str, Any]]] = []\n\n def _insert_message(self, message: Message) -> None:\n \"\"\"Insert a message into the recorded file.\"\"\"\n\n # Exclude GUI messages. This is hacky.\n if not self._filter(message):\n return\n self._messages.append((self._time, message.as_serializable_dict()))\n\n def insert_sleep(self, duration: float) -> None:\n \"\"\"Insert a sleep into the recorded file.\"\"\"\n self._time += duration\n\n def set_loop_start(self) -> None:\n \"\"\"Mark the start of the loop. Messages sent after this point will be\n looped. Should only be called once.\"\"\"\n assert self._loop_start_index is None, \"Loop start already set.\"\n self._loop_start_index = len(self._messages)\n\n def end_and_serialize(self) -> bytes:\n \"\"\"End the recording and serialize contents. Returns the recording as\n bytes, which should generally be written to a file.\"\"\"\n packed_bytes = msgspec.msgpack.encode(\n {\n \"loopStartIndex\": self._loop_start_index,\n \"durationSeconds\": self._time,\n \"messages\": self._messages,\n }\n )\n assert isinstance(packed_bytes, bytes)\n self._handler._record_handle = None\n return gzip.compress(packed_bytes, compresslevel=9)\n\n\nclass WebsockMessageHandler:\n \"\"\"Mix-in for adding message handling to a class.\"\"\"\n\n def __init__(self, thread_executor: ThreadPoolExecutor) -> None:\n self._thread_executor = thread_executor\n self._incoming_handlers: dict[\n type[Message], list[Callable[[ClientId, Message], None]]\n ] = {}\n self._atomic_lock = threading.Lock()\n self._queued_messages: queue.Queue = queue.Queue()\n self._locked_thread_id = -1\n\n # Set to None if not recording.\n self._record_handle: RecordHandle | None = None\n\n def start_recording(self, filter: Callable[[Message], bool]) -> RecordHandle:\n \"\"\"Start recording messages that are sent. Sent messages will be\n serialized and can be used for playback.\"\"\"\n assert self._record_handle is None, \"Already recording.\"\n self._record_handle = RecordHandle(self, filter)\n return self._record_handle\n\n def register_handler(\n self,\n message_cls: type[TMessage],\n callback: Callable[[ClientId, TMessage], Any],\n ) -> None:\n \"\"\"Register a handler for a particular message type.\"\"\"\n if message_cls not in self._incoming_handlers:\n self._incoming_handlers[message_cls] = []\n self._incoming_handlers[message_cls].append(callback) # type: ignore\n\n def unregister_handler(\n self,\n message_cls: type[TMessage],\n callback: Callable[[ClientId, TMessage], Any] | None = None,\n ):\n \"\"\"Unregister a handler for a particular message type.\"\"\"\n assert (\n message_cls in self._incoming_handlers\n ), \"Tried to unregister a handler that hasn't been registered.\"\n if callback is None:\n self._incoming_handlers.pop(message_cls)\n else:\n self._incoming_handlers[message_cls].remove(callback) # type: ignore\n\n def _handle_incoming_message(self, client_id: ClientId, message: Message) -> None:\n \"\"\"Handle incoming messages.\"\"\"\n if type(message) in self._incoming_handlers:\n for cb in self._incoming_handlers[type(message)]:\n cb(client_id, message)\n\n @abc.abstractmethod\n def unsafe_send_message(self, message: Message) -> None: ...\n\n def queue_message(self, message: Message) -> None:\n \"\"\"Wrapped method for sending messages safely.\"\"\"\n if self._record_handle is not None:\n self._record_handle._insert_message(message)\n\n got_lock = self._atomic_lock.acquire(blocking=False)\n if got_lock:\n self.unsafe_send_message(message)\n self._atomic_lock.release()\n else:\n # Send when lock is acquirable, while retaining message order.\n # This could be optimized!\n self._queued_messages.put(message)\n\n def try_again() -> None:\n with self._atomic_lock:\n self.unsafe_send_message(self._queued_messages.get())\n\n self._thread_executor.submit(try_again)\n\n @contextlib.contextmanager\n def atomic(self) -> Generator[None, None, None]:\n \"\"\"Returns a context where: all outgoing messages are grouped and applied by\n clients atomically.\n\n This should be treated as a soft constraint that's helpful for things\n like animations, or when we want position and orientation updates to\n happen synchronously.\n\n Returns:\n Context manager.\n \"\"\"\n # If called multiple times in the same thread, we ignore inner calls.\n thread_id = threading.get_ident()\n if thread_id == self._locked_thread_id:\n got_lock = False\n else:\n self._atomic_lock.acquire()\n self._locked_thread_id = thread_id\n got_lock = True\n\n yield\n\n if got_lock:\n self._atomic_lock.release()\n self._locked_thread_id = -1\n\n\nclass WebsockClientConnection(WebsockMessageHandler):\n \"\"\"Handle for sending messages to and listening to messages from a single\n connected client.\"\"\"\n\n def __init__(\n self,\n client_id: int,\n thread_executor: ThreadPoolExecutor,\n client_state: _ClientHandleState,\n ) -> None:\n self.client_id = client_id\n self._state = client_state\n super().__init__(thread_executor)\n\n @override\n def unsafe_send_message(self, message: Message) -> None:\n \"\"\"Send a message to a specific client.\"\"\"\n self._state.message_buffer.push(message)\n\n\nclass WebsockServer(WebsockMessageHandler):\n \"\"\"Websocket server abstraction. Communicates asynchronously with client\n applications.\n\n By default, all messages are broadcasted to all connected clients.\n\n To send messages to an individual client, we can use `on_client_connect()` to\n retrieve client handles.\n\n Args:\n host: Host to bind server to.\n port: Port to bind server to.\n message_class: Base class for message types. Subclasses of the message type\n should have unique names. This argument is optional currently, but will be\n required in the future.\n http_server_root: Path to root for HTTP server.\n verbose: Toggle for print messages.\n client_api_version: Flag for backwards compatibility. 0 sends individual\n messages. 1 sends windowed messages.\n \"\"\"\n\n def __init__(\n self,\n host: str,\n port: int,\n message_class: type[Message] = Message,\n http_server_root: Path | None = None,\n verbose: bool = True,\n client_api_version: Literal[0, 1] = 0,\n ):\n super().__init__(thread_executor=ThreadPoolExecutor(max_workers=32))\n\n # Track connected clients.\n self._client_connect_cb: list[Callable[[WebsockClientConnection], None]] = []\n self._client_disconnect_cb: list[Callable[[WebsockClientConnection], None]] = []\n\n self._host = host\n self._port = port\n self._message_class = message_class\n self._http_server_root = http_server_root\n self._verbose = verbose\n self._client_api_version: Literal[0, 1] = client_api_version\n self._shutdown_event = threading.Event()\n self._ws_server: websockets.WebSocketServer | None = None\n\n self._client_state_from_id: dict[int, _ClientHandleState] = {}\n\n def start(self) -> None:\n \"\"\"Start the server.\"\"\"\n\n # Start server thread.\n ready_sem = threading.Semaphore(value=1)\n ready_sem.acquire()\n threading.Thread(\n target=lambda: self._background_worker(ready_sem),\n daemon=True,\n ).start()\n\n # Wait for ready signal from the background thread.\n ready_sem.acquire()\n\n # Broadcast buffer should be populated by the background worker.\n assert isinstance(self._broadcast_buffer, AsyncMessageBuffer)\n\n def stop(self) -> None:\n \"\"\"Stop the server.\"\"\"\n assert self._ws_server is not None\n self._ws_server.close()\n self._ws_server = None\n self._thread_executor.shutdown(wait=True)\n\n def on_client_connect(self, cb: Callable[[WebsockClientConnection], Any]) -> None:\n \"\"\"Attach a callback to run for newly connected clients.\"\"\"\n self._client_connect_cb.append(cb)\n\n def on_client_disconnect(\n self, cb: Callable[[WebsockClientConnection], Any]\n ) -> None:\n \"\"\"Attach a callback to run when clients disconnect.\"\"\"\n self._client_disconnect_cb.append(cb)\n\n @override\n def unsafe_send_message(self, message: Message) -> None:\n \"\"\"Pushes a message onto the broadcast queue. Message will be sent to all clients.\n\n Broadcasted messages are persistent: if a new client connects to the server,\n they will receive a buffered set of previously broadcasted messages. The buffer\n is culled using the value of `message.redundancy_key()`.\"\"\"\n self._broadcast_buffer.push(message)\n\n def flush(self) -> None:\n \"\"\"Flush the outgoing message buffer for broadcasted messages. Any buffered\n messages will immediately be sent. (by default they are windowed)\"\"\"\n # TODO: we should add a flush event.\n self._broadcast_buffer.flush()\n\n def flush_client(self, client_id: int) -> None:\n \"\"\"Flush the outgoing message buffer for a particular client. Any buffered\n messages will immediately be sent. (by default they are windowed)\"\"\"\n self._client_state_from_id[client_id].message_buffer.flush()\n\n def _background_worker(self, ready_sem: threading.Semaphore) -> None:\n host = self._host\n port = self._port\n message_class = self._message_class\n http_server_root = self._http_server_root\n\n # Need to make a new event loop for notebook compatbility.\n event_loop = asyncio.new_event_loop()\n asyncio.set_event_loop(event_loop)\n self._broadcast_buffer = AsyncMessageBuffer(\n event_loop, persistent_messages=True\n )\n\n count_lock = asyncio.Lock()\n connection_count = 0\n total_connections = 0\n\n async def serve(websocket: WebSocketServerProtocol) -> None:\n \"\"\"Server loop, run once per connection.\"\"\"\n\n async with count_lock:\n nonlocal connection_count\n client_id = ClientId(connection_count)\n connection_count += 1\n\n nonlocal total_connections\n total_connections += 1\n\n if self._verbose:\n rich.print(\n f\"[bold](viser)[/bold] Connection opened ({client_id},\"\n f\" {total_connections} total),\"\n f\" {len(self._broadcast_buffer.message_from_id)} persistent\"\n \" messages\"\n )\n\n client_state = _ClientHandleState(\n AsyncMessageBuffer(event_loop, persistent_messages=False),\n event_loop,\n )\n client_connection = WebsockClientConnection(\n client_id, self._thread_executor, client_state\n )\n self._client_state_from_id[client_id] = client_state\n\n def handle_incoming(message: Message) -> None:\n self._thread_executor.submit(\n error_print_wrapper(\n lambda: self._handle_incoming_message(client_id, message)\n )\n )\n self._thread_executor.submit(\n error_print_wrapper(\n lambda: client_connection._handle_incoming_message(\n client_id, message\n )\n )\n )\n\n # New connection callbacks.\n for cb in self._client_connect_cb:\n cb(client_connection)\n\n try:\n # For each client: infinite loop over producers (which send messages)\n # and consumers (which receive messages).\n await asyncio.gather(\n _message_producer(\n websocket,\n client_state.message_buffer,\n client_id,\n self._client_api_version,\n ),\n _message_producer(\n websocket,\n self._broadcast_buffer,\n client_id,\n self._client_api_version,\n ),\n _message_consumer(websocket, handle_incoming, message_class),\n )\n except (\n websockets.exceptions.ConnectionClosedOK,\n websockets.exceptions.ConnectionClosedError,\n ):\n # We use a sentinel value to signal that the client producer thread\n # should exit.\n #\n # This is partially cosmetic: it allows us to safely finish pending\n # queue get() tasks, which suppresses a \"Task was destroyed but it is\n # pending\" error.\n client_state.message_buffer.set_done()\n\n # Disconnection callbacks.\n for cb in self._client_disconnect_cb:\n cb(client_connection)\n\n # Cleanup.\n self._client_state_from_id.pop(client_id)\n total_connections -= 1\n if self._verbose:\n rich.print(\n f\"[bold](viser)[/bold] Connection closed ({client_id},\"\n f\" {total_connections} total)\"\n )\n\n # Host client on the same port as the websocket.\n file_cache: dict[Path, bytes] = {}\n file_cache_gzipped: dict[Path, bytes] = {}\n\n async def viser_http_server(\n path: str, request_headers: websockets.datastructures.Headers\n ) -> (\n tuple[http.HTTPStatus, websockets.datastructures.HeadersLike, bytes] | None\n ):\n # Ignore websocket packets.\n if request_headers.get(\"Upgrade\") == \"websocket\":\n return None\n\n # Strip out search params, get relative path.\n path = path.partition(\"?\")[0]\n relpath = str(Path(path).relative_to(\"/\"))\n if relpath == \".\":\n relpath = \"index.html\"\n assert http_server_root is not None\n\n source_path = http_server_root / relpath\n if not source_path.exists():\n return (http.HTTPStatus.NOT_FOUND, {}, b\"404\") # type: ignore\n\n use_gzip = \"gzip\" in request_headers.get(\"Accept-Encoding\", \"\")\n\n mime_type = mimetypes.guess_type(relpath)[0]\n if mime_type is None:\n mime_type = \"application/octet-stream\"\n response_headers = {\n \"Content-Type\": mime_type,\n }\n\n if source_path not in file_cache:\n file_cache[source_path] = source_path.read_bytes()\n if use_gzip:\n response_headers[\"Content-Encoding\"] = \"gzip\"\n if source_path not in file_cache_gzipped:\n file_cache_gzipped[source_path] = gzip.compress(\n file_cache[source_path]\n )\n response_payload = file_cache_gzipped[source_path]\n else:\n response_payload = file_cache[source_path]\n\n # Try to read + send over file.\n return (http.HTTPStatus.OK, response_headers, response_payload)\n\n for _ in range(1000):\n try:\n serve_future = websockets.server.serve(\n serve,\n host,\n port,\n # Compression can be turned off to reduce client-side CPU usage.\n # compression=None,\n process_request=(\n viser_http_server if http_server_root is not None else None\n ),\n )\n self._ws_server = serve_future.ws_server\n event_loop.run_until_complete(serve_future)\n break\n except OSError: # Port not available.\n port += 1\n continue\n\n if self._ws_server is None:\n raise RuntimeError(\"Failed to bind to port!\")\n\n self._port = port\n\n ready_sem.release()\n event_loop.run_forever()\n\n # This will run only when the event loop ends, which happens when the\n # websocket server is closed.\n rich.print(\"[bold](viser)[/bold] Server stopped\")\n\n\nasync def _message_producer(\n websocket: WebSocketServerProtocol,\n buffer: AsyncMessageBuffer,\n client_id: int,\n client_api_version: Literal[0, 1],\n) -> None:\n \"\"\"Infinite loop to broadcast windows of messages from a buffer.\"\"\"\n window_generator = buffer.window_generator(client_id)\n while not buffer.done:\n outgoing = await window_generator.__anext__()\n if client_api_version == 1:\n serialized = msgspec.msgpack.encode(\n tuple(message.as_serializable_dict() for message in outgoing)\n )\n assert isinstance(serialized, bytes)\n await websocket.send(serialized)\n elif client_api_version == 0:\n for msg in outgoing:\n serialized = msgspec.msgpack.encode(msg.as_serializable_dict())\n assert isinstance(serialized, bytes)\n await websocket.send(serialized)\n else:\n assert_never(client_api_version)\n\n\nasync def _message_consumer(\n websocket: WebSocketServerProtocol,\n handle_message: Callable[[Message], None],\n message_class: type[Message],\n) -> None:\n \"\"\"Infinite loop waiting for and then handling incoming messages.\"\"\"\n while True:\n raw = await websocket.recv()\n assert isinstance(raw, bytes)\n message = message_class.deserialize(raw)\n handle_message(message)\n\n\ndef error_print_wrapper(inner: Callable[[], Any]) -> Callable[[], None]:\n \"\"\"Wrap a Callable to print error messages when they happen.\n\n This can be helpful for jobs submitted to ThreadPoolExecutor instances, which, by\n default, will suppress error messages until returned futures are awaited.\n \"\"\"\n\n def wrapped() -> None:\n try:\n inner()\n except Exception as e:\n import traceback as tb\n\n tb.print_exception(type(e), e, e.__traceback__, limit=100)\n\n return wrapped\n" }, { "path": "viser/src/viser/infra/_messages.py", "content": "\"\"\"Message type definitions. For synchronization with the TypeScript definitions, see\n`_typescript_interface_gen.py.`\"\"\"\n\nfrom __future__ import annotations\n\nimport abc\nimport functools\nimport warnings\nfrom typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, TypeVar, cast\n\nimport msgspec\nimport numpy as onp\nfrom typing_extensions import get_args, get_origin, get_type_hints\n\nif TYPE_CHECKING:\n from ._infra import ClientId\nelse:\n ClientId = Any\n\n\ndef _prepare_for_deserialization(value: Any, annotation: Type) -> Any:\n # If annotated as a float but we got an integer, cast to float. These\n # are both `number` in Javascript.\n if annotation is float:\n return float(value)\n elif annotation is int:\n return int(value)\n elif get_origin(annotation) is tuple:\n out = []\n args = get_args(annotation)\n if len(args) >= 2 and args[1] == ...:\n args = (args[0],) * len(value)\n elif len(value) != len(args):\n warnings.warn(f\"[viser] {value} does not match annotation {annotation}\")\n return value\n\n for i, v in enumerate(value):\n out.append(\n # Hack to be OK with wrong type annotations.\n # https://github.com/nerfstudio-project/nerfstudio/pull/1805\n _prepare_for_deserialization(v, args[i]) if i < len(args) else v\n )\n return tuple(out)\n return value\n\n\ndef _prepare_for_serialization(value: Any, annotation: object) -> Any:\n \"\"\"Prepare any special types for serialization.\"\"\"\n if annotation is Any:\n annotation = type(value)\n\n # Coerce some scalar types: if we've annotated as float / int but we get an\n # onp.float32 / onp.int64, for example, we should cast automatically.\n if annotation is float or isinstance(value, onp.floating):\n return float(value)\n if annotation is int or isinstance(value, onp.integer):\n return int(value)\n\n # Recursively handle tuples.\n if isinstance(value, tuple):\n if isinstance(value, onp.ndarray):\n assert False, (\n \"Expected a tuple, but got an array... missing a cast somewhere?\"\n f\" {value}\"\n )\n\n out = []\n if get_origin(annotation) is tuple:\n args = get_args(annotation)\n if len(args) >= 2 and args[1] == ...:\n args = (args[0],) * len(value)\n elif len(value) != len(args):\n warnings.warn(f\"[viser] {value} does not match annotation {annotation}\")\n return value\n else:\n args = [Any] * len(value)\n\n for i, v in enumerate(value):\n out.append(\n # Hack to be OK with wrong type annotations.\n # https://github.com/nerfstudio-project/nerfstudio/pull/1805\n _prepare_for_serialization(v, args[i]) if i < len(args) else v\n )\n return tuple(out)\n\n # For arrays, we serialize underlying data directly. The client is responsible for\n # reading using the correct dtype.\n if isinstance(value, onp.ndarray):\n return value.data if value.data.c_contiguous else value.copy().data\n\n if isinstance(value, dict):\n return {k: _prepare_for_serialization(v, Any) for k, v in value.items()} # type: ignore\n\n return value\n\n\nT = TypeVar(\"T\", bound=\"Message\")\n\n\n@functools.lru_cache(maxsize=None)\ndef get_type_hints_cached(cls: Type[Any]) -> Dict[str, Any]:\n return get_type_hints(cls) # type: ignore\n\n\nclass Message(abc.ABC):\n \"\"\"Base message type for server/client communication.\"\"\"\n\n excluded_self_client: Optional[ClientId] = None\n \"\"\"Don't send this message to a particular client. Useful when a client wants to\n send synchronization information to other clients.\"\"\"\n\n def as_serializable_dict(self) -> Dict[str, Any]:\n \"\"\"Convert a Python Message object into bytes.\"\"\"\n message_type = type(self)\n hints = get_type_hints_cached(message_type)\n out = {\n k: _prepare_for_serialization(v, hints[k]) for k, v in vars(self).items()\n }\n out[\"type\"] = message_type.__name__\n return out\n\n @classmethod\n def _from_serializable_dict(cls, mapping: Dict[str, Any]) -> Dict[str, Any]:\n \"\"\"Convert a dict message back into a Python Message object.\"\"\"\n\n hints = get_type_hints_cached(cls)\n\n mapping = {\n k: _prepare_for_deserialization(v, hints[k]) for k, v in mapping.items()\n }\n return mapping\n\n @classmethod\n def deserialize(cls, message: bytes) -> Message:\n \"\"\"Convert bytes into a Python Message object.\"\"\"\n mapping = msgspec.msgpack.decode(message)\n\n # msgpack deserializes to lists by default, but all of our annotations use\n # tuples.\n def lists_to_tuple(obj: Any) -> Any:\n if isinstance(obj, list):\n return tuple(lists_to_tuple(x) for x in obj)\n elif isinstance(obj, dict):\n return {k: lists_to_tuple(v) for k, v in obj.items()}\n else:\n return obj\n\n mapping = lists_to_tuple(mapping)\n message_type = cls._subclass_from_type_string()[cast(str, mapping.pop(\"type\"))]\n message_kwargs = message_type._from_serializable_dict(mapping)\n return message_type(**message_kwargs)\n\n @classmethod\n @functools.lru_cache(maxsize=100)\n def _subclass_from_type_string(cls: Type[T]) -> Dict[str, Type[T]]:\n subclasses = cls.get_subclasses()\n return {s.__name__: s for s in subclasses}\n\n @classmethod\n def get_subclasses(cls: Type[T]) -> List[Type[T]]:\n \"\"\"Recursively get message subclasses.\"\"\"\n\n def _get_subclasses(typ: Type[T]) -> List[Type[T]]:\n out = []\n for sub in typ.__subclasses__():\n out.append(sub)\n out.extend(_get_subclasses(sub))\n return out\n\n return _get_subclasses(cls)\n\n @abc.abstractmethod\n def redundancy_key(self) -> str:\n \"\"\"Returns a unique key for this message, used for detecting redundant\n messages.\n\n For example: if we send 1000 \"set value\" messages for the same GUI element, we\n should only keep the latest message.\n \"\"\"\n" }, { "path": "viser/src/viser/infra/_typescript_interface_gen.py", "content": "import dataclasses\nfrom collections import defaultdict\nfrom typing import Any, Type, Union, cast\n\nimport numpy as onp\nfrom typing_extensions import (\n Annotated,\n Literal,\n NotRequired,\n get_args,\n get_origin,\n get_type_hints,\n is_typeddict,\n)\n\ntry:\n from typing import Literal as LiteralAlt\nexcept ImportError:\n LiteralAlt = Literal # type: ignore\n\nfrom ._messages import Message\n\n_raw_type_mapping = {\n bool: \"boolean\",\n float: \"number\",\n int: \"number\",\n str: \"string\",\n # For numpy arrays, we directly serialize the underlying data buffer.\n onp.ndarray: \"Uint8Array\",\n bytes: \"Uint8Array\",\n Any: \"any\",\n None: \"null\",\n type(None): \"null\",\n}\n\n\ndef _get_ts_type(typ: Type[Any]) -> str:\n origin_typ = get_origin(typ)\n\n # Look for TypeScriptAnnotationOverride in the annotations.\n if origin_typ is Annotated:\n args = get_args(typ)\n for arg in args[1:]:\n if isinstance(arg, TypeScriptAnnotationOverride):\n return arg.annotation\n\n # If no override is found, just use the unwrapped type.\n origin_typ = args[0]\n\n # Automatic Python => TypeScript conversion.\n if origin_typ is tuple:\n args = get_args(typ)\n if len(args) == 2 and args[1] == ...:\n return _get_ts_type(args[0]) + \"[]\"\n else:\n return \"[\" + \", \".join(map(_get_ts_type, args)) + \"]\"\n elif origin_typ is list:\n args = get_args(typ)\n assert len(args) == 1\n return _get_ts_type(args[0]) + \"[]\"\n elif origin_typ in (Literal, LiteralAlt):\n return \" | \".join(\n map(\n lambda lit: repr(lit).lower() if type(lit) is bool else repr(lit),\n get_args(typ),\n )\n )\n elif origin_typ is Union:\n return (\n \"(\"\n + \" | \".join(\n map(\n _get_ts_type,\n get_args(typ),\n )\n )\n + \")\"\n )\n elif origin_typ is list:\n args = get_args(typ)\n return _get_ts_type(args[0]) + \"[]\"\n elif origin_typ is dict:\n args = get_args(typ)\n assert len(args) == 2\n return \"{ [key: \" + _get_ts_type(args[0]) + \"]: \" + _get_ts_type(args[1]) + \" }\"\n elif is_typeddict(typ):\n hints = get_type_hints(typ)\n optional_keys = getattr(typ, \"__optional_keys__\", [])\n\n def fmt(key):\n val = hints[key]\n optional = key in optional_keys\n if get_origin(val) is NotRequired:\n val = get_args(val)[0]\n ret = f\"'{key}'{'?' if optional else ''}\" + \": \" + _get_ts_type(val)\n return ret\n\n ret = \"{\" + \", \".join(map(fmt, hints)) + \"}\"\n return ret\n else:\n # Like get_origin(), but also supports numpy.typing.NDArray[dtype].\n typ = cast(Any, getattr(typ, \"__origin__\", typ))\n\n assert typ in _raw_type_mapping, f\"Unsupported type {typ}\"\n return _raw_type_mapping[typ]\n\n\n@dataclasses.dataclass(frozen=True)\nclass TypeScriptAnnotationOverride:\n \"\"\"Use with `typing.Annotated[]` to override the automatically-generated\n TypeScript annotation corresponding to a dataclass field.\"\"\"\n\n annotation: str\n\n\ndef generate_typescript_interfaces(message_cls: Type[Message]) -> str:\n \"\"\"Generate TypeScript definitions for all subclasses of a base message class.\"\"\"\n out_lines = []\n message_types = message_cls.get_subclasses()\n tag_map = defaultdict(list)\n\n # Generate interfaces for each specific message.\n for cls in message_types:\n if cls.__doc__ is not None:\n docstring = \"\\n * \".join(\n map(lambda line: line.strip(), cls.__doc__.split(\"\\n\"))\n )\n out_lines.append(f\"/** {docstring}\")\n out_lines.append(\" *\")\n out_lines.append(\" * (automatically generated)\")\n out_lines.append(\" */\")\n\n for tag in getattr(cls, \"_tags\", []):\n tag_map[tag].append(cls.__name__)\n\n out_lines.append(f\"export interface {cls.__name__} \" + \"{\")\n out_lines.append(f' type: \"{cls.__name__}\";')\n field_names = set([f.name for f in dataclasses.fields(cls)]) # type: ignore\n for name, typ in get_type_hints(cls, include_extras=True).items():\n if name in field_names:\n typ = _get_ts_type(typ)\n else:\n continue\n out_lines.append(f\" {name}: {typ};\")\n out_lines.append(\"}\")\n out_lines.append(\"\")\n\n # Generate union type over all messages.\n out_lines.append(\"export type Message = \")\n for cls in message_types:\n out_lines.append(f\" | {cls.__name__}\")\n out_lines[-1] = out_lines[-1] + \";\"\n\n # Generate union type over all tags.\n for tag, cls_names in tag_map.items():\n out_lines.append(f\"export type {tag} = \")\n for cls_name in cls_names:\n out_lines.append(f\" | {cls_name}\")\n out_lines[-1] = out_lines[-1] + \";\"\n\n interfaces = \"\\n\".join(out_lines) + \"\\n\"\n\n # Add header and return.\n return (\n \"\\n\".join(\n [\n (\n \"// AUTOMATICALLY GENERATED message interfaces, from Python\"\n \" dataclass definitions.\"\n ),\n \"// This file should not be manually modified.\",\n \"\",\n ]\n )\n + interfaces\n )\n" }, { "path": "viser/src/viser/py.typed", "content": "" }, { "path": "viser/src/viser/scripts/__init__.py", "content": "" }, { "path": "viser/src/viser/scripts/dev_checks.py", "content": "#!/usr/bin/env python\n\"\"\"Runs formatting, linting, and type checking tests.\"\"\"\n\nimport subprocess\nimport sys\n\nimport tyro\nfrom rich import console\nfrom rich.style import Style\n\nCONSOLE = console.Console()\n\nTYPE_TESTS = [\"pyright .\", \"mypy .\"]\nFORMAT_TESTS = [\"ruff check --fix .\", \"ruff format .\"]\n\n\ndef run_command(command: str, continue_on_fail: bool = False) -> bool:\n \"\"\"Run a command kill actions if it fails\n\n Args:\n command: Command to run.\n continue_on_fail: Whether to continue running commands if the current one fails..\n \"\"\"\n ret_code = subprocess.call(command, shell=True)\n if ret_code != 0:\n CONSOLE.print(f\"[bold red]Error: `{command}` failed.\")\n if not continue_on_fail:\n sys.exit(1)\n return ret_code == 0\n\n\ndef run_code_checks(\n continue_on_fail: bool = False,\n skip_format_checks: bool = False,\n skip_type_checks: bool = False,\n):\n \"\"\"Runs formatting, linting, and type checking tests.\n\n Args:\n continue_on_fail: Whether or not to continue running actions commands if the current one fails.\n skip_format_checks: Whether or not to skip format tests.\n skip_type_checks: Whether or not to skip type tests.\n \"\"\"\n\n success = True\n\n assert (\n not skip_format_checks or not skip_type_checks\n ), \"Cannot skip format and type tests at the same time.\"\n tests = []\n if not skip_format_checks:\n tests += FORMAT_TESTS\n if not skip_type_checks:\n tests += TYPE_TESTS\n\n for test in tests:\n CONSOLE.line()\n CONSOLE.rule(f\"[bold green]Running: {test}\")\n success = success and run_command(test, continue_on_fail=continue_on_fail)\n\n if success:\n CONSOLE.line()\n CONSOLE.rule(characters=\"=\")\n CONSOLE.print(\n \"[bold green]:TADA: :TADA: :TADA: ALL CHECKS PASSED :TADA: :TADA: :TADA:\",\n justify=\"center\",\n )\n CONSOLE.rule(characters=\"=\")\n else:\n CONSOLE.line()\n CONSOLE.rule(characters=\"=\", style=Style(color=\"red\"))\n CONSOLE.print(\n \"[bold red]:skull: :skull: :skull: ERRORS FOUND :skull: :skull: :skull:\",\n justify=\"center\",\n )\n CONSOLE.rule(characters=\"=\", style=Style(color=\"red\"))\n\n\ndef entrypoint():\n \"\"\"Entrypoint for use with pyproject scripts.\"\"\"\n tyro.cli(run_code_checks)\n\n\nif __name__ == \"__main__\":\n entrypoint()\n" }, { "path": "viser/src/viser/theme/__init__.py", "content": "\"\"\":mod:`viser.theme` provides interfaces for themeing the viser\nfrontend from within Python.\n\"\"\"\n\nfrom ._titlebar import TitlebarButton as TitlebarButton\nfrom ._titlebar import TitlebarConfig as TitlebarConfig\nfrom ._titlebar import TitlebarImage as TitlebarImage\n" }, { "path": "viser/src/viser/theme/_titlebar.py", "content": "from typing import Literal, Optional, Tuple, TypedDict\n\n\nclass TitlebarButton(TypedDict):\n \"\"\"A link-only button that appears in the Titlebar.\"\"\"\n\n text: Optional[str]\n icon: Optional[Literal[\"GitHub\", \"Description\", \"Keyboard\"]]\n href: Optional[str]\n\n\nclass TitlebarImage(TypedDict):\n \"\"\"An image that appears on the titlebar.\"\"\"\n\n image_url_light: str\n image_url_dark: Optional[str]\n image_alt: str\n href: Optional[str]\n\n\nclass TitlebarConfig(TypedDict):\n \"\"\"Configure the content that appears in the titlebar.\"\"\"\n\n buttons: Optional[Tuple[TitlebarButton, ...]]\n image: Optional[TitlebarImage]\n" }, { "path": "viser/src/viser/transforms/__init__.py", "content": "\"\"\"Lie group interface for rigid transforms, ported from\n[jaxlie](https://github.com/brentyi/jaxlie). Used by `viser` internally and\nin examples.\n\nImplements SO(2), SO(3), SE(2), and SE(3) Lie groups. Rotations are parameterized\nvia S^1 and S^3.\n\"\"\"\n\nfrom ._base import MatrixLieGroup as MatrixLieGroup\nfrom ._base import SEBase as SEBase\nfrom ._base import SOBase as SOBase\nfrom ._se2 import SE2 as SE2\nfrom ._se3 import SE3 as SE3\nfrom ._so2 import SO2 as SO2\nfrom ._so3 import SO3 as SO3\n" }, { "path": "viser/src/viser/transforms/_base.py", "content": "import abc\nfrom typing import ClassVar, Generic, Tuple, TypeVar, Union, overload\n\nimport numpy as onp\nimport numpy.typing as onpt\nfrom typing_extensions import Self, final, get_args, override\n\n\nclass MatrixLieGroup(abc.ABC):\n \"\"\"Interface definition for matrix Lie groups.\"\"\"\n\n # Class properties.\n # > These will be set in `_utils.register_lie_group()`.\n\n matrix_dim: ClassVar[int]\n \"\"\"Dimension of square matrix output from `.as_matrix()`.\"\"\"\n\n parameters_dim: ClassVar[int]\n \"\"\"Dimension of underlying parameters, `.parameters()`.\"\"\"\n\n tangent_dim: ClassVar[int]\n \"\"\"Dimension of tangent space.\"\"\"\n\n space_dim: ClassVar[int]\n \"\"\"Dimension of coordinates that can be transformed.\"\"\"\n\n def __init__(\n # Notes:\n # - For the constructor signature to be consistent with subclasses, `parameters`\n # should be marked as positional-only. But this isn't possible in Python 3.7.\n # - This method is implicitly overriden by the dataclass decorator and\n # should _not_ be marked abstract.\n self,\n parameters: onp.ndarray,\n ):\n \"\"\"Construct a group object from its underlying parameters.\"\"\"\n raise NotImplementedError()\n\n # Shared implementations.\n\n @overload\n def __matmul__(self, other: Self) -> Self: ...\n\n @overload\n def __matmul__(\n self, other: onpt.NDArray[onp.floating]\n ) -> onpt.NDArray[onp.floating]: ...\n\n def __matmul__(\n self, other: Union[Self, onpt.NDArray[onp.floating]]\n ) -> Union[Self, onpt.NDArray[onp.floating]]:\n \"\"\"Overload for the `@` operator.\n\n Switches between the group action (`.apply()`) and multiplication\n (`.multiply()`) based on the type of `other`.\n \"\"\"\n if isinstance(other, onp.ndarray):\n return self.apply(target=other)\n elif isinstance(other, MatrixLieGroup):\n assert self.space_dim == other.space_dim\n return self.multiply(other=other)\n else:\n assert False, f\"Invalid argument type for `@` operator: {type(other)}\"\n\n # Factory.\n\n @classmethod\n @abc.abstractmethod\n def identity(cls, batch_axes: Tuple[int, ...] = ()) -> Self:\n \"\"\"Returns identity element.\n\n Args:\n batch_axes: Any leading batch axes for the output transform.\n\n Returns:\n Identity element.\n \"\"\"\n\n @classmethod\n @abc.abstractmethod\n def from_matrix(cls, matrix: onpt.NDArray[onp.floating]) -> Self:\n \"\"\"Get group member from matrix representation.\n\n Args:\n matrix: Matrix representaiton.\n\n Returns:\n Group member.\n \"\"\"\n\n # Accessors.\n\n @abc.abstractmethod\n def as_matrix(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Get transformation as a matrix. Homogeneous for SE groups.\"\"\"\n\n @abc.abstractmethod\n def parameters(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Get underlying representation.\"\"\"\n\n # Operations.\n\n @abc.abstractmethod\n def apply(self, target: onpt.NDArray[onp.floating]) -> onpt.NDArray[onp.floating]:\n \"\"\"Applies group action to a point.\n\n Args:\n target: Point to transform.\n\n Returns:\n Transformed point.\n \"\"\"\n\n @abc.abstractmethod\n def multiply(self, other: Self) -> Self:\n \"\"\"Composes this transformation with another.\n\n Returns:\n self @ other\n \"\"\"\n\n @classmethod\n @abc.abstractmethod\n def exp(cls, tangent: onpt.NDArray[onp.floating]) -> Self:\n \"\"\"Computes `expm(wedge(tangent))`.\n\n Args:\n tangent: Tangent vector to take the exponential of.\n\n Returns:\n Output.\n \"\"\"\n\n @abc.abstractmethod\n def log(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Computes `vee(logm(transformation matrix))`.\n\n Returns:\n Output. Shape should be `(tangent_dim,)`.\n \"\"\"\n\n @abc.abstractmethod\n def adjoint(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Computes the adjoint, which transforms tangent vectors between tangent\n spaces.\n\n More precisely, for a transform `GroupType`:\n ```\n GroupType @ exp(omega) = exp(Adj_T @ omega) @ GroupType\n ```\n\n In robotics, typically used for transforming twists, wrenches, and Jacobians\n across different reference frames.\n\n Returns:\n Output. Shape should be `(tangent_dim, tangent_dim)`.\n \"\"\"\n\n @abc.abstractmethod\n def inverse(self) -> Self:\n \"\"\"Computes the inverse of our transform.\n\n Returns:\n Output.\n \"\"\"\n\n @abc.abstractmethod\n def normalize(self) -> Self:\n \"\"\"Normalize/projects values and returns.\n\n Returns:\n Normalized group member.\n \"\"\"\n\n # @classmethod\n # @abc.abstractmethod\n # def sample_uniform(cls, key: onp.ndarray, batch_axes: Tuple[int, ...] = ()) -> Self:\n # \"\"\"Draw a uniform sample from the group. Translations (if applicable) are in the\n # range [-1, 1].\n #\n # Args:\n # key: PRNG key, as returned by `jax.random.PRNGKey()`.\n # batch_axes: Any leading batch axes for the output transforms. Each\n # sampled transform will be different.\n #\n # Returns:\n # Sampled group member.\n # \"\"\"\n\n @final\n def get_batch_axes(self) -> Tuple[int, ...]:\n \"\"\"Return any leading batch axes in contained parameters. If an array of shape\n `(100, 4)` is placed in the wxyz field of an SO3 object, for example, this will\n return `(100,)`.\"\"\"\n return self.parameters().shape[:-1]\n\n\nclass SOBase(MatrixLieGroup):\n \"\"\"Base class for special orthogonal groups.\"\"\"\n\n\nContainedSOType = TypeVar(\"ContainedSOType\", bound=SOBase)\n\n\nclass SEBase(Generic[ContainedSOType], MatrixLieGroup):\n \"\"\"Base class for special Euclidean groups.\n\n Each SE(N) group member contains an SO(N) rotation, as well as an N-dimensional\n translation vector.\n \"\"\"\n\n # SE-specific interface.\n\n @classmethod\n @abc.abstractmethod\n def from_rotation_and_translation(\n cls,\n rotation: ContainedSOType,\n translation: onpt.NDArray[onp.floating],\n ) -> Self:\n \"\"\"Construct a rigid transform from a rotation and a translation.\n\n Args:\n rotation: Rotation term.\n translation: translation term.\n\n Returns:\n Constructed transformation.\n \"\"\"\n\n @final\n @classmethod\n def from_rotation(cls, rotation: ContainedSOType) -> Self:\n return cls.from_rotation_and_translation(\n rotation=rotation,\n translation=onp.zeros(\n (*rotation.get_batch_axes(), cls.space_dim),\n dtype=rotation.parameters().dtype,\n ),\n )\n\n @final\n @classmethod\n def from_translation(cls, translation: onpt.NDArray[onp.floating]) -> Self:\n # Extract rotation class from type parameter.\n assert len(cls.__orig_bases__) == 1 # type: ignore\n return cls.from_rotation_and_translation(\n rotation=get_args(cls.__orig_bases__[0])[0].identity(), # type: ignore\n translation=translation,\n )\n\n @abc.abstractmethod\n def rotation(self) -> ContainedSOType:\n \"\"\"Returns a transform's rotation term.\"\"\"\n\n @abc.abstractmethod\n def translation(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Returns a transform's translation term.\"\"\"\n\n # Overrides.\n\n @final\n @override\n def apply(self, target: onpt.NDArray[onp.floating]) -> onpt.NDArray[onp.floating]:\n return self.rotation() @ target + self.translation() # type: ignore\n\n @final\n @override\n def multiply(self, other: Self) -> Self:\n return type(self).from_rotation_and_translation(\n rotation=self.rotation() @ other.rotation(),\n translation=(self.rotation() @ other.translation()) + self.translation(),\n )\n\n @final\n @override\n def inverse(self) -> Self:\n R_inv = self.rotation().inverse()\n return type(self).from_rotation_and_translation(\n rotation=R_inv,\n translation=-(R_inv @ self.translation()),\n )\n\n @final\n @override\n def normalize(self) -> Self:\n return type(self).from_rotation_and_translation(\n rotation=self.rotation().normalize(),\n translation=self.translation(),\n )\n" }, { "path": "viser/src/viser/transforms/_se2.py", "content": "import dataclasses\nfrom typing import Tuple, cast\n\nimport numpy as onp\nimport numpy.typing as onpt\nfrom typing_extensions import override\n\nfrom . import _base, hints\nfrom ._so2 import SO2\nfrom .utils import broadcast_leading_axes, get_epsilon, register_lie_group\n\n\n@register_lie_group(\n matrix_dim=3,\n parameters_dim=4,\n tangent_dim=3,\n space_dim=2,\n)\n@dataclasses.dataclass(frozen=True)\nclass SE2(_base.SEBase[SO2]):\n \"\"\"Special Euclidean group for proper rigid transforms in 2D. Broadcasting\n rules are the same as for numpy.\n\n Ported to numpy from `jaxlie.SE2`.\n\n Internal parameterization is `(cos, sin, x, y)`. Tangent parameterization is `(vx,\n vy, omega)`.\n \"\"\"\n\n # SE2-specific.\n\n unit_complex_xy: onpt.NDArray[onp.floating]\n \"\"\"Internal parameters. `(cos, sin, x, y)`. Shape should be `(*, 4)`.\"\"\"\n\n @override\n def __repr__(self) -> str:\n unit_complex = onp.round(self.unit_complex_xy[..., :2], 5)\n xy = onp.round(self.unit_complex_xy[..., 2:], 5)\n return f\"{self.__class__.__name__}(unit_complex={unit_complex}, xy={xy})\"\n\n @staticmethod\n def from_xy_theta(x: hints.Scalar, y: hints.Scalar, theta: hints.Scalar) -> \"SE2\":\n \"\"\"Construct a transformation from standard 2D pose parameters.\n\n This is not the same as integrating over a length-3 twist.\n \"\"\"\n cos = onp.cos(theta)\n sin = onp.sin(theta)\n return SE2(unit_complex_xy=onp.stack([cos, sin, x, y], axis=-1))\n\n # SE-specific.\n\n @classmethod\n @override\n def from_rotation_and_translation(\n cls,\n rotation: SO2,\n translation: onpt.NDArray[onp.floating],\n ) -> \"SE2\":\n assert translation.shape[-1:] == (2,)\n rotation, translation = broadcast_leading_axes((rotation, translation))\n return SE2(\n unit_complex_xy=onp.concatenate(\n [rotation.unit_complex, translation], axis=-1\n )\n )\n\n @override\n def rotation(self) -> SO2:\n return SO2(unit_complex=self.unit_complex_xy[..., :2])\n\n @override\n def translation(self) -> onpt.NDArray[onp.floating]:\n return self.unit_complex_xy[..., 2:]\n\n # Factory.\n\n @classmethod\n @override\n def identity(cls, batch_axes: Tuple[int, ...] = ()) -> \"SE2\":\n return SE2(\n unit_complex_xy=onp.broadcast_to(\n onp.array([1.0, 0.0, 0.0, 0.0]), (*batch_axes, 4)\n )\n )\n\n @classmethod\n @override\n def from_matrix(cls, matrix: onpt.NDArray[onp.floating]) -> \"SE2\":\n assert matrix.shape[-2:] == (3, 3) or matrix.shape[-2:] == (2, 3)\n # Currently assumes bottom row is [0, 0, 1].\n return SE2.from_rotation_and_translation(\n rotation=SO2.from_matrix(matrix[..., :2, :2]),\n translation=matrix[..., :2, 2],\n )\n\n # Accessors.\n\n @override\n def parameters(self) -> onpt.NDArray[onp.floating]:\n return self.unit_complex_xy\n\n @override\n def as_matrix(self) -> onpt.NDArray[onp.floating]:\n cos, sin, x, y = onp.moveaxis(self.unit_complex_xy, -1, 0)\n out = onp.stack(\n [\n cos,\n -sin,\n x,\n sin,\n cos,\n y,\n onp.zeros_like(x),\n onp.zeros_like(x),\n onp.ones_like(x),\n ],\n axis=-1,\n ).reshape((*self.get_batch_axes(), 3, 3))\n return out\n\n # Operations.\n\n @classmethod\n @override\n def exp(cls, tangent: onpt.NDArray[onp.floating]) -> \"SE2\":\n # Reference:\n # > https://github.com/strasdat/Sophus/blob/a0fe89a323e20c42d3cecb590937eb7a06b8343a/sophus/se2.hpp#L558\n # Also see:\n # > http://ethaneade.com/lie.pdf\n\n assert tangent.shape[-1:] == (3,)\n\n theta = tangent[..., 2]\n use_taylor = onp.abs(theta) < get_epsilon(tangent.dtype)\n\n # Shim to avoid NaNs in onp.where branches, which cause failures for\n # reverse-mode AD in JAX. This isn't needed for vanilla numpy.\n safe_theta = cast(\n onp.ndarray,\n onp.where(\n use_taylor,\n onp.ones_like(theta), # Any non-zero value should do here.\n theta,\n ),\n )\n\n theta_sq = theta**2\n sin_over_theta = cast(\n onp.ndarray,\n onp.where(\n use_taylor,\n 1.0 - theta_sq / 6.0,\n onp.sin(safe_theta) / safe_theta,\n ),\n )\n one_minus_cos_over_theta = cast(\n onp.ndarray,\n onp.where(\n use_taylor,\n 0.5 * theta - theta * theta_sq / 24.0,\n (1.0 - onp.cos(safe_theta)) / safe_theta,\n ),\n )\n\n V = onp.stack(\n [\n sin_over_theta,\n -one_minus_cos_over_theta,\n one_minus_cos_over_theta,\n sin_over_theta,\n ],\n axis=-1,\n ).reshape((*tangent.shape[:-1], 2, 2))\n return SE2.from_rotation_and_translation(\n rotation=SO2.from_radians(theta),\n translation=onp.einsum(\"...ij,...j->...i\", V, tangent[..., :2]),\n )\n\n @override\n def log(self) -> onpt.NDArray[onp.floating]:\n # Reference:\n # > https://github.com/strasdat/Sophus/blob/a0fe89a323e20c42d3cecb590937eb7a06b8343a/sophus/se2.hpp#L160\n # Also see:\n # > http://ethaneade.com/lie.pdf\n\n theta = self.rotation().log()[..., 0]\n\n cos = onp.cos(theta)\n cos_minus_one = cos - 1.0\n half_theta = theta / 2.0\n use_taylor = onp.abs(cos_minus_one) < get_epsilon(theta.dtype)\n\n # Shim to avoid NaNs in onp.where branches, which cause failures for\n # reverse-mode AD in JAX. This isn't needed for vanilla numpy.\n safe_cos_minus_one = onp.where(\n use_taylor,\n onp.ones_like(cos_minus_one), # Any non-zero value should do here.\n cos_minus_one,\n )\n\n half_theta_over_tan_half_theta = onp.where(\n use_taylor,\n # Taylor approximation.\n 1.0 - theta**2 / 12.0,\n # Default.\n -(half_theta * onp.sin(theta)) / safe_cos_minus_one,\n )\n\n V_inv = onp.stack(\n [\n half_theta_over_tan_half_theta,\n half_theta,\n -half_theta,\n half_theta_over_tan_half_theta,\n ],\n axis=-1,\n ).reshape((*theta.shape, 2, 2))\n\n tangent = onp.concatenate(\n [\n onp.einsum(\"...ij,...j->...i\", V_inv, self.translation()),\n theta[..., None],\n ],\n axis=-1,\n )\n return tangent\n\n @override\n def adjoint(self: \"SE2\") -> onpt.NDArray[onp.floating]:\n cos, sin, x, y = onp.moveaxis(self.unit_complex_xy, -1, 0)\n return onp.stack(\n [\n cos,\n -sin,\n y,\n sin,\n cos,\n -x,\n onp.zeros_like(x),\n onp.zeros_like(x),\n onp.ones_like(x),\n ],\n axis=-1,\n ).reshape((*self.get_batch_axes(), 3, 3))\n\n # @classmethod\n # @override\n # def sample_uniform(\n # cls, key: onp.ndarray, batch_axes: jdc.Static[Tuple[int, ...]] = ()\n # ) -> \"SE2\":\n # key0, key1 = jax.random.split(key)\n # return SE2.from_rotation_and_translation(\n # rotation=SO2.sample_uniform(key0, batch_axes=batch_axes),\n # translation=jax.random.uniform(\n # key=key1,\n # shape=(\n # *batch_axes,\n # 2,\n # ),\n # minval=-1.0,\n # maxval=1.0,\n # ),\n # )\n" }, { "path": "viser/src/viser/transforms/_se3.py", "content": "from __future__ import annotations\n\nimport dataclasses\nfrom typing import Tuple, cast\n\nimport numpy as onp\nimport numpy.typing as onpt\nfrom typing_extensions import override\n\nfrom . import _base\nfrom ._so3 import SO3\nfrom .utils import broadcast_leading_axes, get_epsilon, register_lie_group\n\n\ndef _skew(omega: onpt.NDArray[onp.floating]) -> onpt.NDArray[onp.floating]:\n \"\"\"Returns the skew-symmetric form of a length-3 vector.\"\"\"\n\n wx, wy, wz = onp.moveaxis(omega, -1, 0)\n zeros = onp.zeros_like(wx)\n return onp.stack(\n [zeros, -wz, wy, wz, zeros, -wx, -wy, wx, zeros],\n axis=-1,\n ).reshape((*omega.shape[:-1], 3, 3))\n\n\n@register_lie_group(\n matrix_dim=4,\n parameters_dim=7,\n tangent_dim=6,\n space_dim=3,\n)\n@dataclasses.dataclass(frozen=True)\nclass SE3(_base.SEBase[SO3]):\n \"\"\"Special Euclidean group for proper rigid transforms in 3D. Broadcasting\n rules are the same as for numpy.\n\n Ported to numpy from `jaxlie.SE3`.\n\n Internal parameterization is `(qw, qx, qy, qz, x, y, z)`. Tangent parameterization\n is `(vx, vy, vz, omega_x, omega_y, omega_z)`.\n \"\"\"\n\n # SE3-specific.\n\n wxyz_xyz: onpt.NDArray[onp.floating]\n \"\"\"Internal parameters. wxyz quaternion followed by xyz translation. Shape should be `(*, 7)`.\"\"\"\n\n @override\n def __repr__(self) -> str:\n quat = onp.round(self.wxyz_xyz[..., :4], 5)\n trans = onp.round(self.wxyz_xyz[..., 4:], 5)\n return f\"{self.__class__.__name__}(wxyz={quat}, xyz={trans})\"\n\n # SE-specific.\n\n @classmethod\n @override\n def from_rotation_and_translation(\n cls,\n rotation: SO3,\n translation: onpt.NDArray[onp.floating],\n ) -> SE3:\n assert translation.shape[-1:] == (3,)\n rotation, translation = broadcast_leading_axes((rotation, translation))\n return SE3(wxyz_xyz=onp.concatenate([rotation.wxyz, translation], axis=-1))\n\n @override\n def rotation(self) -> SO3:\n return SO3(wxyz=self.wxyz_xyz[..., :4])\n\n @override\n def translation(self) -> onpt.NDArray[onp.floating]:\n return self.wxyz_xyz[..., 4:]\n\n # Factory.\n\n @classmethod\n @override\n def identity(cls, batch_axes: Tuple[int, ...] = ()) -> SE3:\n return SE3(\n wxyz_xyz=onp.broadcast_to(\n onp.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), (*batch_axes, 7)\n )\n )\n\n @classmethod\n @override\n def from_matrix(cls, matrix: onpt.NDArray[onp.floating]) -> SE3:\n assert matrix.shape[-2:] == (4, 4) or matrix.shape[-2:] == (3, 4)\n # Currently assumes bottom row is [0, 0, 0, 1].\n return SE3.from_rotation_and_translation(\n rotation=SO3.from_matrix(matrix[..., :3, :3]),\n translation=matrix[..., :3, 3],\n )\n\n # Accessors.\n\n @override\n def as_matrix(self) -> onpt.NDArray[onp.floating]:\n out = onp.zeros((*self.get_batch_axes(), 4, 4))\n out[..., :3, :3] = self.rotation().as_matrix()\n out[..., :3, 3] = self.translation()\n out[..., 3, 3] = 1.0\n return out\n\n @override\n def parameters(self) -> onpt.NDArray[onp.floating]:\n return self.wxyz_xyz\n\n # Operations.\n\n @classmethod\n @override\n def exp(cls, tangent: onpt.NDArray[onp.floating]) -> SE3:\n # Reference:\n # > https://github.com/strasdat/Sophus/blob/a0fe89a323e20c42d3cecb590937eb7a06b8343a/sophus/se3.hpp#L761\n\n # (x, y, z, omega_x, omega_y, omega_z)\n assert tangent.shape[-1:] == (6,)\n\n rotation = SO3.exp(tangent[..., 3:])\n\n theta_squared = onp.sum(onp.square(tangent[..., 3:]), axis=-1)\n use_taylor = theta_squared < get_epsilon(theta_squared.dtype)\n\n # Shim to avoid NaNs in onp.where branches, which cause failures for\n # reverse-mode AD in JAX. This isn't needed for vanilla numpy.\n theta_squared_safe = cast(\n onp.ndarray,\n onp.where(\n use_taylor,\n onp.ones_like(theta_squared), # Any non-zero value should do here.\n theta_squared,\n ),\n )\n del theta_squared\n theta_safe = onp.sqrt(theta_squared_safe)\n\n skew_omega = _skew(tangent[..., 3:])\n V = onp.where(\n use_taylor[..., None, None],\n rotation.as_matrix(),\n (\n onp.eye(3)\n + ((1.0 - onp.cos(theta_safe)) / (theta_squared_safe))[..., None, None]\n * skew_omega\n + (\n (theta_safe - onp.sin(theta_safe))\n / (theta_squared_safe * theta_safe)\n )[..., None, None]\n * onp.einsum(\"...ij,...jk->...ik\", skew_omega, skew_omega)\n ),\n )\n\n return SE3.from_rotation_and_translation(\n rotation=rotation,\n translation=onp.einsum(\"...ij,...j->...i\", V, tangent[..., :3]),\n )\n\n @override\n def log(self) -> onpt.NDArray[onp.floating]:\n # Reference:\n # > https://github.com/strasdat/Sophus/blob/a0fe89a323e20c42d3cecb590937eb7a06b8343a/sophus/se3.hpp#L223\n omega = self.rotation().log()\n theta_squared = onp.sum(onp.square(omega), axis=-1)\n use_taylor = theta_squared < get_epsilon(theta_squared.dtype)\n\n skew_omega = _skew(omega)\n\n # Shim to avoid NaNs in onp.where branches, which cause failures for\n # reverse-mode AD in JAX. This isn't needed for vanilla numpy.\n theta_squared_safe = onp.where(\n use_taylor,\n onp.ones_like(theta_squared), # Any non-zero value should do here.\n theta_squared,\n )\n del theta_squared\n theta_safe = onp.sqrt(theta_squared_safe)\n half_theta_safe = theta_safe / 2.0\n\n V_inv = onp.where(\n use_taylor[..., None, None],\n onp.eye(3)\n - 0.5 * skew_omega\n + onp.einsum(\"...ij,...jk->...ik\", skew_omega, skew_omega) / 12.0,\n (\n onp.eye(3)\n - 0.5 * skew_omega\n + (\n (\n 1.0\n - theta_safe\n * onp.cos(half_theta_safe)\n / (2.0 * onp.sin(half_theta_safe))\n )\n / theta_squared_safe\n )[..., None, None]\n * onp.einsum(\"...ij,...jk->...ik\", skew_omega, skew_omega)\n ),\n )\n return onp.concatenate(\n [onp.einsum(\"...ij,...j->...i\", V_inv, self.translation()), omega], axis=-1\n )\n\n @override\n def adjoint(self) -> onpt.NDArray[onp.floating]:\n R = self.rotation().as_matrix()\n return onp.concatenate(\n [\n onp.concatenate(\n [R, onp.einsum(\"...ij,...jk->...ik\", _skew(self.translation()), R)],\n axis=-1,\n ),\n onp.concatenate(\n [onp.zeros((*self.get_batch_axes(), 3, 3)), R], axis=-1\n ),\n ],\n axis=-2,\n )\n\n # @classmethod\n # @override\n # def sample_uniform(\n # cls, key: onp.ndarray, batch_axes: jdc.Static[Tuple[int, ...]] = ()\n # ) -> SE3:\n # key0, key1 = jax.random.split(key)\n # return SE3.from_rotation_and_translation(\n # rotation=SO3.sample_uniform(key0, batch_axes=batch_axes),\n # translation=jax.random.uniform(\n # key=key1, shape=(*batch_axes, 3), minval=-1.0, maxval=1.0\n # ),\n # )\n" }, { "path": "viser/src/viser/transforms/_so2.py", "content": "from __future__ import annotations\n\nimport dataclasses\nfrom typing import Tuple\n\nimport numpy as onp\nimport numpy.typing as onpt\nfrom typing_extensions import override\n\nfrom . import _base, hints\nfrom .utils import broadcast_leading_axes, register_lie_group\n\n\n@register_lie_group(\n matrix_dim=2,\n parameters_dim=2,\n tangent_dim=1,\n space_dim=2,\n)\n@dataclasses.dataclass(frozen=True)\nclass SO2(_base.SOBase):\n \"\"\"Special orthogonal group for 2D rotations. Broadcasting rules are the\n same as for numpy.\n\n Ported to numpy from `jaxlie.SO2`.\n\n Internal parameterization is `(cos, sin)`. Tangent parameterization is `(omega,)`.\n \"\"\"\n\n # SO2-specific.\n\n unit_complex: onpt.NDArray[onp.floating]\n \"\"\"Internal parameters. `(cos, sin)`. Shape should be `(*, 2)`.\"\"\"\n\n @override\n def __repr__(self) -> str:\n unit_complex = onp.round(self.unit_complex, 5)\n return f\"{self.__class__.__name__}(unit_complex={unit_complex})\"\n\n @staticmethod\n def from_radians(theta: hints.Scalar) -> SO2:\n \"\"\"Construct a rotation object from a scalar angle.\"\"\"\n cos = onp.cos(theta)\n sin = onp.sin(theta)\n return SO2(unit_complex=onp.stack([cos, sin], axis=-1))\n\n def as_radians(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Compute a scalar angle from a rotation object.\"\"\"\n radians = self.log()[..., 0]\n return radians\n\n # Factory.\n\n @classmethod\n @override\n def identity(cls, batch_axes: Tuple[int, ...] = ()) -> SO2:\n return SO2(\n unit_complex=onp.stack(\n [onp.ones(batch_axes), onp.zeros(batch_axes)], axis=-1\n )\n )\n\n @classmethod\n @override\n def from_matrix(cls, matrix: onpt.NDArray[onp.floating]) -> SO2:\n assert matrix.shape[-2:] == (2, 2)\n return SO2(unit_complex=onp.asarray(matrix[..., :, 0]))\n\n # Accessors.\n\n @override\n def as_matrix(self) -> onpt.NDArray[onp.floating]:\n cos_sin = self.unit_complex\n out = onp.stack(\n [\n # [cos, -sin],\n cos_sin * onp.array([1, -1]),\n # [sin, cos],\n cos_sin[..., ::-1],\n ],\n axis=-2,\n )\n assert out.shape == (*self.get_batch_axes(), 2, 2)\n return out # type: ignore\n\n @override\n def parameters(self) -> onpt.NDArray[onp.floating]:\n return self.unit_complex\n\n # Operations.\n\n @override\n def apply(self, target: onpt.NDArray[onp.floating]) -> onpt.NDArray[onp.floating]:\n assert target.shape[-1:] == (2,)\n self, target = broadcast_leading_axes((self, target))\n return onp.einsum(\"...ij,...j->...i\", self.as_matrix(), target)\n\n @override\n def multiply(self, other: SO2) -> SO2:\n return SO2(\n unit_complex=onp.einsum(\n \"...ij,...j->...i\", self.as_matrix(), other.unit_complex\n )\n )\n\n @classmethod\n @override\n def exp(cls, tangent: onpt.NDArray[onp.floating]) -> SO2:\n assert tangent.shape[-1] == 1\n cos = onp.cos(tangent)\n sin = onp.sin(tangent)\n return SO2(unit_complex=onp.concatenate([cos, sin], axis=-1))\n\n @override\n def log(self) -> onpt.NDArray[onp.floating]:\n return onp.arctan2(\n self.unit_complex[..., 1, None], self.unit_complex[..., 0, None]\n )\n\n @override\n def adjoint(self) -> onpt.NDArray[onp.floating]:\n return onp.ones((*self.get_batch_axes(), 1, 1))\n\n @override\n def inverse(self) -> SO2:\n return SO2(unit_complex=self.unit_complex * onp.array([1, -1]))\n\n @override\n def normalize(self) -> SO2:\n return SO2(\n unit_complex=self.unit_complex\n / onp.linalg.norm(self.unit_complex, axis=-1, keepdims=True)\n )\n\n # @classmethod\n # @override\n # def sample_uniform(\n # cls, key: onp.ndarray, batch_axes: jdc.Static[Tuple[int, ...]] = ()\n # ) -> SO2:\n # out = SO2.from_radians(\n # jax.random.uniform(\n # key=key, shape=batch_axes, minval=0.0, maxval=2.0 * onp.pi)\n # )\n # assert out.get_batch_axes() == batch_axes\n # return out\n" }, { "path": "viser/src/viser/transforms/_so3.py", "content": "from __future__ import annotations\n\nimport dataclasses\nfrom typing import Tuple\n\nimport numpy as onp\nimport numpy.typing as onpt\nfrom typing_extensions import override\n\nfrom . import _base, hints\nfrom .utils import broadcast_leading_axes, get_epsilon, register_lie_group\n\n\n@dataclasses.dataclass(frozen=True)\nclass RollPitchYaw:\n \"\"\"Struct containing roll, pitch, and yaw Euler angles.\"\"\"\n\n roll: onpt.NDArray[onp.floating]\n pitch: onpt.NDArray[onp.floating]\n yaw: onpt.NDArray[onp.floating]\n\n\n@register_lie_group(\n matrix_dim=3,\n parameters_dim=4,\n tangent_dim=3,\n space_dim=3,\n)\n@dataclasses.dataclass(frozen=True)\nclass SO3(_base.SOBase):\n \"\"\"Special orthogonal group for 3D rotations. Broadcasting rules are the same as\n for numpy.\n\n Ported to numpy from `jaxlie.SO3`.\n\n Internal parameterization is `(qw, qx, qy, qz)`. Tangent parameterization is\n `(omega_x, omega_y, omega_z)`.\n \"\"\"\n\n wxyz: onpt.NDArray[onp.floating]\n \"\"\"Internal parameters. `(w, x, y, z)` quaternion. Shape should be `(*, 4)`.\"\"\"\n\n @override\n def __repr__(self) -> str:\n wxyz = onp.round(self.wxyz, 5)\n return f\"{self.__class__.__name__}(wxyz={wxyz})\"\n\n @staticmethod\n def from_x_radians(theta: hints.Scalar) -> SO3:\n \"\"\"Generates a x-axis rotation.\n\n Args:\n angle: X rotation, in radians.\n\n Returns:\n Output.\n \"\"\"\n zeros = onp.zeros_like(theta)\n return SO3.exp(onp.stack([theta, zeros, zeros], axis=-1))\n\n @staticmethod\n def from_y_radians(theta: hints.Scalar) -> SO3:\n \"\"\"Generates a y-axis rotation.\n\n Args:\n angle: Y rotation, in radians.\n\n Returns:\n Output.\n \"\"\"\n zeros = onp.zeros_like(theta)\n return SO3.exp(onp.stack([zeros, theta, zeros], axis=-1))\n\n @staticmethod\n def from_z_radians(theta: hints.Scalar) -> SO3:\n \"\"\"Generates a z-axis rotation.\n\n Args:\n angle: Z rotation, in radians.\n\n Returns:\n Output.\n \"\"\"\n zeros = onp.zeros_like(theta)\n return SO3.exp(onp.stack([zeros, zeros, theta], axis=-1))\n\n @staticmethod\n def from_rpy_radians(\n roll: hints.Scalar,\n pitch: hints.Scalar,\n yaw: hints.Scalar,\n ) -> SO3:\n \"\"\"Generates a transform from a set of Euler angles. Uses the ZYX mobile robot\n convention.\n\n Args:\n roll: X rotation, in radians. Applied first.\n pitch: Y rotation, in radians. Applied second.\n yaw: Z rotation, in radians. Applied last.\n\n Returns:\n Output.\n \"\"\"\n return (\n SO3.from_z_radians(yaw)\n @ SO3.from_y_radians(pitch)\n @ SO3.from_x_radians(roll)\n )\n\n @staticmethod\n def from_quaternion_xyzw(xyzw: onpt.NDArray[onp.floating]) -> SO3:\n \"\"\"Construct a rotation from an `xyzw` quaternion.\n\n Note that `wxyz` quaternions can be constructed using the default dataclass\n constructor.\n\n Args:\n xyzw: xyzw quaternion. Shape should be (*, 4).\n\n Returns:\n Output.\n \"\"\"\n assert xyzw.shape[-1:] == (4,)\n return SO3(onp.roll(xyzw, axis=-1, shift=1))\n\n def as_quaternion_xyzw(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Grab parameters as xyzw quaternion.\"\"\"\n return onp.roll(self.wxyz, axis=-1, shift=-1)\n\n def as_rpy_radians(self) -> RollPitchYaw:\n \"\"\"Computes roll, pitch, and yaw angles. Uses the ZYX mobile robot convention.\n\n Returns:\n Named tuple containing Euler angles in radians.\n \"\"\"\n return RollPitchYaw(\n roll=self.compute_roll_radians(),\n pitch=self.compute_pitch_radians(),\n yaw=self.compute_yaw_radians(),\n )\n\n def compute_roll_radians(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Compute roll angle. Uses the ZYX mobile robot convention.\n\n Returns:\n Euler angle in radians.\n \"\"\"\n # https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Quaternion_to_Euler_angles_conversion\n q0, q1, q2, q3 = onp.moveaxis(self.wxyz, -1, 0)\n return onp.arctan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1**2 + q2**2))\n\n def compute_pitch_radians(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Compute pitch angle. Uses the ZYX mobile robot convention.\n\n Returns:\n Euler angle in radians.\n \"\"\"\n # https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Quaternion_to_Euler_angles_conversion\n q0, q1, q2, q3 = onp.moveaxis(self.wxyz, -1, 0)\n return onp.arcsin(2 * (q0 * q2 - q3 * q1))\n\n def compute_yaw_radians(self) -> onpt.NDArray[onp.floating]:\n \"\"\"Compute yaw angle. Uses the ZYX mobile robot convention.\n\n Returns:\n Euler angle in radians.\n \"\"\"\n # https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Quaternion_to_Euler_angles_conversion\n q0, q1, q2, q3 = onp.moveaxis(self.wxyz, -1, 0)\n return onp.arctan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2**2 + q3**2))\n\n # Factory.\n\n @classmethod\n @override\n def identity(cls, batch_axes: Tuple[int, ...] = ()) -> SO3:\n return SO3(\n wxyz=onp.broadcast_to(onp.array([1.0, 0.0, 0.0, 0.0]), (*batch_axes, 4))\n )\n\n @classmethod\n @override\n def from_matrix(cls, matrix: onpt.NDArray[onp.floating]) -> SO3:\n assert matrix.shape[-2:] == (3, 3)\n\n # Modified from:\n # > \"Converting a Rotation Matrix to a Quaternion\" from Mike Day\n # > https://d3cw3dd2w32x2b.cloudfront.net/wp-content/uploads/2015/01/matrix-to-quat.pdf\n\n def case0(m):\n t = 1 + m[..., 0, 0] - m[..., 1, 1] - m[..., 2, 2]\n q = onp.stack(\n [\n m[..., 2, 1] - m[..., 1, 2],\n t,\n m[..., 1, 0] + m[..., 0, 1],\n m[..., 0, 2] + m[..., 2, 0],\n ],\n axis=-1,\n )\n return t, q\n\n def case1(m):\n t = 1 - m[..., 0, 0] + m[..., 1, 1] - m[..., 2, 2]\n q = onp.stack(\n [\n m[..., 0, 2] - m[..., 2, 0],\n m[..., 1, 0] + m[..., 0, 1],\n t,\n m[..., 2, 1] + m[..., 1, 2],\n ],\n axis=-1,\n )\n return t, q\n\n def case2(m):\n t = 1 - m[..., 0, 0] - m[..., 1, 1] + m[..., 2, 2]\n q = onp.stack(\n [\n m[..., 1, 0] - m[..., 0, 1],\n m[..., 0, 2] + m[..., 2, 0],\n m[..., 2, 1] + m[..., 1, 2],\n t,\n ],\n axis=-1,\n )\n return t, q\n\n def case3(m):\n t = 1 + m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2]\n q = onp.stack(\n [\n t,\n m[..., 2, 1] - m[..., 1, 2],\n m[..., 0, 2] - m[..., 2, 0],\n m[..., 1, 0] - m[..., 0, 1],\n ],\n axis=-1,\n )\n return t, q\n\n # Compute four cases, then pick the most precise one.\n # Probably worth revisiting this!\n case0_t, case0_q = case0(matrix)\n case1_t, case1_q = case1(matrix)\n case2_t, case2_q = case2(matrix)\n case3_t, case3_q = case3(matrix)\n\n cond0 = matrix[..., 2, 2] < 0\n cond1 = matrix[..., 0, 0] > matrix[..., 1, 1]\n cond2 = matrix[..., 0, 0] < -matrix[..., 1, 1]\n\n t = onp.where(\n cond0,\n onp.where(cond1, case0_t, case1_t),\n onp.where(cond2, case2_t, case3_t),\n )\n q = onp.where(\n cond0[..., None],\n onp.where(cond1[..., None], case0_q, case1_q),\n onp.where(cond2[..., None], case2_q, case3_q),\n )\n\n # We can also choose to branch, but this is slower.\n # t, q = jax.lax.cond(\n # matrix[2, 2] < 0,\n # true_fun=lambda matrix: jax.lax.cond(\n # matrix[0, 0] > matrix[1, 1],\n # true_fun=case0,\n # false_fun=case1,\n # operand=matrix,\n # ),\n # false_fun=lambda matrix: jax.lax.cond(\n # matrix[0, 0] < -matrix[1, 1],\n # true_fun=case2,\n # false_fun=case3,\n # operand=matrix,\n # ),\n # operand=matrix,\n # )\n\n return SO3(wxyz=q * 0.5 / onp.sqrt(t[..., None]))\n\n # Accessors.\n\n @override\n def as_matrix(self) -> onpt.NDArray[onp.floating]:\n norm_sq = onp.sum(onp.square(self.wxyz), axis=-1, keepdims=True)\n q = self.wxyz * onp.sqrt(2.0 / norm_sq) # (*, 4)\n q_outer = onp.einsum(\"...i,...j->...ij\", q, q) # (*, 4, 4)\n return onp.stack(\n [\n 1.0 - q_outer[..., 2, 2] - q_outer[..., 3, 3],\n q_outer[..., 1, 2] - q_outer[..., 3, 0],\n q_outer[..., 1, 3] + q_outer[..., 2, 0],\n q_outer[..., 1, 2] + q_outer[..., 3, 0],\n 1.0 - q_outer[..., 1, 1] - q_outer[..., 3, 3],\n q_outer[..., 2, 3] - q_outer[..., 1, 0],\n q_outer[..., 1, 3] - q_outer[..., 2, 0],\n q_outer[..., 2, 3] + q_outer[..., 1, 0],\n 1.0 - q_outer[..., 1, 1] - q_outer[..., 2, 2],\n ],\n axis=-1,\n ).reshape(*q.shape[:-1], 3, 3)\n\n @override\n def parameters(self) -> onpt.NDArray[onp.floating]:\n return self.wxyz\n\n # Operations.\n\n @override\n def apply(self, target: onpt.NDArray[onp.floating]) -> onpt.NDArray[onp.floating]:\n assert target.shape[-1:] == (3,)\n self, target = broadcast_leading_axes((self, target))\n\n # Compute using quaternion multiplys.\n padded_target = onp.concatenate(\n [onp.zeros((*self.get_batch_axes(), 1)), target], axis=-1\n )\n return (self @ SO3(wxyz=padded_target) @ self.inverse()).wxyz[..., 1:]\n\n @override\n def multiply(self, other: SO3) -> SO3:\n w0, x0, y0, z0 = onp.moveaxis(self.wxyz, -1, 0)\n w1, x1, y1, z1 = onp.moveaxis(other.wxyz, -1, 0)\n return SO3(\n wxyz=onp.stack(\n [\n -x0 * x1 - y0 * y1 - z0 * z1 + w0 * w1,\n x0 * w1 + y0 * z1 - z0 * y1 + w0 * x1,\n -x0 * z1 + y0 * w1 + z0 * x1 + w0 * y1,\n x0 * y1 - y0 * x1 + z0 * w1 + w0 * z1,\n ],\n axis=-1,\n )\n )\n\n @classmethod\n @override\n def exp(cls, tangent: onpt.NDArray[onp.floating]) -> SO3:\n # Reference:\n # > https://github.com/strasdat/Sophus/blob/a0fe89a323e20c42d3cecb590937eb7a06b8343a/sophus/so3.hpp#L583\n\n assert tangent.shape[-1:] == (3,)\n\n theta_squared = onp.sum(onp.square(tangent), axis=-1)\n theta_pow_4 = theta_squared * theta_squared\n use_taylor = theta_squared < get_epsilon(tangent.dtype)\n\n # Shim to avoid NaNs in onp.where branches, which cause failures for\n # reverse-mode AD in JAX. This isn't needed for vanilla numpy.\n safe_theta = onp.sqrt(\n onp.where(\n use_taylor,\n onp.ones_like(theta_squared), # Any constant value should do here.\n theta_squared,\n )\n )\n safe_half_theta = 0.5 * safe_theta\n\n real_factor = onp.where(\n use_taylor,\n 1.0 - theta_squared / 8.0 + theta_pow_4 / 384.0,\n onp.cos(safe_half_theta),\n )\n\n imaginary_factor = onp.where(\n use_taylor,\n 0.5 - theta_squared / 48.0 + theta_pow_4 / 3840.0,\n onp.sin(safe_half_theta) / safe_theta,\n )\n\n return SO3(\n wxyz=onp.concatenate(\n [\n real_factor[..., None],\n imaginary_factor[..., None] * tangent,\n ],\n axis=-1,\n )\n )\n\n @override\n def log(self) -> onpt.NDArray[onp.floating]:\n # Reference:\n # > https://github.com/strasdat/Sophus/blob/a0fe89a323e20c42d3cecb590937eb7a06b8343a/sophus/so3.hpp#L247\n\n w = self.wxyz[..., 0]\n norm_sq = onp.sum(onp.square(self.wxyz[..., 1:]), axis=-1)\n use_taylor = norm_sq < get_epsilon(norm_sq.dtype)\n\n # Shim to avoid NaNs in onp.where branches, which cause failures for\n # reverse-mode AD in JAX. This isn't needed for vanilla numpy.\n norm_safe = onp.sqrt(\n onp.where(\n use_taylor,\n 1.0, # Any non-zero value should do here.\n norm_sq,\n )\n )\n w_safe = onp.where(use_taylor, w, 1.0)\n atan_n_over_w = onp.arctan2(\n onp.where(w < 0, -norm_safe, norm_safe),\n onp.abs(w),\n )\n atan_factor = onp.where(\n use_taylor,\n 2.0 / w_safe - 2.0 / 3.0 * norm_sq / w_safe**3,\n onp.where(\n onp.abs(w) < get_epsilon(w.dtype),\n onp.where(w > 0, 1.0, -1.0) * onp.pi / norm_safe,\n 2.0 * atan_n_over_w / norm_safe,\n ),\n )\n\n return atan_factor[..., None] * self.wxyz[..., 1:] # type: ignore\n\n @override\n def adjoint(self) -> onpt.NDArray[onp.floating]:\n return self.as_matrix()\n\n @override\n def inverse(self) -> SO3:\n # Negate complex terms.\n return SO3(wxyz=self.wxyz * onp.array([1, -1, -1, -1]))\n\n @override\n def normalize(self) -> SO3:\n return SO3(wxyz=self.wxyz / onp.linalg.norm(self.wxyz, axis=-1, keepdims=True))\n\n # @classmethod\n # @override\n # def sample_uniform(\n # cls, key: onp.ndarray, batch_axes: jdc.Static[Tuple[int, ...]] = ()\n # ) -> SO3:\n # # Uniformly sample over S^3.\n # # > Reference: http://planning.cs.uiuc.edu/node198.html\n # u1, u2, u3 = onp.moveaxis(\n # jax.random.uniform(\n # key=key,\n # shape=(*batch_axes, 3),\n # minval=onp.zeros(3),\n # maxval=onp.array([1.0, 2.0 * onp.pi, 2.0 * onp.pi]),\n # ),\n # -1,\n # 0,\n # )\n # a = onp.sqrt(1.0 - u1)\n # b = onp.sqrt(u1)\n #\n # return SO3(\n # wxyz=onp.stack(\n # [\n # a * onp.sin(u2),\n # a * onp.cos(u2),\n # b * onp.sin(u3),\n # b * onp.cos(u3),\n # ],\n # axis=-1,\n # )\n # )\n" }, { "path": "viser/src/viser/transforms/hints/__init__.py", "content": "from typing import Union\n\nimport numpy as onp\nimport numpy.typing as onpt\n\n# Type aliases Numpy arrays; primarily for function inputs.\n\nScalar = Union[float, onpt.NDArray[onp.floating]]\n\"\"\"Type alias for `Union[float, Array]`.\"\"\"\n\n\n__all__ = [\n \"Scalar\",\n]\n" }, { "path": "viser/src/viser/transforms/utils/__init__.py", "content": "from ._utils import broadcast_leading_axes, get_epsilon, register_lie_group\n\n__all__ = [\"get_epsilon\", \"register_lie_group\", \"broadcast_leading_axes\"]\n" }, { "path": "viser/src/viser/transforms/utils/_utils.py", "content": "from typing import TYPE_CHECKING, Callable, Tuple, Type, TypeVar, Union, cast\n\nimport numpy as onp\n\nif TYPE_CHECKING:\n from .._base import MatrixLieGroup\n\n\nT = TypeVar(\"T\", bound=\"MatrixLieGroup\")\n\n\ndef get_epsilon(dtype: onp.dtype) -> float:\n \"\"\"Helper for grabbing type-specific precision constants.\n\n Args:\n dtype: Datatype.\n\n Returns:\n Output float.\n \"\"\"\n if dtype == onp.float32:\n return 1e-5\n elif dtype == onp.float64:\n return 1e-10\n else:\n assert False\n\n\ndef register_lie_group(\n *,\n matrix_dim: int,\n parameters_dim: int,\n tangent_dim: int,\n space_dim: int,\n) -> Callable[[Type[T]], Type[T]]:\n \"\"\"Decorator for registering Lie group dataclasses.\n\n Sets dimensionality class variables, and marks all methods for JIT compilation.\n \"\"\"\n\n def _wrap(cls: Type[T]) -> Type[T]:\n # Register dimensions as class attributes.\n cls.matrix_dim = matrix_dim\n cls.parameters_dim = parameters_dim\n cls.tangent_dim = tangent_dim\n cls.space_dim = space_dim\n return cls\n\n return _wrap\n\n\nTupleOfBroadcastable = TypeVar(\n \"TupleOfBroadcastable\",\n bound=\"Tuple[Union[MatrixLieGroup, onp.ndarray], ...]\",\n)\n\n\ndef broadcast_leading_axes(inputs: TupleOfBroadcastable) -> TupleOfBroadcastable:\n \"\"\"Broadcast leading axes of arrays. Takes tuples of either:\n - an array, which we assume has shape (*, D).\n - a Lie group object.\"\"\"\n\n from .._base import MatrixLieGroup\n\n array_inputs = [\n (\n (x.parameters(), (x.parameters_dim,))\n if isinstance(x, MatrixLieGroup)\n else (x, x.shape[-1:])\n )\n for x in inputs\n ]\n for array, shape_suffix in array_inputs:\n assert array.shape[-len(shape_suffix) :] == shape_suffix\n batch_axes = onp.broadcast_shapes(\n *[array.shape[: -len(suffix)] for array, suffix in array_inputs]\n )\n broadcasted_arrays = tuple(\n onp.broadcast_to(array, batch_axes + shape_suffix)\n for (array, shape_suffix) in array_inputs\n )\n return cast(\n TupleOfBroadcastable,\n tuple(\n array if not isinstance(inp, MatrixLieGroup) else type(inp)(array)\n for array, inp in zip(broadcasted_arrays, inputs)\n ),\n )\n" }, { "path": "viser/sync_message_defs.py", "content": "\"\"\"Generate typescript message definitions from Python dataclasses.\"\"\"\n\nimport pathlib\nimport subprocess\n\nimport viser.infra\nfrom viser._messages import Message\n\nif __name__ == \"__main__\":\n # Generate typescript source.\n defs = viser.infra.generate_typescript_interfaces(Message)\n\n # Write to file.\n target_path = pathlib.Path(__file__).parent / pathlib.Path(\n \"src/viser/client/src/WebsocketMessages.tsx\"\n )\n assert target_path.exists()\n target_path.write_text(defs)\n print(f\"Wrote to {target_path}\")\n\n # Run prettier.\n subprocess.run(args=[\"npx\", \"prettier\", \"-w\", str(target_path)], check=False)\n" }, { "path": "viser/visualize_megasam.py", "content": "import time\nimport sys\nimport argparse\nfrom pathlib import Path\n\nimport numpy as onp\nimport tyro\nfrom tqdm.auto import tqdm\n\nimport viser\nimport viser.extras\nimport viser.transforms as tf\nimport matplotlib.cm as cm # For colormap\n\nimport os\nimport cv2\nimport numpy as np\nimport argparse\n\ndef main(\n data: Path = \"./demo_tmp/NULL.npz\",\n downsample_factor: int = 1,\n max_frames: int = 300,\n share: bool = False,\n conf_threshold: float = 0.,\n foreground_conf_threshold: float = 0.,\n point_size: float = 0.001,\n camera_frustum_scale: float = 0.02,\n no_mask: bool = True,\n xyzw: bool = True,\n axes_scale: float = 0.25,\n bg_downsample_factor: int = 1,\n init_conf: bool = False,\n cam_thickness: float = 1.5,\n) -> None:\n from pathlib import Path # <-- Import Path here if not already imported\n\n data = np.load(data)\n \n server = viser.ViserServer()\n if share:\n server.request_share_url()\n\n server.scene.set_up_direction('-z')\n if no_mask: # not using dynamic / static mask\n init_conf = True # must use init_conf map, to avoid depth cleaning\n fg_conf_thre = conf_threshold # now fg_conf_thre is the same as conf_thre\n print(\"Loading frames!\")\n loader = viser.extras.Record3dLoader_Customized_Megasam(\n data,\n conf_threshold=conf_threshold,\n foreground_conf_threshold=foreground_conf_threshold,\n no_mask=no_mask,\n xyzw=xyzw,\n init_conf=init_conf,\n )\n num_frames = min(max_frames, loader.num_frames())\n\n # Add playback UI.\n with server.gui.add_folder(\"Playback\"):\n gui_timestep = server.gui.add_slider(\n \"Timestep\",\n min=0,\n max=num_frames - 1,\n step=1,\n initial_value=0,\n disabled=True,\n )\n gui_next_frame = server.gui.add_button(\"Next Frame\", disabled=True)\n gui_prev_frame = server.gui.add_button(\"Prev Frame\", disabled=True)\n gui_playing = server.gui.add_checkbox(\"Playing\", True)\n gui_framerate = server.gui.add_slider(\n \"FPS\", min=1, max=60, step=0.1, initial_value=loader.fps\n )\n gui_framerate_options = server.gui.add_button_group(\n \"FPS options\", (\"10\", \"20\", \"30\", \"60\")\n )\n gui_show_all_frames = server.gui.add_checkbox(\"Show all frames\", False)\n gui_stride = server.gui.add_slider(\n \"Stride\",\n min=1,\n max=num_frames,\n step=1,\n initial_value=1,\n disabled=True, # Initially disabled\n )\n\n # Add recording UI.\n with server.gui.add_folder(\"Recording\"):\n gui_record_scene = server.gui.add_button(\"Record Scene\")\n\n # Frame step buttons.\n @gui_next_frame.on_click\n def _(_) -> None:\n gui_timestep.value = (gui_timestep.value + 1) % num_frames\n\n @gui_prev_frame.on_click\n def _(_) -> None:\n gui_timestep.value = (gui_timestep.value - 1) % num_frames\n\n # Disable frame controls when we're playing.\n @gui_playing.on_update\n def _(_) -> None:\n gui_timestep.disabled = gui_playing.value or gui_show_all_frames.value\n gui_next_frame.disabled = gui_playing.value or gui_show_all_frames.value\n gui_prev_frame.disabled = gui_playing.value or gui_show_all_frames.value\n\n # Toggle frame visibility when the timestep slider changes.\n @gui_timestep.on_update\n def _(_) -> None:\n nonlocal prev_timestep\n current_timestep = gui_timestep.value\n if not gui_show_all_frames.value:\n with server.atomic():\n frame_nodes[current_timestep].visible = True\n frame_nodes[prev_timestep].visible = False\n prev_timestep = current_timestep\n server.flush() # Optional!\n\n # Show or hide all frames based on the checkbox.\n @gui_show_all_frames.on_update\n def _(_) -> None:\n gui_stride.disabled = not gui_show_all_frames.value # Enable/disable stride slider\n if gui_show_all_frames.value:\n # Show frames with stride\n stride = gui_stride.value\n with server.atomic():\n for i, frame_node in enumerate(frame_nodes):\n frame_node.visible = (i % stride == 0)\n # Disable playback controls\n gui_playing.disabled = True\n gui_timestep.disabled = True\n gui_next_frame.disabled = True\n gui_prev_frame.disabled = True\n else:\n # Show only the current frame\n current_timestep = gui_timestep.value\n with server.atomic():\n for i, frame_node in enumerate(frame_nodes):\n frame_node.visible = i == current_timestep\n # Re-enable playback controls\n gui_playing.disabled = False\n gui_timestep.disabled = gui_playing.value\n gui_next_frame.disabled = gui_playing.value\n gui_prev_frame.disabled = gui_playing.value\n\n # Update frame visibility when the stride changes.\n @gui_stride.on_update\n def _(_) -> None:\n if gui_show_all_frames.value:\n # Update frame visibility based on new stride\n stride = gui_stride.value\n with server.atomic():\n for i, frame_node in enumerate(frame_nodes):\n frame_node.visible = (i % stride == 0)\n\n # Recording handler\n @gui_record_scene.on_click\n def _(_):\n gui_record_scene.disabled = True\n\n # Save the original frame visibility state\n original_visibility = [frame_node.visible for frame_node in frame_nodes]\n\n rec = server._start_scene_recording()\n rec.set_loop_start()\n \n # Determine sleep duration based on current FPS\n sleep_duration = 1.0 / gui_framerate.value if gui_framerate.value > 0 else 0.033 # Default to ~30 FPS\n \n if gui_show_all_frames.value:\n # Record all frames according to the stride\n stride = gui_stride.value\n frames_to_record = [i for i in range(num_frames) if i % stride == 0]\n else:\n # Record the frames in sequence\n frames_to_record = range(num_frames)\n \n for t in frames_to_record:\n # Update the scene to show frame t\n with server.atomic():\n for i, frame_node in enumerate(frame_nodes):\n frame_node.visible = (i == t) if not gui_show_all_frames.value else (i % gui_stride.value == 0)\n server.flush()\n rec.insert_sleep(sleep_duration)\n\n # set all invisible\n with server.atomic():\n for frame_node in frame_nodes:\n frame_node.visible = False\n \n # Finish recording\n bs = rec.end_and_serialize()\n \n # Save the recording to a file\n output_path = Path(f\"./viser_result/recording_{str(data).split('/')[-1]}.viser\")\n # make sure the output directory exists\n output_path.parent.mkdir(parents=True, exist_ok=True)\n output_path.write_bytes(bs)\n print(f\"Recording saved to {output_path.resolve()}\")\n \n # Restore the original frame visibility state\n with server.atomic():\n for frame_node, visibility in zip(frame_nodes, original_visibility):\n frame_node.visible = visibility\n server.flush()\n \n gui_record_scene.disabled = False\n\n # Load in frames.\n server.scene.add_frame(\n \"/frames\",\n wxyz=tf.SO3.exp(onp.array([onp.pi / 2.0, 0.0, 0.0])).wxyz,\n position=(0, 0, 0),\n show_axes=False,\n )\n frame_nodes: list[viser.FrameHandle] = []\n bg_positions = []\n bg_colors = []\n for i in tqdm(range(num_frames)):\n frame = loader.get_frame(i)\n position, color, bg_position, bg_color = frame.get_point_cloud(downsample_factor, bg_downsample_factor)\n\n bg_positions.append(bg_position)\n bg_colors.append(bg_color)\n\n # Add base frame.\n frame_nodes.append(server.scene.add_frame(f\"/frames/t{i}\", show_axes=False))\n\n # Place the point cloud in the frame.\n server.scene.add_point_cloud(\n name=f\"/frames/t{i}/point_cloud\",\n points=position,\n colors=color,\n point_size=point_size,\n point_shape=\"rounded\",\n )\n\n # Compute color for frustum based on frame index.\n norm_i = i / (num_frames - 1) if num_frames > 1 else 0 # Normalize index to [0, 1]\n color_rgba = cm.viridis(norm_i) # Get RGBA color from colormap\n color_rgb = color_rgba[:3] # Use RGB components\n\n # Place the frustum with the computed color.\n fov = 2 * onp.arctan2(frame.rgb.shape[0] / 2, frame.K[0, 0])\n aspect = frame.rgb.shape[1] / frame.rgb.shape[0]\n server.scene.add_camera_frustum(\n f\"/frames/t{i}/frustum\",\n fov=fov,\n aspect=aspect,\n scale=camera_frustum_scale,\n image=frame.rgb[::downsample_factor, ::downsample_factor],\n wxyz=tf.SO3.from_matrix(frame.T_world_camera[:3, :3]).wxyz,\n position=frame.T_world_camera[:3, 3],\n color=color_rgb, # Set the color for the frustum\n thickness=cam_thickness,\n )\n\n # Add some axes. (Commented out to hide coordinate axes)\n # server.scene.add_frame(\n # f\"/frames/t{i}/frustum/axes\",\n # axes_length=camera_frustum_scale * axes_scale * 10,\n # axes_radius=camera_frustum_scale * axes_scale,\n # )\n\n # Initialize frame visibility.\n for i, frame_node in enumerate(frame_nodes):\n if gui_show_all_frames.value:\n frame_node.visible = (i % gui_stride.value == 0)\n else:\n frame_node.visible = i == gui_timestep.value\n\n # Add background frame.\n bg_positions = onp.concatenate(bg_positions, axis=0)\n bg_colors = onp.concatenate(bg_colors, axis=0)\n server.scene.add_point_cloud(\n name=f\"/frames/background\",\n points=bg_positions,\n colors=bg_colors,\n point_size=point_size,\n point_shape=\"rounded\",\n )\n\n # Playback update loop.\n prev_timestep = gui_timestep.value\n while True:\n if gui_playing.value and not gui_show_all_frames.value:\n gui_timestep.value = (gui_timestep.value + 1) % num_frames\n time.sleep(1.0 / gui_framerate.value)\n\n\nif __name__ == \"__main__\":\n\n tyro.cli(main)\n" }, { "path": "viser/visualize_pose.py", "content": "#!/usr/bin/env python3\n\"\"\"\nCamera Pose Visualization Module\n\nThis module provides comprehensive tools for visualizing camera poses and trajectories\nin 3D space using Plotly. It supports both static and animated visualizations with\nautomatic camera view optimization.\n\nAdapted from: https://huggingface.co/datasets/nvidia/dynpose-100k/blob/main/scripts/visualize_pose.py\n\"\"\"\n\nimport argparse\nimport matplotlib\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport plotly.graph_objs as go\nimport plotly.io as pio\nfrom tqdm import tqdm\nimport einops\nimport torch\n\n# Use non-interactive backend for matplotlib to avoid display issues\nmatplotlib.use(\"agg\")\n\n\nclass Pose:\n \"\"\"\n A class of operations on camera poses (numpy arrays with shape [...,3,4]).\n Each [3,4] camera pose takes the form of [R|t].\n \"\"\"\n\n def __call__(self, R=None, t=None):\n \"\"\"\n Construct a camera pose from the given rotation matrix R and/or translation vector t.\n \"\"\"\n assert R is not None or t is not None\n if R is None:\n if not isinstance(t, np.ndarray):\n t = np.array(t)\n R = np.eye(3).repeat(*t.shape[:-1], 1, 1)\n elif t is None:\n if not isinstance(R, np.ndarray):\n R = np.array(R)\n t = np.zeros(R.shape[:-1])\n else:\n if not isinstance(R, np.ndarray):\n R = np.array(R)\n if not isinstance(t, np.ndarray):\n t = np.array(t)\n assert R.shape[:-1] == t.shape and R.shape[-2:] == (3, 3)\n R = R.astype(np.float32)\n t = t.astype(np.float32)\n pose = np.concatenate([R, t[..., None]], axis=-1) # [...,3,4]\n assert pose.shape[-2:] == (3, 4)\n return pose\n\n def invert(self, pose, use_inverse=False):\n \"\"\"\n Invert a camera pose.\n \"\"\"\n R, t = pose[..., :3], pose[..., 3:]\n R_inv = np.linalg.inv(R) if use_inverse else R.transpose(0, 2, 1)\n t_inv = (-R_inv @ t)[..., 0]\n pose_inv = self(R=R_inv, t=t_inv)\n return pose_inv\n\n def compose(self, pose_list):\n \"\"\"\n Compose a sequence of poses together.\n pose_new(x) = poseN o ... o pose2 o pose1(x)\n \"\"\"\n pose_new = pose_list[0]\n for pose in pose_list[1:]:\n pose_new = self.compose_pair(pose_new, pose)\n return pose_new\n\n def compose_pair(self, pose_a, pose_b):\n \"\"\"\n Compose two poses together.\n \"\"\"\n R_a, t_a = pose_a[..., :3], pose_a[..., 3:]\n R_b, t_b = pose_b[..., :3], pose_b[..., 3:]\n R_new = R_b @ R_a\n t_new = (R_b @ t_a + t_b)[..., 0]\n pose_new = self(R=R_new, t=t_new)\n return pose_new\n\n def scale_center(self, pose, scale):\n \"\"\"\n Scale the camera center from the origin.\n 0 = R@c+t --> c = -R^T@t (camera center in world coordinates)\n 0 = R@(sc)+t' --> t' = -R@(sc) = -R@(-R^T@st) = st\n \"\"\"\n R, t = pose[..., :3], pose[..., 3:]\n pose_new = np.concatenate([R, t * scale], axis=-1)\n return pose_new\n\n\ndef to_hom(X):\n \"\"\"\n Convert points to homogeneous coordinates by appending ones.\n \"\"\"\n X_hom = np.concatenate([X, np.ones_like(X[..., :1])], axis=-1)\n return X_hom\n\n\ndef cam2world(X, pose):\n \"\"\"\n Transform points from camera coordinates to world coordinates.\n \"\"\"\n X_hom = to_hom(X)\n pose_inv = Pose().invert(pose)\n return X_hom @ pose_inv.transpose(0, 2, 1)\n\n\ndef get_camera_mesh(pose, depth=1):\n \"\"\"\n Create a 3D mesh representation of camera frustums for visualization.\n \"\"\"\n # Define camera frustum geometry: 4 corners of image plane + camera center\n vertices = (\n np.array(\n [[-0.5, -0.5, 1], [0.5, -0.5, 1], [0.5, 0.5, 1], [-0.5, 0.5, 1], [0, 0, 0]]\n )\n * depth\n ) # Shape: [5, 3] - 4 image plane corners + camera center\n\n # Define triangular faces for the camera frustum mesh\n faces = np.array(\n [[0, 1, 2], [0, 2, 3], [0, 1, 4], [1, 2, 4], [2, 3, 4], [3, 0, 4]]\n ) # Shape: [6, 3] - 6 triangular faces forming the pyramid\n\n # Transform vertices from camera space to world space\n vertices = cam2world(vertices[None], pose) # Shape: [N, 5, 3]\n\n # Create wireframe lines connecting: corners -> center -> next corner\n wireframe = vertices[:, [0, 1, 2, 3, 0, 4, 1, 2, 4, 3]] # Shape: [N, 10, 3]\n\n return vertices, faces, wireframe\n\n\n# def merge_xyz_indicators_plotly(xyz):\n# \"\"\"Merge xyz coordinate indicators for plotly visualization.\"\"\"\n# xyz = xyz[:, [[-1, 0], [-1, 1], [-1, 2]]] # [N,3,2,3]\n# xyz_0, xyz_1 = unbind_np(xyz, axis=2) # [N,3,3]\n# xyz_dummy = xyz_0 * np.nan\n# xyz_merged = np.stack([xyz_0, xyz_1, xyz_dummy], axis=2) # [N,3,3,3]\n# xyz_merged = xyz_merged.reshape(-1, 3)\n# return xyz_merged\n\n\n# def get_xyz_indicators(pose, length=0.1):\n# \"\"\"Get xyz coordinate axis indicators for a camera pose.\"\"\"\n# xyz = np.eye(4, 3)[None] * length\n# xyz = cam2world(xyz, pose)\n# return xyz\n\n\ndef merge_wireframes_plotly(wireframe):\n \"\"\"\n Merge camera wireframes for efficient Plotly visualization.\n \"\"\"\n wf_dummy = wireframe[:, :1] * np.nan # Create NaN separators\n wireframe_merged = np.concatenate([wireframe, wf_dummy], axis=1).reshape(-1, 3)\n return wireframe_merged\n\n\ndef merge_meshes(vertices, faces):\n \"\"\"\n Merge multiple camera meshes into a single mesh for efficient rendering.\n \"\"\"\n mesh_N, vertex_N = vertices.shape[:2]\n # Adjust face indices for each mesh by adding vertex offset\n faces_merged = np.concatenate([faces + i * vertex_N for i in range(mesh_N)], axis=0)\n # Flatten all vertices into single array\n vertices_merged = vertices.reshape(-1, vertices.shape[-1])\n return vertices_merged, faces_merged\n\n\ndef unbind_np(array, axis=0):\n \"\"\"\n Split numpy array along specified axis into a list of arrays.\n \"\"\"\n if axis == 0:\n return [array[i, :] for i in range(array.shape[0])]\n elif axis == 1 or (len(array.shape) == 2 and axis == -1):\n return [array[:, j] for j in range(array.shape[1])]\n elif axis == 2 or (len(array.shape) == 3 and axis == -1):\n return [array[:, :, j] for j in range(array.shape[2])]\n else:\n raise ValueError(\"Invalid axis. Use 0 for rows, 1 for columns, or 2 for depth.\")\n\n\ndef plotly_visualize_pose(\n poses, vis_depth=0.5, xyz_length=0.5, center_size=2, xyz_width=5, mesh_opacity=0.05\n):\n \"\"\"\n Create comprehensive Plotly visualization traces for camera poses.\n \"\"\"\n N = len(poses)\n\n # Calculate camera centers in world coordinates\n centers_cam = np.zeros([N, 1, 3]) # Camera centers in camera space (origin)\n centers_world = cam2world(centers_cam, poses) # Transform to world space\n centers_world = centers_world[:, 0] # Remove extra dimension [N, 3]\n\n # Generate camera frustum geometry\n vertices, faces, wireframe = get_camera_mesh(poses, depth=vis_depth)\n\n # Merge all camera meshes into single arrays for efficient rendering\n vertices_merged, faces_merged = merge_meshes(vertices, faces)\n wireframe_merged = merge_wireframes_plotly(wireframe)\n\n # Extract x, y, z coordinates for Plotly\n wireframe_x, wireframe_y, wireframe_z = unbind_np(wireframe_merged, axis=-1)\n centers_x, centers_y, centers_z = unbind_np(centers_world, axis=-1)\n vertices_x, vertices_y, vertices_z = unbind_np(vertices_merged, axis=-1)\n\n # Set up rainbow color mapping for trajectory progression\n color_map = plt.get_cmap(\"gist_rainbow\") # red -> yellow -> green -> blue -> purple\n center_color = []\n faces_merged_color = []\n wireframe_color = []\n\n # Determine quarter positions for emphasis (start, 1/3, 2/3, end)\n quarter_indices = set([0]) # Always include start\n if N >= 3:\n quarter_indices.add(N // 3)\n quarter_indices.add(2 * N // 3)\n quarter_indices.add(N - 1) # Always include end\n\n # Apply colors with emphasis on key trajectory points\n for i in range(N):\n # Emphasize quarter positions with higher opacity and brightness\n is_quarter = i in quarter_indices\n alpha = 6.0 if is_quarter else 0.4 # Higher opacity for key points\n\n # Generate color from rainbow colormap\n r, g, b, _ = color_map(i / (N - 1))\n rgb = np.array([r, g, b]) * (1.2 if is_quarter else 0.8) # Brighten key points\n rgba = np.concatenate([rgb, [alpha]])\n\n # Apply colors to all visualization elements\n wireframe_color += [rgba] * 11 # 11 line segments per camera wireframe\n center_color += [rgba]\n faces_merged_color += [rgba] * 6 # 6 triangular faces per camera frustum\n\n # Create Plotly trace objects\n plotly_traces = [\n # Camera wireframe outlines\n go.Scatter3d(\n x=wireframe_x,\n y=wireframe_y,\n z=wireframe_z,\n mode=\"lines\",\n line=dict(color=wireframe_color, width=1),\n name=\"Camera Wireframes\",\n ),\n # Camera center points\n go.Scatter3d(\n x=centers_x,\n y=centers_y,\n z=centers_z,\n mode=\"markers\",\n marker=dict(color=center_color, size=center_size, opacity=1),\n name=\"Camera Centers\",\n ),\n # Camera frustum mesh faces\n go.Mesh3d(\n x=vertices_x,\n y=vertices_y,\n z=vertices_z,\n i=[f[0] for f in faces_merged],\n j=[f[1] for f in faces_merged],\n k=[f[2] for f in faces_merged],\n facecolor=faces_merged_color,\n opacity=mesh_opacity,\n name=\"Camera Frustums\",\n ),\n ]\n return plotly_traces\n\n\ndef compute_optimal_camera_view(poses):\n \"\"\"\n Compute optimal camera view parameters to ensure the entire trajectory is visible\n and aesthetically pleasing.\n \"\"\"\n # Calculate all camera positions in world coordinates\n centers_cam = np.zeros([len(poses), 1, 3])\n centers_world = cam2world(centers_cam, poses)[:, 0]\n\n # Compute bounding box of the trajectory\n min_coords = np.min(centers_world, axis=0)\n max_coords = np.max(centers_world, axis=0)\n ranges = max_coords - min_coords\n\n # Calculate trajectory center point\n trajectory_center = (min_coords + max_coords) / 2\n\n # Calculate maximum range for adaptive scaling\n max_range = np.max(ranges)\n\n # Set minimum range to avoid division by zero for very small trajectories\n if max_range < 1e-6:\n max_range = 1.0\n ranges = np.ones(3)\n\n # Calculate principal direction of trajectory using PCA (Principal Component Analysis)\n if len(centers_world) > 1:\n # Center the points by subtracting the mean\n centered_points = centers_world - trajectory_center\n\n # Compute covariance matrix for PCA\n cov_matrix = np.cov(centered_points.T)\n\n # Calculate eigenvalues and eigenvectors\n eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)\n\n # Sort by eigenvalues in descending order\n idx = np.argsort(eigenvalues)[::-1]\n eigenvalues = eigenvalues[idx]\n eigenvectors = eigenvectors[:, idx]\n\n # Main direction is the first eigenvector (highest variance)\n main_direction = eigenvectors[:, 0]\n\n # Ensure main direction points towards trajectory's positive direction\n start_to_end = centers_world[-1] - centers_world[0]\n if np.dot(main_direction, start_to_end) < 0:\n main_direction = -main_direction\n\n else:\n # Default direction for single pose or insufficient data\n main_direction = np.array([1, 0, 0])\n\n # Calculate optimal camera distance\n # Based on trajectory range and field of view, using smaller factor for better screen filling\n fov_factor = (\n 0.8 # Reduced field of view factor to make trajectory occupy more screen space\n )\n base_distance = max_range * fov_factor\n\n # Consider trajectory aspect ratio and adjust distance accordingly\n aspect_ratios = ranges / max_range\n distance_scale = 1.0 + 0.1 * np.std(\n aspect_ratios\n ) # Reduced distance adjustment magnitude\n camera_distance = base_distance * distance_scale\n\n # Calculate optimal camera position\n # Method 1: Diagonal viewing angle based on main direction\n up_vector = np.array([0, 0, 1]) # World up direction (Z-axis)\n\n # Adjust strategy if main direction is nearly vertical\n if abs(np.dot(main_direction, up_vector)) > 0.9:\n # Main direction is nearly vertical, use side view\n view_direction = np.cross(main_direction, np.array([1, 0, 0]))\n if np.linalg.norm(view_direction) < 0.1:\n view_direction = np.cross(main_direction, np.array([0, 1, 0]))\n view_direction = view_direction / np.linalg.norm(view_direction)\n else:\n # Calculate diagonal view direction perpendicular to main direction\n # Combine horizontal component of main direction with tilt angle\n horizontal_component = (\n main_direction - np.dot(main_direction, up_vector) * up_vector\n )\n horizontal_component = horizontal_component / (\n np.linalg.norm(horizontal_component) + 1e-8\n )\n\n # Add some tilt angles for better 3D perspective\n elevation_angle = np.pi / 6 # 30 degrees elevation angle\n azimuth_offset = np.pi / 4 # 45 degrees azimuth offset\n\n # Create tilted view direction for optimal 3D perspective\n view_direction = (\n horizontal_component * np.cos(azimuth_offset) * np.cos(elevation_angle)\n + np.cross(horizontal_component, up_vector)\n * np.sin(azimuth_offset)\n * np.cos(elevation_angle)\n + up_vector * np.sin(elevation_angle)\n )\n\n # Calculate camera eye position\n camera_eye = trajectory_center + view_direction * camera_distance\n\n # Fine-tune camera position to ensure entire trajectory is within view\n # Calculate vectors from camera position to all trajectory points\n view_vectors = centers_world - camera_eye\n view_distances = np.linalg.norm(view_vectors, axis=1)\n\n # Adjust camera distance moderately if some points are too close\n min_distance = camera_distance * 0.3 # Reduced minimum distance ratio\n if np.min(view_distances) < min_distance:\n distance_adjustment = min_distance / np.min(view_distances)\n # Limit adjustment magnitude to avoid excessive scaling\n distance_adjustment = min(\n distance_adjustment, 1.2\n ) # Further limit adjustment range\n camera_eye = (\n trajectory_center + view_direction * camera_distance * distance_adjustment\n )\n\n # Calculate adaptive parameters with appropriate proportions\n auto_vis_depth = max_range * 0.08 # Moderately reduced camera frustum size\n auto_center_size = max_range * 1.5 # Moderately reduced center point size\n\n # Ensure parameters are within reasonable bounds\n auto_vis_depth = max(0.01, min(auto_vis_depth, max_range * 0.2))\n auto_center_size = max(0.1, min(auto_center_size, max_range * 2.0))\n\n return {\n \"camera_eye\": camera_eye,\n \"trajectory_center\": trajectory_center,\n \"auto_vis_depth\": auto_vis_depth,\n \"auto_center_size\": auto_center_size,\n \"max_range\": max_range,\n \"ranges\": ranges,\n \"main_direction\": main_direction,\n }\n\n\ndef compute_multiple_camera_views(poses):\n \"\"\"\n Compute multiple optimized camera view angles, providing different viewing options.\n \"\"\"\n base_params = compute_optimal_camera_view(poses)\n\n trajectory_center = base_params[\"trajectory_center\"]\n max_range = base_params[\"max_range\"]\n main_direction = base_params[\"main_direction\"]\n\n # Calculate multiple view options\n views = {}\n\n # 1. Best automatic view (original optimal view)\n views[\"optimal\"] = base_params\n\n # 2. Top-down bird's eye view\n top_distance = max_range * 1.5 # Further reduced top-down view distance\n views[\"top\"] = {\n **base_params,\n \"camera_eye\": trajectory_center + np.array([0, 0, top_distance]),\n \"description\": \"Top-down view\",\n }\n\n # 3. Side view perspective\n side_distance = max_range * 1.3 # Further reduced side view distance\n side_direction = np.cross(main_direction, np.array([0, 0, 1]))\n if np.linalg.norm(side_direction) < 0.1:\n side_direction = np.array([1, 0, 0])\n else:\n side_direction = side_direction / np.linalg.norm(side_direction)\n\n views[\"side\"] = {\n **base_params,\n \"camera_eye\": trajectory_center + side_direction * side_distance,\n \"description\": \"Side view\",\n }\n\n # 4. Diagonal view (45-degree elevation)\n diagonal_distance = max_range * 1.4 # Further reduced diagonal view distance\n elevation = np.pi / 4 # 45 degrees elevation\n azimuth = np.pi / 4 # 45 degrees azimuth angle\n\n diagonal_direction = np.array(\n [\n np.cos(elevation) * np.cos(azimuth),\n np.cos(elevation) * np.sin(azimuth),\n np.sin(elevation),\n ]\n )\n\n views[\"diagonal\"] = {\n **base_params,\n \"camera_eye\": trajectory_center + diagonal_direction * diagonal_distance,\n \"description\": \"Diagonal view (45° elevation)\",\n }\n\n # 5. Trajectory start-oriented view\n if len(poses) > 1:\n start_to_center = trajectory_center - base_params[\"camera_eye\"]\n start_distance = max_range * 1.2 # Further reduced start view distance\n start_direction = start_to_center / (np.linalg.norm(start_to_center) + 1e-8)\n\n views[\"trajectory_start\"] = {\n **base_params,\n \"camera_eye\": trajectory_center + start_direction * start_distance,\n \"description\": \"View from trajectory start direction\",\n }\n\n # 6. Compact view - ensure entire trajectory is fully visible\n fit_distance = max_range * 0.6 # Very compact distance for close-up view\n fit_direction = np.array([0.7, 0.7, 0.5]) # Stable viewing direction\n fit_direction = fit_direction / np.linalg.norm(fit_direction)\n\n views[\"fit_all\"] = {\n **base_params,\n \"camera_eye\": trajectory_center + fit_direction * fit_distance,\n \"description\": \"Fit all trajectory in view\",\n }\n\n return views\n\n\ndef add_view_selector_to_html(html_str, views):\n \"\"\"\n Add interactive view selector to HTML visualization.\n\n This function injects JavaScript code into the HTML to provide an interactive\n interface for switching between different camera views and enabling auto-rotation.\n\n Args:\n html_str: Original HTML string containing the Plotly visualization\n views: Dictionary of view configurations\n\n Returns:\n str: Enhanced HTML string with view selector and controls\n \"\"\"\n\n # Generate JavaScript code for view selector\n view_selector_js = \"\"\"\n
\n \n
\n \n \n \"\"\"\n\n # Add view selector to the beginning of HTML\n return view_selector_js + html_str\n\n\ndef write_html(poses, file, vis_depth=1, xyz_length=0.2, center_size=0.01, xyz_width=2):\n \"\"\"\n Write camera pose visualization to HTML file with optimized camera view.\n \"\"\"\n # Calculate basic optimal view parameters\n base_view = compute_optimal_camera_view(poses)\n\n # Extract trajectory information\n trajectory_center = base_view[\"trajectory_center\"]\n max_range = base_view[\"max_range\"]\n ranges = base_view[\"ranges\"]\n auto_vis_depth = base_view[\"auto_vis_depth\"]\n auto_center_size = base_view[\"auto_center_size\"]\n\n # Calculate optimal view to see entire trajectory\n # Use larger distance to ensure entire trajectory is visible with better angles\n optimal_distance = (\n max_range * 1.8 * 10\n ) # Increase distance by 10x for better overall view\n\n # Choose ideal angle that can see the full trajectory\n # Use combination of 45-degree elevation and azimuth for good 3D perspective\n elevation = np.pi / 4 # 45-degree elevation angle\n azimuth = np.pi / 4 # 45-degree azimuth angle\n\n # Calculate optimal viewing direction\n optimal_direction = np.array(\n [\n np.cos(elevation) * np.cos(azimuth),\n np.cos(elevation) * np.sin(azimuth),\n np.sin(elevation),\n ]\n )\n\n # Calculate optimal camera position\n camera_eye = trajectory_center + optimal_direction * optimal_distance\n\n # Verify view coverage - ensure all trajectory points are within reasonable distance\n centers_cam = np.zeros([len(poses), 1, 3])\n centers_world = cam2world(centers_cam, poses)[:, 0]\n\n # Calculate distances from optimal camera position to all trajectory points\n distances_to_points = np.linalg.norm(centers_world - camera_eye, axis=1)\n max_distance_to_point = np.max(distances_to_points)\n min_distance_to_point = np.min(distances_to_points)\n\n # If distance variation is too large, the view might not be ideal, adjust accordingly\n if max_distance_to_point / min_distance_to_point > 3.0:\n # Recalculate more balanced distance\n optimal_distance = max_range * 2.2 * 10 # Further increase distance (10x)\n camera_eye = trajectory_center + optimal_direction * optimal_distance\n\n # Create view dictionary with only optimal view for Auto Rotate\n views = {\n \"fit_all\": {\n \"camera_eye\": camera_eye,\n \"trajectory_center\": trajectory_center,\n \"auto_vis_depth\": auto_vis_depth,\n \"auto_center_size\": auto_center_size,\n \"max_range\": max_range,\n \"ranges\": ranges,\n \"description\": \"Optimal view to see entire trajectory\",\n }\n }\n\n print(f\"Trajectory ranges: x={ranges[0]:.3f}, y={ranges[1]:.3f}, z={ranges[2]:.3f}\")\n print(f\"Max range: {max_range:.3f}\")\n print(f\"Auto vis_depth: {auto_vis_depth:.3f}, center_size: {auto_center_size:.3f}\")\n print(\n f\"Trajectory center: ({trajectory_center[0]:.3f}, {trajectory_center[1]:.3f}, {trajectory_center[2]:.3f})\"\n )\n print(\n f\"Optimal camera position for full trajectory view: ({camera_eye[0]:.3f}, {camera_eye[1]:.3f}, {camera_eye[2]:.3f})\"\n )\n print(f\"Camera distance from trajectory center: {optimal_distance:.3f}\")\n print(\n f\"Distance range to trajectory points: {min_distance_to_point:.3f} - {max_distance_to_point:.3f}\"\n )\n\n xyz_length = xyz_length / 3\n xyz_width = xyz_width\n vis_depth = auto_vis_depth # Use automatically computed depth\n center_size = auto_center_size # Use automatically computed size\n\n traces_poses = plotly_visualize_pose(\n poses,\n vis_depth=vis_depth,\n xyz_length=xyz_length,\n center_size=center_size,\n xyz_width=xyz_width,\n mesh_opacity=0.05,\n )\n traces_all2 = traces_poses\n layout2 = go.Layout(\n scene=dict(\n xaxis=dict(visible=False),\n yaxis=dict(visible=False),\n zaxis=dict(visible=False),\n dragmode=\"orbit\",\n aspectratio=dict(x=1, y=1, z=1),\n aspectmode=\"data\",\n # Set initial camera view to fully see the trajectory with optimized positioning\n camera=dict(\n eye=dict(x=camera_eye[0], y=camera_eye[1], z=camera_eye[2]),\n center=dict(\n x=trajectory_center[0],\n y=trajectory_center[1],\n z=trajectory_center[2],\n ),\n up=dict(x=0, y=0, z=1),\n ),\n ),\n height=800,\n width=1200,\n showlegend=False,\n )\n\n fig2 = go.Figure(data=traces_all2, layout=layout2)\n html_str2 = pio.to_html(fig2, full_html=False)\n\n # Add real-time camera view display functionality\n camera_info_html = \"\"\"\n
\n

Camera Info

\n
Eye:
\n
x: 2.000
\n
y: 2.000
\n
z: 1.000
\n
\n
Center:
\n
x: 0.000
\n
y: 0.000
\n
z: 0.000
\n
\n
Up:
\n
x: 0.000
\n
y: 0.000
\n
z: 1.000
\n
\n \n
\n \n \n \"\"\"\n\n # Add view selector and camera info to HTML\n enhanced_html = add_view_selector_to_html(camera_info_html + html_str2, views)\n\n file.write(enhanced_html)\n\n print(f\"Enhanced visualized poses are saved to {file.name}\")\n # Removed redundant view options printing\n\n\ndef plotly_visualize_pose_animated(\n poses_full,\n vis_depth=0.5,\n xyz_length=0.5,\n center_size=2,\n xyz_width=5,\n mesh_opacity=0.05,\n):\n \"\"\"\n Create plotly visualization traces for camera poses, frame by frame for animation.\n Now shows the full trajectory with future poses as completely transparent.\n \"\"\"\n N_total = len(poses_full)\n plotly_frames = []\n\n # Pre-compute data for all poses to ensure consistent layout\n centers_cam = np.zeros([N_total, 1, 3])\n centers_world = cam2world(centers_cam, poses_full)\n centers_world = centers_world[:, 0]\n # Get the camera wireframes for all poses\n vertices, faces, wireframe = get_camera_mesh(poses_full, depth=vis_depth)\n vertices_merged, faces_merged = merge_meshes(vertices, faces)\n wireframe_merged = merge_wireframes_plotly(wireframe)\n # Break up (x,y,z) coordinates.\n wireframe_x, wireframe_y, wireframe_z = unbind_np(wireframe_merged, axis=-1)\n centers_x, centers_y, centers_z = unbind_np(centers_world, axis=-1)\n vertices_x, vertices_y, vertices_z = unbind_np(vertices_merged, axis=-1)\n\n # Initial frame showing all poses with appropriate transparency\n initial_data = []\n\n for i in tqdm(range(1, N_total + 1), desc=\"Generating animation frames\"):\n current_frame = i - 1 # Current frame index (0-based)\n\n # Set the color map for the camera trajectory\n color_map = plt.get_cmap(\"gist_rainbow\")\n center_color = []\n faces_merged_color = []\n wireframe_color = []\n\n for k in range(N_total): # Process all poses\n # Set the camera pose colors (with a smooth gradient color map).\n r, g, b, _ = color_map(k / (N_total - 1))\n rgb = np.array([r, g, b]) * 0.8\n\n # Set transparency based on current frame\n if k < current_frame: # Past poses - visible with reduced opacity\n # Set transparency based on temporal distance, more distant = more transparent\n time_distance = (current_frame - k) / max(current_frame, 1)\n alpha = 0.15 + 0.25 * (1 - time_distance) # Transparency range 0.15-0.4\n wireframe_alpha = alpha\n mesh_alpha = alpha * 0.4\n elif k == current_frame: # Current pose - fully visible\n alpha = 0.8 # Fully opaque, dark display\n wireframe_alpha = 0.8\n mesh_alpha = 0.6\n else: # Future poses - completely transparent\n alpha = 0.0 # Completely transparent\n wireframe_alpha = 0.0\n mesh_alpha = 0.0\n\n # Set colors and transparency\n wireframe_color += [np.concatenate([rgb, [wireframe_alpha]])] * 11\n center_color += [np.concatenate([rgb, [alpha]])]\n faces_merged_color += [np.concatenate([rgb, [mesh_alpha]])] * 6\n\n frame_data = [\n go.Scatter3d(\n x=wireframe_x,\n y=wireframe_y,\n z=wireframe_z,\n mode=\"lines\",\n line=dict(color=wireframe_color, width=1),\n ),\n go.Scatter3d(\n x=centers_x,\n y=centers_y,\n z=centers_z,\n mode=\"markers\",\n marker=dict(color=center_color, size=center_size),\n ),\n go.Mesh3d(\n x=vertices_x,\n y=vertices_y,\n z=vertices_z,\n i=[f[0] for f in faces_merged],\n j=[f[1] for f in faces_merged],\n k=[f[2] for f in faces_merged],\n facecolor=faces_merged_color,\n opacity=0.6, # Set base opacity for mesh\n ),\n ]\n\n if i == 1: # Set initial data for the first frame\n initial_data = frame_data\n\n plotly_frames.append(go.Frame(data=frame_data, name=str(i)))\n\n return initial_data, plotly_frames\n\n\ndef write_html_animated(\n poses, file, vis_depth=1, xyz_length=0.2, center_size=0.01, xyz_width=2\n):\n \"\"\"\n Write camera pose visualization with animation to HTML file with optimized camera view.\n \"\"\"\n # Calculate basic optimal view parameters\n base_view = compute_optimal_camera_view(poses)\n\n # Extract trajectory information\n trajectory_center = base_view[\"trajectory_center\"]\n max_range = base_view[\"max_range\"]\n ranges = base_view[\"ranges\"]\n auto_vis_depth = base_view[\"auto_vis_depth\"]\n auto_center_size = base_view[\"auto_center_size\"]\n\n # Calculate optimal view to see entire trajectory\n # Use larger distance to ensure entire trajectory is visible with better angles\n optimal_distance = (\n max_range * 1.8 * 10\n ) # Increase distance by 10x for better overall view\n\n # Choose ideal angle that can see the full trajectory\n # Use combination of 45-degree elevation and azimuth for good 3D perspective\n elevation = np.pi / 4 # 45-degree elevation angle\n azimuth = np.pi / 4 # 45-degree azimuth angle\n\n # Calculate optimal viewing direction\n optimal_direction = np.array(\n [\n np.cos(elevation) * np.cos(azimuth),\n np.cos(elevation) * np.sin(azimuth),\n np.sin(elevation),\n ]\n )\n\n # Calculate optimal camera position\n camera_eye = trajectory_center + optimal_direction * optimal_distance\n\n # Verify view coverage - ensure all trajectory points are within reasonable distance\n centers_cam = np.zeros([len(poses), 1, 3])\n centers_world = cam2world(centers_cam, poses)[:, 0]\n\n # Calculate distances from optimal camera position to all trajectory points\n distances_to_points = np.linalg.norm(centers_world - camera_eye, axis=1)\n max_distance_to_point = np.max(distances_to_points)\n min_distance_to_point = np.min(distances_to_points)\n\n # If distance variation is too large, the view might not be ideal, adjust accordingly\n if max_distance_to_point / min_distance_to_point > 3.0:\n # Recalculate more balanced distance\n optimal_distance = max_range * 2.2 * 10 # Further increase distance (10x)\n camera_eye = trajectory_center + optimal_direction * optimal_distance\n\n # Adjust parameters for animation\n xyz_length = xyz_length / 3\n xyz_width = xyz_width\n vis_depth = auto_vis_depth # Use automatically computed depth\n center_size = auto_center_size # Use automatically computed size\n\n print(\n f\"Animation - Trajectory ranges: x={ranges[0]:.3f}, y={ranges[1]:.3f}, z={ranges[2]:.3f}\"\n )\n print(f\"Animation - Max range: {max_range:.3f}\")\n print(\n f\"Animation - Auto vis_depth: {auto_vis_depth:.3f}, center_size: {auto_center_size:.3f}\"\n )\n print(\n f\"Animation - Trajectory center: ({trajectory_center[0]:.3f}, {trajectory_center[1]:.3f}, {trajectory_center[2]:.3f})\"\n )\n print(\n f\"Animation - Optimal camera position for full trajectory view: ({camera_eye[0]:.3f}, {camera_eye[1]:.3f}, {camera_eye[2]:.3f})\"\n )\n print(f\"Animation - Camera distance from trajectory center: {optimal_distance:.3f}\")\n print(\n f\"Animation - Distance range to trajectory points: {min_distance_to_point:.3f} - {max_distance_to_point:.3f}\"\n )\n\n initial_data, plotly_frames = plotly_visualize_pose_animated(\n poses,\n vis_depth=vis_depth,\n xyz_length=xyz_length,\n center_size=center_size,\n xyz_width=xyz_width,\n mesh_opacity=0.05,\n )\n\n layout = go.Layout(\n scene=dict(\n xaxis=dict(visible=False),\n yaxis=dict(visible=False),\n zaxis=dict(visible=False),\n dragmode=\"orbit\",\n aspectratio=dict(x=1, y=1, z=1),\n aspectmode=\"data\",\n # Use optimized camera view settings (same 10x distance as write_html)\n camera=dict(\n eye=dict(x=camera_eye[0], y=camera_eye[1], z=camera_eye[2]),\n center=dict(\n x=trajectory_center[0],\n y=trajectory_center[1],\n z=trajectory_center[2],\n ),\n up=dict(x=0, y=0, z=1),\n ),\n ),\n height=800, # Increased height for better animation display\n width=1200, # Increased width for better animation display\n showlegend=False,\n updatemenus=[\n dict(\n type=\"buttons\",\n buttons=[\n dict(\n label=\"Play\",\n method=\"animate\",\n args=[\n None,\n {\n \"frame\": {\"duration\": 50, \"redraw\": True},\n \"fromcurrent\": True,\n \"transition\": {\"duration\": 0},\n },\n ],\n )\n ],\n )\n ],\n )\n\n fig = go.Figure(data=initial_data, layout=layout, frames=plotly_frames)\n html_str = pio.to_html(fig, full_html=False)\n file.write(html_str)\n\n print(f\"Visualized poses are saved to {file}\")\n\n\ndef quaternion_to_matrix(quaternions, eps: float = 1e-8):\n \"\"\"\n Convert 4-dimensional quaternions to 3x3 rotation matrices.\n\n Reference:\n https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py\n \"\"\"\n\n # Order changed to match scipy format: (i, j, k, r)\n i, j, k, r = torch.unbind(quaternions, dim=-1)\n two_s = 2 / ((quaternions * quaternions).sum(dim=-1) + eps)\n\n # Construct rotation matrix elements using quaternion algebra\n o = torch.stack(\n (\n 1 - two_s * (j * j + k * k), # R[0,0]\n two_s * (i * j - k * r), # R[0,1]\n two_s * (i * k + j * r), # R[0,2]\n two_s * (i * j + k * r), # R[1,0]\n 1 - two_s * (i * i + k * k), # R[1,1]\n two_s * (j * k - i * r), # R[1,2]\n two_s * (i * k - j * r), # R[2,0]\n two_s * (j * k + i * r), # R[2,1]\n 1 - two_s * (i * i + j * j), # R[2,2]\n ),\n -1,\n )\n return einops.rearrange(o, \"... (i j) -> ... i j\", i=3, j=3)\n\n\ndef pose_from_quaternion(pose):\n \"\"\"\n Convert quaternion-based pose representation to 4x4 transformation matrices.\n\n Reference:\n https://github.com/pointrix-project/Geomotion/blob/6ab0c364f1b44ab4ea190085dbf068f62b42727c/geomotion/model/cameras.py#L6\n \"\"\"\n # Convert numpy array to torch tensor if needed\n if type(pose) == np.ndarray:\n pose = torch.tensor(pose)\n # Add batch dimension if input is 1D\n if len(pose.shape) == 1:\n pose = pose[None]\n # Extract translation and quaternion components\n quat_t = pose[..., :3] # Translation components [tx, ty, tz]\n quat_r = pose[..., 3:] # Quaternion components [qi, qj, qk, qr]\n\n # Initialize world-to-camera transformation matrix\n w2c_matrix = torch.zeros((*list(pose.shape)[:-1], 3, 4), device=pose.device)\n w2c_matrix[..., :3, 3] = quat_t # Set translation part\n w2c_matrix[..., :3, :3] = quaternion_to_matrix(quat_r) # Set rotation part\n return w2c_matrix\n\n\ndef viz_poses(i, pth, file, args):\n \"\"\"\n Visualize camera poses for a sequence and write to HTML file.\n \"\"\"\n file.write(f\"{i} {pth}
\")\n\n # Load pose data from file\n pose = np.load(pth)\n\n # Convert quaternion poses to transformation matrices\n poses = pose_from_quaternion(pose) # Input: (N,7), Output: (N,3,4) w2c matrices\n poses = poses.cpu().numpy()\n\n # Scale camera positions to reduce distance between camera frustums for better visualization\n scale_factor = getattr(\n args, \"scale_factor\", 0.3\n ) # Default scale factor 0.3, adjustable via command line parameter\n\n # Apply scaling to translation part (camera positions) while keeping rotation unchanged\n # Create scaled copy of poses\n poses_scaled = poses.copy()\n poses_scaled[..., :3, 3] = poses[..., :3, 3] * scale_factor\n\n print(f\"Original poses shape: {poses.shape}\")\n print(f\"Applied scale factor: {scale_factor}\")\n\n # Generate visualization based on dynamic flag\n if args.dynamic:\n write_html_animated(poses_scaled, file, vis_depth=args.vis_depth)\n else:\n write_html(poses_scaled, file, vis_depth=args.vis_depth)\n\n\nif __name__ == \"__main__\":\n # Set up command-line argument parser\n parser = argparse.ArgumentParser(\n description=\"Visualize camera poses with interactive 3D plots\",\n formatter_class=argparse.RawDescriptionHelpFormatter,\n )\n\n parser.add_argument(\n \"--datas\",\n type=str,\n nargs=\"+\",\n required=True,\n help=\"List of pose file paths (.npy format) to visualize.\",\n )\n parser.add_argument(\n \"--vis_depth\",\n type=float,\n default=0.2,\n help=\"Depth of camera frustum visualization (default: 0.2).\",\n )\n parser.add_argument(\n \"--scale_factor\",\n type=float,\n default=0.3,\n help=\"Scale factor to reduce distance between cameras - smaller values bring cameras closer together (default: 0.3).\",\n )\n parser.add_argument(\n \"--outdir\",\n type=str,\n default=\"./visualize\",\n help=\"Output directory to save HTML visualization files (default: ./visualize).\",\n )\n parser.add_argument(\n \"--dynamic\",\n action=\"store_true\",\n help=\"Create animated visualization showing camera trajectory progression over time.\",\n )\n\n # Parse command-line arguments\n args = parser.parse_args()\n\n # Create output directory and process pose files\n os.makedirs(args.outdir, exist_ok=True)\n\n print(f\"Processing {len(args.datas)} pose file(s)...\")\n print(f\"Output directory: {args.outdir}\")\n print(f\"Visualization type: {'Animated' if args.dynamic else 'Static'}\")\n\n with open(f\"{args.outdir}/visualize.html\", \"w\") as file:\n for i, pth in enumerate(tqdm(args.datas, desc=\"Processing pose files\")):\n if not os.path.exists(pth):\n print(f\"Warning: Path {pth} does not exist, skipping.\")\n continue\n print(f\"Processing: {pth} (#{i+1})\")\n viz_poses(i, pth, file, args)\n\n print(\n f\"Visualization complete! Open {args.outdir}/visualize.html in your browser to view results.\"\n )\n" } ]