diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..cf1afbfd4b598bb8ed20ec4648f70f88f4347d30 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +examples/* filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..83c4ef79ab58dd3f450172e283531970e9fd6257 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +checkpoints/ +.gradio/ \ No newline at end of file diff --git a/MoGe/.gitignore b/MoGe/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..eee13ac82659361f18f3595742fd0355de307e39 --- /dev/null +++ b/MoGe/.gitignore @@ -0,0 +1,423 @@ +## Ignore Visual Studio temporary files, build results, and +## files generated by popular Visual Studio add-ons. +## +## Get latest from https://github.com/github/gitignore/blob/main/VisualStudio.gitignore + +# User-specific files +*.rsuser +*.suo +*.user +*.userosscache +*.sln.docstates + +# User-specific files (MonoDevelop/Xamarin Studio) +*.userprefs + +# Mono auto generated files +mono_crash.* + +# Build results +[Dd]ebug/ +[Dd]ebugPublic/ +[Rr]elease/ +[Rr]eleases/ +x64/ +x86/ +[Ww][Ii][Nn]32/ +[Aa][Rr][Mm]/ +[Aa][Rr][Mm]64/ +bld/ +[Bb]in/ +[Oo]bj/ +[Ll]og/ +[Ll]ogs/ + +# Visual Studio 2015/2017 cache/options directory +.vs/ +# Uncomment if you have tasks that create the project's static files in wwwroot +#wwwroot/ + +# Visual Studio 2017 auto generated files +Generated\ Files/ + +# MSTest test Results +[Tt]est[Rr]esult*/ +[Bb]uild[Ll]og.* + +# NUnit +*.VisualState.xml +TestResult.xml +nunit-*.xml + +# Build Results of an ATL Project +[Dd]ebugPS/ +[Rr]eleasePS/ +dlldata.c + +# Benchmark Results +BenchmarkDotNet.Artifacts/ + +# .NET Core +project.lock.json +project.fragment.lock.json +artifacts/ + +# ASP.NET Scaffolding +ScaffoldingReadMe.txt + +# StyleCop +StyleCopReport.xml + +# Files built by Visual Studio +*_i.c +*_p.c +*_h.h +*.ilk +*.meta +*.obj +*.iobj +*.pch +*.pdb +*.ipdb +*.pgc +*.pgd +*.rsp +*.sbr +*.tlb +*.tli +*.tlh +*.tmp +*.tmp_proj +*_wpftmp.csproj +*.log +*.tlog +*.vspscc +*.vssscc +.builds +*.pidb +*.svclog +*.scc + +# Chutzpah Test files +_Chutzpah* + +# Visual C++ cache files +ipch/ +*.aps +*.ncb +*.opendb +*.opensdf +*.sdf +*.cachefile +*.VC.db +*.VC.VC.opendb + +# Visual Studio profiler +*.psess +*.vsp +*.vspx +*.sap + +# Visual Studio Trace Files +*.e2e + +# TFS 2012 Local Workspace +$tf/ + +# Guidance Automation Toolkit +*.gpState + +# ReSharper is a .NET coding add-in +_ReSharper*/ +*.[Rr]e[Ss]harper +*.DotSettings.user + +# TeamCity is a build add-in +_TeamCity* + +# DotCover is a Code Coverage Tool +*.dotCover + +# AxoCover is a Code Coverage Tool +.axoCover/* +!.axoCover/settings.json + +# Coverlet is a free, cross platform Code Coverage Tool +coverage*.json +coverage*.xml +coverage*.info + +# Visual Studio code coverage results +*.coverage +*.coveragexml + +# NCrunch +_NCrunch_* +.*crunch*.local.xml +nCrunchTemp_* + +# MightyMoose +*.mm.* +AutoTest.Net/ + +# Web workbench (sass) +.sass-cache/ + +# Installshield output folder +[Ee]xpress/ + +# DocProject is a documentation generator add-in +DocProject/buildhelp/ +DocProject/Help/*.HxT +DocProject/Help/*.HxC +DocProject/Help/*.hhc +DocProject/Help/*.hhk +DocProject/Help/*.hhp +DocProject/Help/Html2 +DocProject/Help/html + +# Click-Once directory +publish/ + +# Publish Web Output +*.[Pp]ublish.xml +*.azurePubxml +# Note: Comment the next line if you want to checkin your web deploy settings, +# but database connection strings (with potential passwords) will be unencrypted +*.pubxml +*.publishproj + +# Microsoft Azure Web App publish settings. Comment the next line if you want to +# checkin your Azure Web App publish settings, but sensitive information contained +# in these scripts will be unencrypted +PublishScripts/ + +# NuGet Packages +*.nupkg +# NuGet Symbol Packages +*.snupkg +# The packages folder can be ignored because of Package Restore +**/[Pp]ackages/* +# except build/, which is used as an MSBuild target. +!**/[Pp]ackages/build/ +# Uncomment if necessary however generally it will be regenerated when needed +#!**/[Pp]ackages/repositories.config +# NuGet v3's project.json files produces more ignorable files +*.nuget.props +*.nuget.targets + +# Microsoft Azure Build Output +csx/ +*.build.csdef + +# Microsoft Azure Emulator +ecf/ +rcf/ + +# Windows Store app package directories and files +AppPackages/ +BundleArtifacts/ +Package.StoreAssociation.xml +_pkginfo.txt +*.appx +*.appxbundle +*.appxupload + +# Visual Studio cache files +# files ending in .cache can be ignored +*.[Cc]ache +# but keep track of directories ending in .cache +!?*.[Cc]ache/ + +# Others +ClientBin/ +~$* +*~ +*.dbmdl +*.dbproj.schemaview +*.jfm +*.pfx +*.publishsettings +orleans.codegen.cs + +# Including strong name files can present a security risk +# (https://github.com/github/gitignore/pull/2483#issue-259490424) +#*.snk + +# Since there are multiple workflows, uncomment next line to ignore bower_components +# (https://github.com/github/gitignore/pull/1529#issuecomment-104372622) +#bower_components/ + +# RIA/Silverlight projects +Generated_Code/ + +# Backup & report files from converting an old project file +# to a newer Visual Studio version. Backup files are not needed, +# because we have git ;-) +_UpgradeReport_Files/ +Backup*/ +UpgradeLog*.XML +UpgradeLog*.htm +ServiceFabricBackup/ +*.rptproj.bak + +# SQL Server files +*.mdf +*.ldf +*.ndf + +# Business Intelligence projects +*.rdl.data +*.bim.layout +*.bim_*.settings +*.rptproj.rsuser +*- [Bb]ackup.rdl +*- [Bb]ackup ([0-9]).rdl +*- [Bb]ackup ([0-9][0-9]).rdl + +# Microsoft Fakes +FakesAssemblies/ + +# GhostDoc plugin setting file +*.GhostDoc.xml + +# Node.js Tools for Visual Studio +.ntvs_analysis.dat +node_modules/ + +# Visual Studio 6 build log +*.plg + +# Visual Studio 6 workspace options file +*.opt + +# Visual Studio 6 auto-generated workspace file (contains which files were open etc.) +*.vbw + +# Visual Studio 6 auto-generated project file (contains which files were open etc.) +*.vbp + +# Visual Studio 6 workspace and project file (working project files containing files to include in project) +*.dsw +*.dsp + +# Visual Studio 6 technical files +*.ncb +*.aps + +# Visual Studio LightSwitch build output +**/*.HTMLClient/GeneratedArtifacts +**/*.DesktopClient/GeneratedArtifacts +**/*.DesktopClient/ModelManifest.xml +**/*.Server/GeneratedArtifacts +**/*.Server/ModelManifest.xml +_Pvt_Extensions + +# Paket dependency manager +.paket/paket.exe +paket-files/ + +# FAKE - F# Make +.fake/ + +# CodeRush personal settings +.cr/personal + +# Python Tools for Visual Studio (PTVS) +__pycache__/ +*.pyc + +# Cake - Uncomment if you are using it +# tools/** +# !tools/packages.config + +# Tabs Studio +*.tss + +# Telerik's JustMock configuration file +*.jmconfig + +# BizTalk build output +*.btp.cs +*.btm.cs +*.odx.cs +*.xsd.cs + +# OpenCover UI analysis results +OpenCover/ + +# Azure Stream Analytics local run output +ASALocalRun/ + +# MSBuild Binary and Structured Log +*.binlog + +# NVidia Nsight GPU debugger configuration file +*.nvuser + +# MFractors (Xamarin productivity tool) working folder +.mfractor/ + +# Local History for Visual Studio +.localhistory/ + +# Visual Studio History (VSHistory) files +.vshistory/ + +# BeatPulse healthcheck temp database +healthchecksdb + +# Backup folder for Package Reference Convert tool in Visual Studio 2017 +MigrationBackup/ + +# Ionide (cross platform F# VS Code tools) working folder +.ionide/ + +# Fody - auto-generated XML schema +FodyWeavers.xsd + +# VS Code files for those working on multiple tools +.vscode/* +!.vscode/settings.json +!.vscode/tasks.json +!.vscode/launch.json +!.vscode/extensions.json +*.code-workspace + +# Local History for Visual Studio Code +.history/ + +# Windows Installer files from build outputs +*.cab +*.msi +*.msix +*.msm +*.msp + +# JetBrains Rider +*.sln.iml + +# Python +*.egg-info/ +/build + +# MoGe +/data* +/download +/extract +/debug +/workspace +/mlruns +/infer_output +/video_output +/eval_output +/.blobcache +/test_images +/test_videos +/vis +/videos +/blobmnt +/eval_dump +/pretrained +/.gradio +/tmp \ No newline at end of file diff --git a/MoGe/CHANGELOG.md b/MoGe/CHANGELOG.md new file mode 100644 index 0000000000000000000000000000000000000000..4eba4b494d8395d49b071bb92ce268253ecc9067 --- /dev/null +++ b/MoGe/CHANGELOG.md @@ -0,0 +1,28 @@ +## 2024-11-28 +### Added +- Supported user-provided camera FOV. See [scripts/infer.py](scripts/infer.py) --fov_x. + - Related issues: [#25](https://github.com/microsoft/MoGe/issues/25) and [#24](https://github.com/microsoft/MoGe/issues/24). +- Added inference scripts for panorama images. See [scripts/infer_panorama.py](scripts/infer_panorama.py). + - Related issue: [#19](https://github.com/microsoft/MoGe/issues/19). + +### Fixed +- Suppressed unnecessary numpy runtime warnings. +- Specified recommended versions of requirements. + - Related issue: [#21](https://github.com/microsoft/MoGe/issues/21). + +### Changed +- Moved `app.py` and `infer.py` to [scripts/](scripts/) +- Improved edge removal. + +## 2025-03-18 +### Added +- Training and evaluation code. See [docs/train.md](docs/train.md) and [docs/eval.md](docs/eval.md). +- Supported installation via pip. Thanks to @fabiencastan and @jgoueslard + for commits in the [#47](https://github.com/microsoft/MoGe/pull/47) +- Supported command-line usage when installed. + +### Changed +- Moved `scripts/` into `moge/` for package installation and command-line usage. +- Renamed `moge.model.moge_model` to `moge.model.v1` for version management. + Now you can import the model class through `from moge.model.v1 import MoGeModel` or `from moge.model import import_model_class_by_version; MoGeModel = import_model_class_by_version('v1')`. +- Exposed `num_tokens` parameter in MoGe model. \ No newline at end of file diff --git a/MoGe/CODE_OF_CONDUCT.md b/MoGe/CODE_OF_CONDUCT.md new file mode 100644 index 0000000000000000000000000000000000000000..f9ba8cf65f3e3104dd061c178066ec8247811f33 --- /dev/null +++ b/MoGe/CODE_OF_CONDUCT.md @@ -0,0 +1,9 @@ +# Microsoft Open Source Code of Conduct + +This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/). + +Resources: + +- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/) +- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) +- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns diff --git a/MoGe/LICENSE b/MoGe/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..3458b5ccd398afed340e17a4d0615c9a8666bb5d --- /dev/null +++ b/MoGe/LICENSE @@ -0,0 +1,224 @@ + MIT License + + Copyright (c) Microsoft Corporation. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/MoGe/README.md b/MoGe/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b17228e3e3472a92eba5042e59877e383a3e0a5a --- /dev/null +++ b/MoGe/README.md @@ -0,0 +1,220 @@ +
+ +# MoGe: Unlocking Accurate Monocular Geometry Estimation for Open-Domain Images with Optimal Training Supervision + +arXiv +Project Page + + +
+ +Method overview + +MoGe is a powerful model for recovering 3D geometry from monocular open-domain images. The model consists of a ViT encoder and a convolutional decoder. It directly predicts an affine-invariant point map as well as a mask that excludes regions with undefined geometry (e.g., sky), from which the camera shift, camera focal length and depth map can be further derived. + +***Check our [website](https://wangrc.site/MoGePage) for videos and interactive results!*** + +## Features + +* **Accurate 3D geometry estimation**: Estimate point maps from single images with high precision. Capable of capturing depth variations up to 1000×, ensuring a comprehensive scene representation. +* **Optional ground-truth FOV input**: Enhance model accuracy further by providing the true field of view. +* **Flexible resolution support**: Works seamlessly with various resolutions and aspect ratios, from 2:1 to 1:2. +* **Optimized for speed**: Achieves <0.1s latency per image on an A100 / RTX 3090 GPU with fp16, and 0.2s with fp32. + +## TODO List + +- [x] Release inference code & ViT-Large model. +- [x] Release evaluation and training code. +- [ ] Release ViT-Base and ViT-Giant models. + +🌟*Updated on 2025/03/18* [CHANGELOG](CHANGELOG.md) + - **Training and evaluation code released!** + - Installation via pip and CLI usage supported. + +## Installation + +### Install via pip + +```bash +pip install git+https://github.com/microsoft/MoGe.git +``` + +### Or clone this repository + +```bash +git clone https://github.com/microsoft/MoGe.git +cd MoGe +``` + +and install the requirements + +```bash +pip install -r requirements.txt +``` + +MoGe should be compatible with most requirements versions. Please check the `requirements.txt` for more details if you have concerns. + +## Usage + +### Pretrained model + +The ViT-Large model has been uploaded to Hugging Face hub at [Ruicheng/moge-vitl](https://huggingface.co/Ruicheng/moge-vitl). +You may load the model via `MoGeModel.from_pretrained("Ruicheng/moge-vitl")` without manually downloading. + +If loading the model from a local file is preferred, you may manually download the model from the huggingface hub and load it via `MoGeModel.from_pretrained("PATH_TO_LOCAL_MODEL.pt")`. + +### Minimal code example + +Here is a minimal example for loading the model and inferring on a single image. + +```python +import cv2 +import torch +from moge.model.v1 import MoGeModel + +device = torch.device("cuda") + +# Load the model from huggingface hub (or load from local). +model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(device) + +# Read the input image and convert to tensor (3, H, W) and normalize to [0, 1] +input_image = cv2.cvtColor(cv2.imread("PATH_TO_IMAGE.jpg"), cv2.COLOR_BGR2RGB) +input_image = torch.tensor(input_image / 255, dtype=torch.float32, device=device).permute(2, 0, 1) + +# Infer +output = model.infer(input_image) +# `output` has keys "points", "depth", "mask" and "intrinsics", +# The maps are in the same size as the input image. +# { +# "points": (H, W, 3), # scale-invariant point map in OpenCV camera coordinate system (x right, y down, z forward) +# "depth": (H, W), # scale-invariant depth map +# "mask": (H, W), # a binary mask for valid pixels. +# "intrinsics": (3, 3), # normalized camera intrinsics +# } +# For more usage details, see the `MoGeModel.infer` docstring. +``` + +### Gradio demo | `moge app` + +The demo is also available at our [Hugging Face space](https://huggingface.co/spaces/Ruicheng/MoGe). + +```bash +# Using the command line tool +moge app + +# In this repo +python moge/scripts/app.py # --share for Gradio public sharing +``` + +See also [`moge/scripts/app.py`](moge/scripts/app.py) + + +### Inference | `moge infer` + +Run the script `moge/scripts/infer.py` via the following command: + +```bash +# Save the output [maps], [glb] and [ply] files +moge infer -i IMAGES_FOLDER_OR_IMAGE_PATH --o OUTPUT_FOLDER --maps --glb --ply + +# Show the result in a window (requires pyglet < 2.0, e.g. pip install pyglet==1.5.29) +moge infer -i IMAGES_FOLDER_OR_IMAGE_PATH --o OUTPUT_FOLDER --show +``` + +For detailed options, run `moge infer --help`: + +``` +Usage: moge infer [OPTIONS] + + Inference script for the MoGe model. + +Options: + -i, --input PATH Input image or folder path. "jpg" and "png" are + supported. + --fov_x FLOAT If camera parameters are known, set the + horizontal field of view in degrees. Otherwise, + MoGe will estimate it. + -o, --output PATH Output folder path + --pretrained TEXT Pretrained model name or path. Defaults to + "Ruicheng/moge-vitl" + --device TEXT Device name (e.g. "cuda", "cuda:0", "cpu"). + Defaults to "cuda" + --fp16 Use fp16 precision for 2x faster inference. + --resize INTEGER Resize the image(s) & output maps to a specific + size. Defaults to None (no resizing). + --resolution_level INTEGER An integer [0-9] for the resolution level for + inference. Higher value means more tokens and + the finer details will be captured, but + inference can be slower. Defaults to 9. Note + that it is irrelevant to the output size, which + is always the same as the input size. + `resolution_level` actually controls + `num_tokens`. See `num_tokens` for more details. + --num_tokens INTEGER number of tokens used for inference. A integer + in the (suggested) range of `[1200, 2500]`. + `resolution_level` will be ignored if + `num_tokens` is provided. Default: None + --threshold FLOAT Threshold for removing edges. Defaults to 0.03. + Smaller value removes more edges. "inf" means no + thresholding. + --maps Whether to save the output maps and fov(image, + depth, mask, points, fov). + --glb Whether to save the output as a.glb file. The + color will be saved as a texture. + --ply Whether to save the output as a.ply file. The + color will be saved as vertex colors. + --show Whether show the output in a window. Note that + this requires pyglet<2 installed as required by + trimesh. + --help Show this message and exit. +``` + +See also [`moge/scripts/infer.py`](moge/scripts/infer.py) + +### 360° panorama images | `moge infer_panorama` + +> *NOTE: This is an experimental extension of MoGe.* + +The script will split the 360-degree panorama image into multiple perspective views and infer on each view separately. +The output maps will be combined to produce a panorama depth map and point map. + +Note that the panorama image must have spherical parameterization (e.g., environment maps or equirectangular images). Other formats must be converted to spherical format before using this script. Run `moge infer_panorama --help` for detailed options. + + +
+ + +The photo is from [this URL](https://commons.wikimedia.org/wiki/Category:360%C2%B0_panoramas_with_equirectangular_projection#/media/File:Braunschweig_Sankt-%C3%84gidien_Panorama_02.jpg) +
+ +See also [`moge/scripts/infer_panorama.py`](moge/scripts/infer_panorama.py) + +## Training & Finetuning + +See [docs/train.md](docs/train.md) + +## Evaluation + +See [docs/eval.md](docs/eval.md) + +## License + +MoGe code is released under the MIT license, except for DINOv2 code in `moge/model/dinov2` which is released by Meta AI under the Apache 2.0 license. +See [LICENSE](LICENSE) for more details. + + +## Citation + +If you find our work useful in your research, we gratefully request that you consider citing our paper: + +``` +@misc{wang2024moge, + title={MoGe: Unlocking Accurate Monocular Geometry Estimation for Open-Domain Images with Optimal Training Supervision}, + author={Wang, Ruicheng and Xu, Sicheng and Dai, Cassie and Xiang, Jianfeng and Deng, Yu and Tong, Xin and Yang, Jiaolong}, + year={2024}, + eprint={2410.19115}, + archivePrefix={arXiv}, + primaryClass={cs.CV}, + url={https://arxiv.org/abs/2410.19115}, +} +``` diff --git a/MoGe/SECURITY.md b/MoGe/SECURITY.md new file mode 100644 index 0000000000000000000000000000000000000000..b3c89efc852e22f71eabf5dfbc6ac62493425eb6 --- /dev/null +++ b/MoGe/SECURITY.md @@ -0,0 +1,41 @@ + + +## Security + +Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet) and [Xamarin](https://github.com/xamarin). + +If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://aka.ms/security.md/definition), please report it to us as described below. + +## Reporting Security Issues + +**Please do not report security vulnerabilities through public GitHub issues.** + +Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://aka.ms/security.md/msrc/create-report). + +If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com). If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://aka.ms/security.md/msrc/pgp). + +You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc). + +Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue: + + * Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.) + * Full paths of source file(s) related to the manifestation of the issue + * The location of the affected source code (tag/branch/commit or direct URL) + * Any special configuration required to reproduce the issue + * Step-by-step instructions to reproduce the issue + * Proof-of-concept or exploit code (if possible) + * Impact of the issue, including how an attacker might exploit the issue + +This information will help us triage your report more quickly. + +If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://aka.ms/security.md/msrc/bounty) page for more details about our active programs. + +## Preferred Languages + +We prefer all communications to be in English. + +## Policy + +Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://aka.ms/security.md/cvd). + + diff --git a/MoGe/SUPPORT.md b/MoGe/SUPPORT.md new file mode 100644 index 0000000000000000000000000000000000000000..291d4d43733f4c15a81ff598ec1c99fd6c18f64c --- /dev/null +++ b/MoGe/SUPPORT.md @@ -0,0 +1,25 @@ +# TODO: The maintainer of this repo has not yet edited this file + +**REPO OWNER**: Do you want Customer Service & Support (CSS) support for this product/project? + +- **No CSS support:** Fill out this template with information about how to file issues and get help. +- **Yes CSS support:** Fill out an intake form at [aka.ms/onboardsupport](https://aka.ms/onboardsupport). CSS will work with/help you to determine next steps. +- **Not sure?** Fill out an intake as though the answer were "Yes". CSS will help you decide. + +*Then remove this first heading from this SUPPORT.MD file before publishing your repo.* + +# Support + +## How to file issues and get help + +This project uses GitHub Issues to track bugs and feature requests. Please search the existing +issues before filing new issues to avoid duplicates. For new issues, file your bug or +feature request as a new Issue. + +For help and questions about using this project, please **REPO MAINTAINER: INSERT INSTRUCTIONS HERE +FOR HOW TO ENGAGE REPO OWNERS OR COMMUNITY FOR HELP. COULD BE A STACK OVERFLOW TAG OR OTHER +CHANNEL. WHERE WILL YOU HELP PEOPLE?**. + +## Microsoft Support Policy + +Support for this **PROJECT or PRODUCT** is limited to the resources listed above. diff --git a/MoGe/baselines/da_v2.py b/MoGe/baselines/da_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..bca560a75514bdfa38c9a28c8d36ea0e006dab1e --- /dev/null +++ b/MoGe/baselines/da_v2.py @@ -0,0 +1,88 @@ +# Reference: https://github.com/DepthAnything/Depth-Anything-V2 +import os +import sys +from typing import * +from pathlib import Path + +import click +import torch +import torch.nn.functional as F +import torchvision.transforms as T +import torchvision.transforms.functional as TF + +from moge.test.baseline import MGEBaselineInterface + + +class Baseline(MGEBaselineInterface): + def __init__(self, repo_path: str, backbone: str, num_tokens: int, device: Union[torch.device, str]): + # Create from repo + repo_path = os.path.abspath(repo_path) + if repo_path not in sys.path: + sys.path.append(repo_path) + if not Path(repo_path).exists(): + raise FileNotFoundError(f'Cannot find the Depth-Anything repository at {repo_path}. Please clone the repository and provide the path to it using the --repo option.') + from depth_anything_v2.dpt import DepthAnythingV2 + + device = torch.device(device) + + # Instantiate model + model = DepthAnythingV2(encoder=backbone, features=256, out_channels=[256, 512, 1024, 1024]) + + # Load checkpoint + checkpoint_path = os.path.join(repo_path, f'checkpoints/depth_anything_v2_{backbone}.pth') + if not os.path.exists(checkpoint_path): + raise FileNotFoundError(f'Cannot find the checkpoint file at {checkpoint_path}. Please download the checkpoint file and place it in the checkpoints directory.') + checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True) + model.load_state_dict(checkpoint) + + model.to(device).eval() + self.model = model + self.num_tokens = num_tokens + self.device = device + + @click.command() + @click.option('--repo', 'repo_path', type=click.Path(), default='../Depth-Anything-V2', help='Path to the Depth-Anything repository.') + @click.option('--backbone', type=click.Choice(['vits', 'vitb', 'vitl']), default='vitl', help='Encoder architecture.') + @click.option('--num_tokens', type=int, default=None, help='Number of tokens to use for the input image.') + @click.option('--device', type=str, default='cuda', help='Device to use for inference.') + @staticmethod + def load(repo_path: str, backbone, num_tokens: int, device: torch.device = 'cuda'): + return Baseline(repo_path, backbone, num_tokens, device) + + @torch.inference_mode() + def infer(self, image: torch.Tensor, intrinsics: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + original_height, original_width = image.shape[-2:] + + assert intrinsics is None, "Depth-Anything-V2 does not support camera intrinsics input" + + if image.ndim == 3: + image = image.unsqueeze(0) + omit_batch_dim = True + else: + omit_batch_dim = False + + if self.num_tokens is None: + resize_factor = 518 / min(original_height, original_width) + expected_width = round(original_width * resize_factor / 14) * 14 + expected_height = round(original_height * resize_factor / 14) * 14 + else: + aspect_ratio = original_width / original_height + tokens_rows = round((self.num_tokens * aspect_ratio) ** 0.5) + tokens_cols = round((self.num_tokens / aspect_ratio) ** 0.5) + expected_width = tokens_cols * 14 + expected_height = tokens_rows * 14 + image = TF.resize(image, (expected_height, expected_width), interpolation=T.InterpolationMode.BICUBIC, antialias=True) + + image = TF.normalize(image, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + disparity = self.model(image) + + disparity = F.interpolate(disparity[:, None], size=(original_height, original_width), mode='bilinear', align_corners=False, antialias=False)[:, 0] + + if omit_batch_dim: + disparity = disparity.squeeze(0) + + return { + 'disparity_affine_invariant': disparity + } + diff --git a/MoGe/baselines/da_v2_metric.py b/MoGe/baselines/da_v2_metric.py new file mode 100644 index 0000000000000000000000000000000000000000..ee4c70d8c6634babf165d2982a692230f5adeac6 --- /dev/null +++ b/MoGe/baselines/da_v2_metric.py @@ -0,0 +1,99 @@ +# Reference https://github.com/DepthAnything/Depth-Anything-V2/metric_depth +import os +import sys +from typing import * +from pathlib import Path + +import click +import torch +import torch.nn.functional as F +import torchvision.transforms as T +import torchvision.transforms.functional as TF +import cv2 + +from moge.test.baseline import MGEBaselineInterface + + +class Baseline(MGEBaselineInterface): + + def __init__(self, repo_path: str, backbone: str, domain: str, num_tokens: int, device: str): + device = torch.device(device) + repo_path = os.path.abspath(repo_path) + if not Path(repo_path).exists(): + raise FileNotFoundError(f'Cannot find the Depth-Anything repository at {repo_path}. Please clone the repository and provide the path to it using the --repo option.') + sys.path.append(os.path.join(repo_path, 'metric_depth')) + from depth_anything_v2.dpt import DepthAnythingV2 + + model_configs = { + 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}, + 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, + 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]} + } + + if domain == 'indoor': + dataset = 'hypersim' + max_depth = 20 + elif domain == 'outdoor': + dataset = 'vkitti' + max_depth = 80 + else: + raise ValueError(f"Invalid domain: {domain}") + + model = DepthAnythingV2(**model_configs[backbone], max_depth=max_depth) + checkpoint_path = os.path.join(repo_path, f'checkpoints/depth_anything_v2_metric_{dataset}_{backbone}.pth') + if not os.path.exists(checkpoint_path): + raise FileNotFoundError(f'Cannot find the checkpoint file at {checkpoint_path}. Please download the checkpoint file and place it in the checkpoints directory.') + model.load_state_dict(torch.load(checkpoint_path, map_location='cpu', weights_only=True)) + model.eval().to(device) + + self.model = model + self.num_tokens = num_tokens + self.device = device + + @click.command() + @click.option('--repo', 'repo_path', type=click.Path(), default='../Depth-Anything-V2', help='Path to the Depth-Anything repository.') + @click.option('--backbone', type=click.Choice(['vits', 'vitb', 'vitl']), default='vitl', help='Backbone architecture.') + @click.option('--domain', type=click.Choice(['indoor', 'outdoor']), help='Domain of the dataset.') + @click.option('--num_tokens', type=int, default=None, help='Number of tokens for the ViT model') + @click.option('--device', type=str, default='cuda', help='Device to use for inference.') + @staticmethod + def load(repo_path: str, backbone: str, domain: str, num_tokens: int, device: str): + return Baseline(repo_path, backbone, domain, num_tokens, device) + + @torch.inference_mode() + def infer(self, image: torch.Tensor, intrinsics: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + original_height, original_width = image.shape[-2:] + + assert intrinsics is None, "Depth-Anything-V2 does not support camera intrinsics input" + + if image.ndim == 3: + image = image.unsqueeze(0) + omit_batch_dim = True + else: + omit_batch_dim = False + + if self.num_tokens is None: + resize_factor = 518 / min(original_height, original_width) + expected_width = round(original_width * resize_factor / 14) * 14 + expected_height = round(original_height * resize_factor / 14) * 14 + else: + aspect_ratio = original_width / original_height + tokens_rows = round((self.num_tokens * aspect_ratio) ** 0.5) + tokens_cols = round((self.num_tokens / aspect_ratio) ** 0.5) + expected_width = tokens_cols * 14 + expected_height = tokens_rows * 14 + image = TF.resize(image, (expected_height, expected_width), interpolation=T.InterpolationMode.BICUBIC, antialias=True) + + image = TF.normalize(image, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + depth = self.model(image) + + depth = F.interpolate(depth[:, None], size=(original_height, original_width), mode='bilinear', align_corners=False, antialias=False)[:, 0] + + if omit_batch_dim: + depth = depth.squeeze(0) + + return { + 'depth_metric': depth + } + diff --git a/MoGe/baselines/metric3d_v2.py b/MoGe/baselines/metric3d_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..661ed5ddfd3bf6f34f53dcb19bcc5a9889a8e377 --- /dev/null +++ b/MoGe/baselines/metric3d_v2.py @@ -0,0 +1,117 @@ +# Reference: https://github.com/YvanYin/Metric3D +import os +import sys +from typing import * + +import click +import torch +import torch.nn.functional as F +import cv2 + +from moge.test.baseline import MGEBaselineInterface + + +class Baseline(MGEBaselineInterface): + def __init__(self, backbone: Literal['vits', 'vitl', 'vitg'], device): + backbone_map = { + 'vits': 'metric3d_vit_small', + 'vitl': 'metric3d_vit_large', + 'vitg': 'metric3d_vit_giant2' + } + + device = torch.device(device) + model = torch.hub.load('yvanyin/metric3d', backbone_map[backbone], pretrain=True) + model.to(device).eval() + + self.model = model + self.device = device + + @click.command() + @click.option('--backbone', type=click.Choice(['vits', 'vitl', 'vitg']), default='vitl', help='Encoder architecture.') + @click.option('--device', type=str, default='cuda', help='Device to use.') + @staticmethod + def load(backbone: str = 'vitl', device: torch.device = 'cuda'): + return Baseline(backbone, device) + + @torch.inference_mode() + def inference_one_image(self, image: torch.Tensor, intrinsics: torch.Tensor = None): + # Reference: https://github.com/YvanYin/Metric3D/blob/main/mono/utils/do_test.py + + # rgb_origin: RGB, 0-255, uint8 + rgb_origin = image.cpu().numpy().transpose((1, 2, 0)) * 255 + + # keep ratio resize + input_size = (616, 1064) # for vit model + h, w = rgb_origin.shape[:2] + scale = min(input_size[0] / h, input_size[1] / w) + rgb = cv2.resize(rgb_origin, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_LINEAR) + if intrinsics is not None: + focal = intrinsics[0, 0] * int(w * scale) + + # padding to input_size + padding = [123.675, 116.28, 103.53] + h, w = rgb.shape[:2] + pad_h = input_size[0] - h + pad_w = input_size[1] - w + pad_h_half = pad_h // 2 + pad_w_half = pad_w // 2 + rgb = cv2.copyMakeBorder(rgb, pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half, cv2.BORDER_CONSTANT, value=padding) + pad_info = [pad_h_half, pad_h - pad_h_half, pad_w_half, pad_w - pad_w_half] + + # normalize rgb + mean = torch.tensor([123.675, 116.28, 103.53]).float()[:, None, None] + std = torch.tensor([58.395, 57.12, 57.375]).float()[:, None, None] + rgb = torch.from_numpy(rgb.transpose((2, 0, 1))).float() + rgb = torch.div((rgb - mean), std) + rgb = rgb[None, :, :, :].cuda() + + # inference + pred_depth, confidence, output_dict = self.model.inference({'input': rgb}) + + # un pad + pred_depth = pred_depth.squeeze() + pred_depth = pred_depth[pad_info[0] : pred_depth.shape[0] - pad_info[1], pad_info[2] : pred_depth.shape[1] - pad_info[3]] + pred_depth = pred_depth.clamp_min(0.5) # clamp to 0.5m, since metric3d could yield very small depth values, resulting in crashed the scale shift alignment. + + # upsample to original size + pred_depth = F.interpolate(pred_depth[None, None, :, :], image.shape[-2:], mode='bilinear').squeeze() + + if intrinsics is not None: + # de-canonical transform + canonical_to_real_scale = focal / 1000.0 # 1000.0 is the focal length of canonical camera + pred_depth = pred_depth * canonical_to_real_scale # now the depth is metric + pred_depth = torch.clamp(pred_depth, 0, 300) + + pred_normal, normal_confidence = output_dict['prediction_normal'].split([3, 1], dim=1) # see https://arxiv.org/abs/2109.09881 for details + + # un pad and resize to some size if needed + pred_normal = pred_normal.squeeze(0) + pred_normal = pred_normal[:, pad_info[0] : pred_normal.shape[1] - pad_info[1], pad_info[2] : pred_normal.shape[2] - pad_info[3]] + + # you can now do anything with the normal + pred_normal = F.interpolate(pred_normal[None, :, :, :], image.shape[-2:], mode='bilinear').squeeze(0) + pred_normal = F.normalize(pred_normal, p=2, dim=0) + + return pred_depth, pred_normal.permute(1, 2, 0) + + @torch.inference_mode() + def infer(self, image: torch.Tensor, intrinsics: torch.Tensor = None): + # image: (B, H, W, 3) or (H, W, 3) + if image.ndim == 3: + pred_depth, pred_normal = self.inference_one_image(image, intrinsics) + else: + for i in range(image.shape[0]): + pred_depth_i, pred_normal_i = self.inference_one_image(image[i], intrinsics[i] if intrinsics is not None else None) + pred_depth.append(pred_depth_i) + pred_normal.append(pred_normal_i) + pred_depth = torch.stack(pred_depth, dim=0) + pred_normal = torch.stack(pred_normal, dim=0) + + if intrinsics is not None: + return { + "depth_metric": pred_depth, + } + else: + return { + "depth_scale_invariant": pred_depth, + } diff --git a/MoGe/baselines/moge.py b/MoGe/baselines/moge.py new file mode 100644 index 0000000000000000000000000000000000000000..9fdfaae24d01649230b107503275318971888409 --- /dev/null +++ b/MoGe/baselines/moge.py @@ -0,0 +1,83 @@ +import os +import sys +from typing import * +import importlib + +import click +import torch +import utils3d + +from moge.test.baseline import MGEBaselineInterface + + +class Baseline(MGEBaselineInterface): + + def __init__(self, num_tokens: int, resolution_level: int, pretrained_model_name_or_path: str, use_fp16: bool, device: str = 'cuda:0', version: str = 'v1'): + super().__init__() + from moge.model import import_model_class_by_version + MoGeModel = import_model_class_by_version(version) + self.version = version + + self.model = MoGeModel.from_pretrained(pretrained_model_name_or_path).to(device).eval() + + self.device = torch.device(device) + self.num_tokens = num_tokens + self.resolution_level = resolution_level + self.use_fp16 = use_fp16 + + @click.command() + @click.option('--num_tokens', type=int, default=None) + @click.option('--resolution_level', type=int, default=9) + @click.option('--pretrained', 'pretrained_model_name_or_path', type=str, default='Ruicheng/moge-vitl') + @click.option('--fp16', 'use_fp16', is_flag=True) + @click.option('--device', type=str, default='cuda:0') + @click.option('--version', type=str, default='v1') + @staticmethod + def load(num_tokens: int, resolution_level: int, pretrained_model_name_or_path: str, use_fp16: bool, device: str = 'cuda:0', version: str = 'v1'): + return Baseline(num_tokens, resolution_level, pretrained_model_name_or_path, use_fp16, device, version) + + # Implementation for inference + @torch.inference_mode() + def infer(self, image: torch.FloatTensor, intrinsics: Optional[torch.FloatTensor] = None): + if intrinsics is not None: + fov_x, _ = utils3d.torch.intrinsics_to_fov(intrinsics) + fov_x = torch.rad2deg(fov_x) + else: + fov_x = None + output = self.model.infer(image, fov_x=fov_x, apply_mask=True, num_tokens=self.num_tokens) + + if self.version == 'v1': + return { + 'points_scale_invariant': output['points'], + 'depth_scale_invariant': output['depth'], + 'intrinsics': output['intrinsics'], + } + else: + return { + 'points_metric': output['points'], + 'depth_metric': output['depth'], + 'intrinsics': output['intrinsics'], + } + + @torch.inference_mode() + def infer_for_evaluation(self, image: torch.FloatTensor, intrinsics: torch.FloatTensor = None): + if intrinsics is not None: + fov_x, _ = utils3d.torch.intrinsics_to_fov(intrinsics) + fov_x = torch.rad2deg(fov_x) + else: + fov_x = None + output = self.model.infer(image, fov_x=fov_x, apply_mask=False, num_tokens=self.num_tokens, use_fp16=self.use_fp16) + + if self.version == 'v1': + return { + 'points_scale_invariant': output['points'], + 'depth_scale_invariant': output['depth'], + 'intrinsics': output['intrinsics'], + } + else: + return { + 'points_metric': output['points'], + 'depth_metric': output['depth'], + 'intrinsics': output['intrinsics'], + } + diff --git a/MoGe/configs/eval/all_benchmarks.json b/MoGe/configs/eval/all_benchmarks.json new file mode 100644 index 0000000000000000000000000000000000000000..94c0fc4605f3a3472d7d39d4d8e40eb9e3d784b7 --- /dev/null +++ b/MoGe/configs/eval/all_benchmarks.json @@ -0,0 +1,78 @@ +{ + "NYUv2": { + "path": "data/eval/NYUv2", + "width": 640, + "height": 480, + "split": ".index.txt", + "depth_unit": 1.0 + }, + "KITTI": { + "path": "data/eval/KITTI", + "width": 750, + "height": 375, + "split": ".index.txt", + "depth_unit": 1 + }, + "ETH3D": { + "path": "data/eval/ETH3D", + "width": 2048, + "height": 1365, + "split": ".index.txt", + "include_segmentation": true, + "depth_unit": 1 + }, + "iBims-1": { + "path": "data/eval/iBims-1", + "width": 640, + "height": 480, + "split": ".index.txt", + "has_sharp_boundary": true, + "include_segmentation": true, + "depth_unit": 1.0 + }, + "GSO": { + "path": "data/eval/GSO", + "width": 512, + "height": 512, + "split": ".index.txt" + }, + "Sintel": { + "path": "data/eval/Sintel", + "width": 872, + "height": 436, + "split": ".index.txt", + "has_sharp_boundary": true, + "include_segmentation": true + }, + "DDAD": { + "path": "data/eval/DDAD", + "width": 1400, + "height": 700, + "include_segmentation": true, + "split": ".index.txt", + "depth_unit": 1.0 + }, + "DIODE": { + "path": "data/eval/DIODE", + "width": 1024, + "height": 768, + "split": ".index.txt", + "include_segmentation": true, + "depth_unit": 1.0 + }, + "Spring": { + "path": "data/eval/Spring", + "width": 1920, + "height": 1080, + "split": ".index.txt", + "has_sharp_boundary": true + }, + "HAMMER": { + "path": "data/eval/HAMMER", + "width": 1664, + "height": 832, + "split": ".index.txt", + "depth_unit": 1, + "has_sharp_boundary": true + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/ddad.json b/MoGe/configs/eval/benchmarks/ddad.json new file mode 100644 index 0000000000000000000000000000000000000000..09dd4d74bbccbb46a4013afd9fee1e717d606a53 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/ddad.json @@ -0,0 +1,9 @@ +{ + "DDAD": { + "path": "data/eval/DDAD", + "width": 1400, + "height": 700, + "include_segmentation": true, + "split": ".index.txt" + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/diode.json b/MoGe/configs/eval/benchmarks/diode.json new file mode 100644 index 0000000000000000000000000000000000000000..679ca6ee13ddf5e5bcab93f453b2f11279781a2f --- /dev/null +++ b/MoGe/configs/eval/benchmarks/diode.json @@ -0,0 +1,9 @@ +{ + "DIODE": { + "path": "data/eval/DIODE", + "width": 1024, + "height": 768, + "split": ".index.txt", + "include_segmentation": true + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/eth3d.json b/MoGe/configs/eval/benchmarks/eth3d.json new file mode 100644 index 0000000000000000000000000000000000000000..88a3a1b291dcde3f2959c0d36d7ebbc33213fc84 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/eth3d.json @@ -0,0 +1,10 @@ +{ + "ETH3D": { + "path": "data/eval/ETH3D", + "width": 2048, + "height": 1365, + "split": ".index.txt", + "include_segmentation": true, + "depth_unit": 1 + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/gso.json b/MoGe/configs/eval/benchmarks/gso.json new file mode 100644 index 0000000000000000000000000000000000000000..ee1aefff7ae3453b0cdddf7ab3369301d2e8d924 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/gso.json @@ -0,0 +1,8 @@ +{ + "GSO": { + "path": "data/eval/GSO", + "width": 512, + "height": 512, + "split": ".index.txt" + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/hammer.json b/MoGe/configs/eval/benchmarks/hammer.json new file mode 100644 index 0000000000000000000000000000000000000000..41838db6bfcf2ea6f3ed230b6c7ee3315ec3fbfe --- /dev/null +++ b/MoGe/configs/eval/benchmarks/hammer.json @@ -0,0 +1,10 @@ +{ + "HAMMER": { + "path": "data/eval/HAMMER", + "width": 1664, + "height": 832, + "split": ".index.txt", + "depth_unit": 1, + "has_sharp_boundary": true + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/ibims-1.json b/MoGe/configs/eval/benchmarks/ibims-1.json new file mode 100644 index 0000000000000000000000000000000000000000..a6f0a0387891deb09bcae61bc4e4098e04db7307 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/ibims-1.json @@ -0,0 +1,10 @@ +{ + "iBims-1": { + "path": "data/eval/iBims-1", + "width": 640, + "height": 480, + "split": ".index.txt", + "include_segmentation": true, + "has_sharp_boundary": true + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/kitti.json b/MoGe/configs/eval/benchmarks/kitti.json new file mode 100644 index 0000000000000000000000000000000000000000..10ca7c3eb560649ce25edbf4ed5c835e90396cb8 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/kitti.json @@ -0,0 +1,9 @@ +{ + "KITTI": { + "path": "data/eval/KITTI", + "width": 750, + "height": 375, + "split": ".index.txt", + "depth_unit": 1 + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/nyu.json b/MoGe/configs/eval/benchmarks/nyu.json new file mode 100644 index 0000000000000000000000000000000000000000..62841335b17f508ca903634b51b70f3e8a576186 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/nyu.json @@ -0,0 +1,8 @@ +{ + "NYUv2": { + "path": "data/eval/NYUv2", + "width": 640, + "height": 480, + "split": ".test.txt" + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/sintel.json b/MoGe/configs/eval/benchmarks/sintel.json new file mode 100644 index 0000000000000000000000000000000000000000..fde872e282e260f987208168bdcd166a104732d3 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/sintel.json @@ -0,0 +1,10 @@ +{ + "Sintel": { + "path": "data/eval/Sintel", + "width": 872, + "height": 436, + "split": ".index.txt", + "include_segmentation": true, + "has_sharp_boundary": true + } +} \ No newline at end of file diff --git a/MoGe/configs/eval/benchmarks/spring.json b/MoGe/configs/eval/benchmarks/spring.json new file mode 100644 index 0000000000000000000000000000000000000000..a18e51a969fe5b605c03ed1f0a4714dec9379539 --- /dev/null +++ b/MoGe/configs/eval/benchmarks/spring.json @@ -0,0 +1,9 @@ +{ + "Spring": { + "path": "data/eval/Spring", + "width": 1920, + "height": 1080, + "split": ".test.txt", + "has_sharp_boundary": true + } +} \ No newline at end of file diff --git a/MoGe/configs/train/v1.json b/MoGe/configs/train/v1.json new file mode 100644 index 0000000000000000000000000000000000000000..f87f38944129e60da58eed20395d36f4404d7164 --- /dev/null +++ b/MoGe/configs/train/v1.json @@ -0,0 +1,77 @@ +{ + "data": { + "aspect_ratio_range": [0.5, 2.0], + "area_range": [250000, 1000000], + "clamp_max_depth": 1000.0, + "center_augmentation": 0.5, + "fov_range_absolute": [1, 179], + "fov_range_relative": [0.01, 1.0], + "image_augmentation": ["jittering", "jpeg_loss", "blurring"], + "datasets": [ + { + "name": "TartanAir", + "path": "blobmnt/data_v3/TartanAir", + "label_type": "synthetic", + "index": ".index.txt", + "depth": "depth.png", + "weight": 4.8, + "center_augmentation": 0.25, + "fov_range_absolute": [30, 150], + "fov_range_relative": [0.5, 1.0], + "image_augmentation": ["jittering", "jpeg_loss", "blurring", "shot_noise"] + } + ] + }, + "model_version": "v1", + "model": { + "encoder": "dinov2_vitl14", + "remap_output": "exp", + "intermediate_layers": 4, + "dim_upsample": [256, 128, 64], + "dim_times_res_block_hidden": 2, + "num_res_blocks": 2, + "num_tokens_range": [1200, 2500], + "last_conv_channels": 32, + "last_conv_size": 1 + }, + "optimizer": { + "type": "AdamW", + "params": [ + {"params": {"include": ["*"], "exclude": ["*backbone.*"]}, "lr": 1e-4}, + {"params": {"include": ["*backbone.*"]}, "lr": 1e-5} + ] + }, + "lr_scheduler": { + "type": "SequentialLR", + "params": { + "schedulers": [ + {"type": "LambdaLR", "params": {"lr_lambda": ["1.0", "max(0.0, min(1.0, (epoch - 1000) / 1000))"]}}, + {"type": "StepLR", "params": {"step_size": 25000, "gamma": 0.5}} + ], + "milestones": [2000] + } + }, + "low_resolution_training_steps": 50000, + "loss": { + "invalid": {}, + "synthetic": { + "global": {"function": "affine_invariant_global_loss", "weight": 1.0, "params": {"align_resolution": 32}}, + "patch_4": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 4, "align_resolution": 16, "num_patches": 16}}, + "patch_16": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 16, "align_resolution": 8, "num_patches": 256}}, + "patch_64": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 64, "align_resolution": 4, "num_patches": 4096}}, + "normal": {"function": "normal_loss", "weight": 1.0}, + "mask": {"function": "mask_l2_loss", "weight": 1.0} + }, + "sfm": { + "global": {"function": "affine_invariant_global_loss", "weight": 1.0, "params": {"align_resolution": 32}}, + "patch_4": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 4, "align_resolution": 16, "num_patches": 16}}, + "patch_16": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 16, "align_resolution": 8, "num_patches": 256}}, + "mask": {"function": "mask_l2_loss", "weight": 1.0} + }, + "lidar": { + "global": {"function": "affine_invariant_global_loss", "weight": 1.0, "params": {"align_resolution": 32}}, + "patch_4": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 4, "align_resolution": 16, "num_patches": 16}}, + "mask": {"function": "mask_l2_loss", "weight": 1.0} + } + } +} \ No newline at end of file diff --git a/MoGe/docs/eval.md b/MoGe/docs/eval.md new file mode 100644 index 0000000000000000000000000000000000000000..a9d93e4a540c6df1c06aaa5694c8377e67ba468f --- /dev/null +++ b/MoGe/docs/eval.md @@ -0,0 +1,77 @@ +# Evaluation + +We provide a unified evaluation script that runs baselines on multiple benchmarks. It takes a baseline model and evaluation configurations, evaluates on-the-fly, and reports results instantly in a JSON file. + +## Benchmarks + +Donwload the processed datasets from [Huggingface Datasets](https://huggingface.co/datasets/Ruicheng/monocular-geometry-evaluation) and put them in the `data/eval` directory, using `huggingface-cli`: + +```bash +mkdir -p data/eval +huggingface-cli download Ruicheng/monocular-geometry-evaluation --repo-type dataset --local-dir data/eval --local-dir-use-symlinks False +``` + +Then unzip the downloaded files: + +```bash +cd data/eval +unzip '*.zip' +# rm *.zip # if you don't keep the zip files +``` + +## Configuration + +See [`configs/eval/all_benchmarks.json`](../configs/eval/all_benchmarks.json) for an example of evaluation configurations on all benchmarks. You can modify this file to evaluate on different benchmarks or different baselines. + +## Baseline + +Some examples of baselines are provided in [`baselines/`](../baselines/). Pass the path to the baseline model python code to the `--baseline` argument of the evaluation script. + +## Run Evaluation + +Run the script [`moge/scripts/eval_baseline.py`](../moge/scripts/eval_baseline.py). +For example, + +```bash +# Evaluate MoGe on the 10 benchmarks +python moge/scripts/eval_baseline.py --baseline baselines/moge.py --config configs/eval/all_benchmarks.json --output eval_output/moge.json --pretrained Ruicheng/moge-vitl --resolution_level 9 + +# Evaluate Depth Anything V2 on the 10 benchmarks. (NOTE: affine disparity) +python moge/scripts/eval_baseline.py --baseline baselines/da_v2.py --config configs/eval/all_benchmarks.json --output eval_output/da_v2.json +``` + +The `--baselies` `--input` `--output` arguments are for the inference script. The rest arguments, e.g. `--pretrained` `--resolution_level`, are custormized for loading the baseline model. + +Details of the arguments: + +``` +Usage: eval_baseline.py [OPTIONS] + + Evaluation script. + +Options: + --baseline PATH Path to the baseline model python code. + --config PATH Path to the evaluation configurations. Defaults to + "configs/eval/all_benchmarks.json". + --output PATH Path to the output json file. + --oracle Use oracle mode for evaluation, i.e., use the GT intrinsics + input. + --dump_pred Dump predition results. + --dump_gt Dump ground truth. + --help Show this message and exit. +``` + + + +## Wrap a Customized Baseline + +Wrap any baseline method with [`moge.test.baseline.MGEBaselineInterface`](../moge/test/baseline.py). +See [`baselines/`](../baselines/) for more examples. + +It is a good idea to check the correctness of the baseline implementation by running inference on a small set of images via [`moge/scripts/infer_baselines.py`](../moge/scripts/infer_baselines.py): + +```base +python moge/scripts/infer_baselines.py --baseline baselines/moge.py --input example_images/ --output infer_outupt/moge --pretrained Ruicheng/moge-vitl --maps --ply +``` + + diff --git a/MoGe/docs/train.md b/MoGe/docs/train.md new file mode 100644 index 0000000000000000000000000000000000000000..170abb80e08ac5eb25badedc2b05138c21bb33f2 --- /dev/null +++ b/MoGe/docs/train.md @@ -0,0 +1,181 @@ + +# Training + +This document provides instructions for training and finetuning the MoGe model. + +## Additional Requirements + +The following packages other than those listed in [`pyproject.toml`](../pyproject.toml) are required for training and finetuning the MoGe model: + +``` +accelerate +sympy +mlflow +``` + +## Data preparation + +### Dataset format + +Each dataset should be organized as follows: + +``` +somedataset +├── .index.txt # A list of instance paths +├── folder1 +│ ├── instance1 # Each instance is in a folder +│ │ ├── image.jpg # RGB image. +│ │ ├── depth.png # 16-bit depth. See moge/utils/io.py for details +│ │ ├── meta.json # Stores "intrinsics" as a 3x3 matrix +│ │ └── ... # Other componests such as segmentation mask, normal map etc. +... +``` + +* `.index.txt` is placed at top directory to store a list of instance paths in this dataset. The dataloader will look for instances in this list. You may also use a custom split, e.g. `.train.txt`, `.val.txt` and specify it in the configuration file. + +* For depth images, it is recommended to use `read_depth()` and `write_depth()` in [`moge/utils/io.py`](../moge/utils/io.py) to read and write depth images. The depth is stored in logarithmic scale in 16-bit PNG format, offering a balanced precision, dynamic range and compression ratio compared to 16-bit and 32-bit EXR and linear depth formats. It also encodes `NaN` and `Inf` values for invalid depth values. + +* The `meta.json` should be a dictionary containing the key `intrinsics`, which are **normalized** camera parameters. You may put more metadata. + +* We also support reading and storing segementation masks for evaluation data (see paper evaluation of local points), which are saved in PNG format with semantic labels stored in png metadata as JSON strings. See `read_segmentation()` and `write_segmentation()` in [`moge/utils/io.py`](../moge/utils/io.py) for details. + + +### Visual inspection + +We provide a script to visualize the data and check the data quality. It will export the instance as a PLY file for visualization of point cloud. + +```bash +python moge/scripts/vis_data.py PATH_TO_INSTANCE --ply [-o SOMEWHERE_ELSE_TO_SAVE_VIS] +``` + +### DataLoader + +Our training dataloaders is customized to handle loading data, performing perspective crop, and augmentation in a multithreading pipeline. Please refer to [`moge/train/dataloader.py`](../moge/train/dataloader.py) if you have any concern. + + +## Configuration + +See [`configs/train/v1.json`](../configs/train/v1.json) for an example configuration file. The configuration file defines the hyperparameters for training the MoGe model. +Here is a commented configuration for reference: + +```json +{ + "data": { + "aspect_ratio_range": [0.5, 2.0], # Range of aspect ratio of sampled images + "area_range": [250000, 1000000], # Range of sampled image area in pixels + "clamp_max_depth": 1000.0, # Maximum far/near + "center_augmentation": 0.5, # Ratio of center crop augmentation + "fov_range_absolute": [1, 179], # Absolute range of FOV in degrees + "fov_range_relative": [0.01, 1.0], # Relative range of FOV to the original FOV + "image_augmentation": ["jittering", "jpeg_loss", "blurring"], # List of image augmentation techniques + "datasets": [ + { + "name": "TartanAir", # Name of the dataset. Name it as you like. + "path": "data/TartanAir", # Path to the dataset + "label_type": "synthetic", # Label type for this dataset. Losses will be applied accordingly. see "loss" config + "weight": 4.8, # Probability of sampling this dataset + "index": ".index.txt", # File name of the index file. Defaults to .index.txt + "depth": "depth.png", # File name of depth images. Defaults to depth.png + "center_augmentation": 0.25, # Below are dataset-specific hyperparameters. Overriding the global ones above. + "fov_range_absolute": [30, 150], + "fov_range_relative": [0.5, 1.0], + "image_augmentation": ["jittering", "jpeg_loss", "blurring", "shot_noise"] + } + ] + }, + "model_version": "v1", # Model version. If you have multiple model variants, you can use this to switch between them. + "model": { # Model hyperparameters. Will be passed to Model __init__() as kwargs. + "encoder": "dinov2_vitl14", + "remap_output": "exp", + "intermediate_layers": 4, + "dim_upsample": [256, 128, 64], + "dim_times_res_block_hidden": 2, + "num_res_blocks": 2, + "num_tokens_range": [1200, 2500], + "last_conv_channels": 32, + "last_conv_size": 1 + }, + "optimizer": { # Reflection-like optimizer configurations. See moge.train.utils.py build_optimizer() for details. + "type": "AdamW", + "params": [ + {"params": {"include": ["*"], "exclude": ["*backbone.*"]}, "lr": 1e-4}, + {"params": {"include": ["*backbone.*"]}, "lr": 1e-5} + ] + }, + "lr_scheduler": { # Reflection-like lr_scheduler configurations. See moge.train.utils.py build_lr_scheduler() for details. + "type": "SequentialLR", + "params": { + "schedulers": [ + {"type": "LambdaLR", "params": {"lr_lambda": ["1.0", "max(0.0, min(1.0, (epoch - 1000) / 1000))"]}}, + {"type": "StepLR", "params": {"step_size": 25000, "gamma": 0.5}} + ], + "milestones": [2000] + } + }, + "low_resolution_training_steps": 50000, # Total number of low-resolution training steps. It makes the early stage training faster. Later stage training on varying size images will be slower. + "loss": { + "invalid": {}, # invalid instance due to runtime error when loading data + "synthetic": { # Below are loss hyperparameters + "global": {"function": "affine_invariant_global_loss", "weight": 1.0, "params": {"align_resolution": 32}}, + "patch_4": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 4, "align_resolution": 16, "num_patches": 16}}, + "patch_16": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 16, "align_resolution": 8, "num_patches": 256}}, + "patch_64": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 64, "align_resolution": 4, "num_patches": 4096}}, + "normal": {"function": "normal_loss", "weight": 1.0}, + "mask": {"function": "mask_l2_loss", "weight": 1.0} + }, + "sfm": { + "global": {"function": "affine_invariant_global_loss", "weight": 1.0, "params": {"align_resolution": 32}}, + "patch_4": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 4, "align_resolution": 16, "num_patches": 16}}, + "patch_16": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 16, "align_resolution": 8, "num_patches": 256}}, + "mask": {"function": "mask_l2_loss", "weight": 1.0} + }, + "lidar": { + "global": {"function": "affine_invariant_global_loss", "weight": 1.0, "params": {"align_resolution": 32}}, + "patch_4": {"function": "affine_invariant_local_loss", "weight": 1.0, "params": {"level": 4, "align_resolution": 16, "num_patches": 16}}, + "mask": {"function": "mask_l2_loss", "weight": 1.0} + } + } +} +``` + +## Run Training + +Launch the training script [`moge/scripts/train.py`](../moge/scripts/train.py). Note that we use [`accelerate`](https://github.com/huggingface/accelerate) for distributed training. + +```bash +accelerate launch \ + --num_processes 8 \ + moge/scripts/train.py \ + --config configs/train/v1.json \ + --workspace workspace/debug \ + --gradient_accumulation_steps 2 \ + --batch_size_forward 2 \ + --checkpoint latest \ + --enable_gradient_checkpointing True \ + --vis_every 1000 \ + --enable_mlflow True +``` + + +## Finetuning + +To finetune the pre-trained MoGe model, download the model checkpoint and put it in a local directory, e.g. `pretrained/moge-vitl.pt`. + +> NOTE: when finetuning pretrained MoGe model, a much lower learning rate is required. +The suggested learning rate for finetuning is not greater than 1e-5 for the head and 1e-6 for the backbone. +And the batch size is recommended to be 32 at least. +The settings in default configuration are not optimal for specific datasets and may require further tuning. + +```bash +accelerate launch \ + --num_processes 8 \ + moge/scripts/train.py \ + --config configs/train/v1.json \ + --workspace workspace/debug \ + --gradient_accumulation_steps 2 \ + --batch_size_forward 2 \ + --checkpoint pretrained/moge-vitl.pt \ + --enable_gradient_checkpointing True \ + --vis_every 1000 \ + --enable_mlflow True +``` diff --git a/MoGe/moge/__init__.py b/MoGe/moge/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MoGe/moge/model/__init__.py b/MoGe/moge/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b0159256956327e554db1385714cf715080e66f --- /dev/null +++ b/MoGe/moge/model/__init__.py @@ -0,0 +1,17 @@ +import importlib +from typing import * + +if TYPE_CHECKING: + from .v1 import MoGeModel as MoGeModelV1 + + +def import_model_class_by_version(version: str) -> Type[Union['MoGeModelV1']]: + assert version in ['v1'], f'Unsupported model version: {version}' + + try: + module = importlib.import_module(f'.{version}', __package__) + except ModuleNotFoundError: + raise ValueError(f'Model version "{version}" not found.') + + cls = getattr(module, 'MoGeModel') + return cls diff --git a/MoGe/moge/model/dinov2/__init__.py b/MoGe/moge/model/dinov2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ae847e46898077fe3d8701b8a181d7b4e3d41cd9 --- /dev/null +++ b/MoGe/moge/model/dinov2/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +__version__ = "0.0.1" diff --git a/MoGe/moge/model/dinov2/hub/__init__.py b/MoGe/moge/model/dinov2/hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b88da6bf80be92af00b72dfdb0a806fa64a7a2d9 --- /dev/null +++ b/MoGe/moge/model/dinov2/hub/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. diff --git a/MoGe/moge/model/dinov2/hub/backbones.py b/MoGe/moge/model/dinov2/hub/backbones.py new file mode 100644 index 0000000000000000000000000000000000000000..53fe83719d5107eb77a8f25ef1814c3d73446002 --- /dev/null +++ b/MoGe/moge/model/dinov2/hub/backbones.py @@ -0,0 +1,156 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from enum import Enum +from typing import Union + +import torch + +from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name + + +class Weights(Enum): + LVD142M = "LVD142M" + + +def _make_dinov2_model( + *, + arch_name: str = "vit_large", + img_size: int = 518, + patch_size: int = 14, + init_values: float = 1.0, + ffn_layer: str = "mlp", + block_chunks: int = 0, + num_register_tokens: int = 0, + interpolate_antialias: bool = False, + interpolate_offset: float = 0.1, + pretrained: bool = True, + weights: Union[Weights, str] = Weights.LVD142M, + **kwargs, +): + from ..models import vision_transformer as vits + + if isinstance(weights, str): + try: + weights = Weights[weights] + except KeyError: + raise AssertionError(f"Unsupported weights: {weights}") + + model_base_name = _make_dinov2_model_name(arch_name, patch_size) + vit_kwargs = dict( + img_size=img_size, + patch_size=patch_size, + init_values=init_values, + ffn_layer=ffn_layer, + block_chunks=block_chunks, + num_register_tokens=num_register_tokens, + interpolate_antialias=interpolate_antialias, + interpolate_offset=interpolate_offset, + ) + vit_kwargs.update(**kwargs) + model = vits.__dict__[arch_name](**vit_kwargs) + + if pretrained: + model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens) + url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_pretrain.pth" + state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu") + model.load_state_dict(state_dict, strict=True) + + return model + + +def dinov2_vits14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model(arch_name="vit_small", pretrained=pretrained, weights=weights, **kwargs) + + +def dinov2_vitb14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model(arch_name="vit_base", pretrained=pretrained, weights=weights, **kwargs) + + +def dinov2_vitl14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model(arch_name="vit_large", pretrained=pretrained, weights=weights, **kwargs) + + +def dinov2_vitg14(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model( + arch_name="vit_giant2", + ffn_layer="swiglufused", + weights=weights, + pretrained=pretrained, + **kwargs, + ) + + +def dinov2_vits14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model( + arch_name="vit_small", + pretrained=pretrained, + weights=weights, + num_register_tokens=4, + interpolate_antialias=True, + interpolate_offset=0.0, + **kwargs, + ) + + +def dinov2_vitb14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model( + arch_name="vit_base", + pretrained=pretrained, + weights=weights, + num_register_tokens=4, + interpolate_antialias=True, + interpolate_offset=0.0, + **kwargs, + ) + + +def dinov2_vitl14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model( + arch_name="vit_large", + pretrained=pretrained, + weights=weights, + num_register_tokens=4, + interpolate_antialias=True, + interpolate_offset=0.0, + **kwargs, + ) + + +def dinov2_vitg14_reg(*, pretrained: bool = True, weights: Union[Weights, str] = Weights.LVD142M, **kwargs): + """ + DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset. + """ + return _make_dinov2_model( + arch_name="vit_giant2", + ffn_layer="swiglufused", + weights=weights, + pretrained=pretrained, + num_register_tokens=4, + interpolate_antialias=True, + interpolate_offset=0.0, + **kwargs, + ) diff --git a/MoGe/moge/model/dinov2/hub/utils.py b/MoGe/moge/model/dinov2/hub/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9c6641404093652d5a2f19b4cf283d976ec39e64 --- /dev/null +++ b/MoGe/moge/model/dinov2/hub/utils.py @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import itertools +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +_DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2" + + +def _make_dinov2_model_name(arch_name: str, patch_size: int, num_register_tokens: int = 0) -> str: + compact_arch_name = arch_name.replace("_", "")[:4] + registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else "" + return f"dinov2_{compact_arch_name}{patch_size}{registers_suffix}" + + +class CenterPadding(nn.Module): + def __init__(self, multiple): + super().__init__() + self.multiple = multiple + + def _get_pad(self, size): + new_size = math.ceil(size / self.multiple) * self.multiple + pad_size = new_size - size + pad_size_left = pad_size // 2 + pad_size_right = pad_size - pad_size_left + return pad_size_left, pad_size_right + + @torch.inference_mode() + def forward(self, x): + pads = list(itertools.chain.from_iterable(self._get_pad(m) for m in x.shape[:1:-1])) + output = F.pad(x, pads) + return output diff --git a/MoGe/moge/model/dinov2/layers/__init__.py b/MoGe/moge/model/dinov2/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..05a0b61868e43abb821ca05a813bab2b8b43629e --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from .dino_head import DINOHead +from .mlp import Mlp +from .patch_embed import PatchEmbed +from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused +from .block import NestedTensorBlock +from .attention import MemEffAttention diff --git a/MoGe/moge/model/dinov2/layers/attention.py b/MoGe/moge/model/dinov2/layers/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..3fed573116d5c837be46a7525d8acf77422c2400 --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/attention.py @@ -0,0 +1,89 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +import logging +import os +import warnings + +from torch import Tensor +from torch import nn + + +logger = logging.getLogger("dinov2") + + +XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None +try: + if XFORMERS_ENABLED: + from xformers.ops import memory_efficient_attention, unbind + + XFORMERS_AVAILABLE = True + # warnings.warn("xFormers is available (Attention)") + else: + # warnings.warn("xFormers is disabled (Attention)") + raise ImportError +except ImportError: + XFORMERS_AVAILABLE = False + # warnings.warn("xFormers is not available (Attention)") + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ) -> None: + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: Tensor, attn_bias=None) -> Tensor: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + + q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] + attn = q @ k.transpose(-2, -1) + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class MemEffAttention(Attention): + def forward(self, x: Tensor, attn_bias=None) -> Tensor: + if not XFORMERS_AVAILABLE: + if attn_bias is not None: + raise AssertionError("xFormers is required for using nested tensors") + return super().forward(x) + + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) + + q, k, v = unbind(qkv, 2) + + x = memory_efficient_attention(q, k, v, attn_bias=attn_bias) + x = x.reshape([B, N, C]) + + x = self.proj(x) + x = self.proj_drop(x) + return x diff --git a/MoGe/moge/model/dinov2/layers/block.py b/MoGe/moge/model/dinov2/layers/block.py new file mode 100644 index 0000000000000000000000000000000000000000..fd5b8a7bb8527b74186af7c1e060e37bdb52c73d --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/block.py @@ -0,0 +1,259 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +import logging +import os +from typing import Callable, List, Any, Tuple, Dict +import warnings + +import torch +from torch import nn, Tensor + +from .attention import Attention, MemEffAttention +from .drop_path import DropPath +from .layer_scale import LayerScale +from .mlp import Mlp + + +logger = logging.getLogger("dinov2") + + +XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None +try: + if XFORMERS_ENABLED: + from xformers.ops import fmha, scaled_index_add, index_select_cat + + XFORMERS_AVAILABLE = True + # warnings.warn("xFormers is available (Block)") + else: + # warnings.warn("xFormers is disabled (Block)") + raise ImportError +except ImportError: + XFORMERS_AVAILABLE = False + # warnings.warn("xFormers is not available (Block)") + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + proj_bias: bool = True, + ffn_bias: bool = True, + drop: float = 0.0, + attn_drop: float = 0.0, + init_values=None, + drop_path: float = 0.0, + act_layer: Callable[..., nn.Module] = nn.GELU, + norm_layer: Callable[..., nn.Module] = nn.LayerNorm, + attn_class: Callable[..., nn.Module] = Attention, + ffn_layer: Callable[..., nn.Module] = Mlp, + ) -> None: + super().__init__() + # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}") + self.norm1 = norm_layer(dim) + self.attn = attn_class( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + attn_drop=attn_drop, + proj_drop=drop, + ) + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = ffn_layer( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop, + bias=ffn_bias, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.sample_drop_ratio = drop_path + + def forward(self, x: Tensor) -> Tensor: + def attn_residual_func(x: Tensor) -> Tensor: + return self.ls1(self.attn(self.norm1(x))) + + def ffn_residual_func(x: Tensor) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + if self.training and self.sample_drop_ratio > 0.1: + # the overhead is compensated only for a drop path rate larger than 0.1 + x = drop_add_residual_stochastic_depth( + x, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + x = drop_add_residual_stochastic_depth( + x, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + ) + elif self.training and self.sample_drop_ratio > 0.0: + x = x + self.drop_path1(attn_residual_func(x)) + x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 + else: + x = x + attn_residual_func(x) + x = x + ffn_residual_func(x) + return x + + +def drop_add_residual_stochastic_depth( + x: Tensor, + residual_func: Callable[[Tensor], Tensor], + sample_drop_ratio: float = 0.0, +) -> Tensor: + # 1) extract subset using permutation + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + x_subset = x[brange] + + # 2) apply residual_func to get residual + residual = residual_func(x_subset) + + x_flat = x.flatten(1) + residual = residual.flatten(1) + + residual_scale_factor = b / sample_subset_size + + # 3) add the residual + x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor) + return x_plus_residual.view_as(x) + + +def get_branges_scales(x, sample_drop_ratio=0.0): + b, n, d = x.shape + sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) + brange = (torch.randperm(b, device=x.device))[:sample_subset_size] + residual_scale_factor = b / sample_subset_size + return brange, residual_scale_factor + + +def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None): + if scaling_vector is None: + x_flat = x.flatten(1) + residual = residual.flatten(1) + x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor) + else: + x_plus_residual = scaled_index_add( + x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor + ) + return x_plus_residual + + +attn_bias_cache: Dict[Tuple, Any] = {} + + +def get_attn_bias_and_cat(x_list, branges=None): + """ + this will perform the index select, cat the tensors, and provide the attn_bias from cache + """ + batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list] + all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list)) + if all_shapes not in attn_bias_cache.keys(): + seqlens = [] + for b, x in zip(batch_sizes, x_list): + for _ in range(b): + seqlens.append(x.shape[1]) + attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens) + attn_bias._batch_sizes = batch_sizes + attn_bias_cache[all_shapes] = attn_bias + + if branges is not None: + cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1]) + else: + tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list) + cat_tensors = torch.cat(tensors_bs1, dim=1) + + return attn_bias_cache[all_shapes], cat_tensors + + +def drop_add_residual_stochastic_depth_list( + x_list: List[Tensor], + residual_func: Callable[[Tensor, Any], Tensor], + sample_drop_ratio: float = 0.0, + scaling_vector=None, +) -> Tensor: + # 1) generate random set of indices for dropping samples in the batch + branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list] + branges = [s[0] for s in branges_scales] + residual_scale_factors = [s[1] for s in branges_scales] + + # 2) get attention bias and index+concat the tensors + attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges) + + # 3) apply residual_func to get residual, and split the result + residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore + + outputs = [] + for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors): + outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x)) + return outputs + + +class NestedTensorBlock(Block): + def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]: + """ + x_list contains a list of tensors to nest together and run + """ + assert isinstance(self.attn, MemEffAttention) + + if self.training and self.sample_drop_ratio > 0.0: + + def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.attn(self.norm1(x), attn_bias=attn_bias) + + def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.mlp(self.norm2(x)) + + x_list = drop_add_residual_stochastic_depth_list( + x_list, + residual_func=attn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None, + ) + x_list = drop_add_residual_stochastic_depth_list( + x_list, + residual_func=ffn_residual_func, + sample_drop_ratio=self.sample_drop_ratio, + scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None, + ) + return x_list + else: + + def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias)) + + def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor: + return self.ls2(self.mlp(self.norm2(x))) + + attn_bias, x = get_attn_bias_and_cat(x_list) + x = x + attn_residual_func(x, attn_bias=attn_bias) + x = x + ffn_residual_func(x) + return attn_bias.split(x) + + def forward(self, x_or_x_list): + if isinstance(x_or_x_list, Tensor): + return super().forward(x_or_x_list) + elif isinstance(x_or_x_list, list): + if not XFORMERS_AVAILABLE: + raise AssertionError("xFormers is required for using nested tensors") + return self.forward_nested(x_or_x_list) + else: + raise AssertionError diff --git a/MoGe/moge/model/dinov2/layers/dino_head.py b/MoGe/moge/model/dinov2/layers/dino_head.py new file mode 100644 index 0000000000000000000000000000000000000000..0ace8ffd6297a1dd480b19db407b662a6ea0f565 --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/dino_head.py @@ -0,0 +1,58 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +from torch.nn.init import trunc_normal_ +from torch.nn.utils import weight_norm + + +class DINOHead(nn.Module): + def __init__( + self, + in_dim, + out_dim, + use_bn=False, + nlayers=3, + hidden_dim=2048, + bottleneck_dim=256, + mlp_bias=True, + ): + super().__init__() + nlayers = max(nlayers, 1) + self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias) + self.apply(self._init_weights) + self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False)) + self.last_layer.weight_g.data.fill_(1) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + + def forward(self, x): + x = self.mlp(x) + eps = 1e-6 if x.dtype == torch.float16 else 1e-12 + x = nn.functional.normalize(x, dim=-1, p=2, eps=eps) + x = self.last_layer(x) + return x + + +def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True): + if nlayers == 1: + return nn.Linear(in_dim, bottleneck_dim, bias=bias) + else: + layers = [nn.Linear(in_dim, hidden_dim, bias=bias)] + if use_bn: + layers.append(nn.BatchNorm1d(hidden_dim)) + layers.append(nn.GELU()) + for _ in range(nlayers - 2): + layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias)) + if use_bn: + layers.append(nn.BatchNorm1d(hidden_dim)) + layers.append(nn.GELU()) + layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias)) + return nn.Sequential(*layers) diff --git a/MoGe/moge/model/dinov2/layers/drop_path.py b/MoGe/moge/model/dinov2/layers/drop_path.py new file mode 100644 index 0000000000000000000000000000000000000000..1d640e0b969b8dcba96260243473700b4e5b24b5 --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/drop_path.py @@ -0,0 +1,34 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py + + +from torch import nn + + +def drop_path(x, drop_prob: float = 0.0, training: bool = False): + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0: + random_tensor.div_(keep_prob) + output = x * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) diff --git a/MoGe/moge/model/dinov2/layers/layer_scale.py b/MoGe/moge/model/dinov2/layers/layer_scale.py new file mode 100644 index 0000000000000000000000000000000000000000..51df0d7ce61f2b41fa9e6369f52391dd7fe7d386 --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/layer_scale.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 + +from typing import Union + +import torch +from torch import Tensor +from torch import nn + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: Union[float, Tensor] = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma diff --git a/MoGe/moge/model/dinov2/layers/mlp.py b/MoGe/moge/model/dinov2/layers/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..bbf9432aae9258612caeae910a7bde17999e328e --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/mlp.py @@ -0,0 +1,40 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py + + +from typing import Callable, Optional + +from torch import Tensor, nn + + +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = nn.GELU, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self.drop = nn.Dropout(drop) + + def forward(self, x: Tensor) -> Tensor: + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x diff --git a/MoGe/moge/model/dinov2/layers/patch_embed.py b/MoGe/moge/model/dinov2/layers/patch_embed.py new file mode 100644 index 0000000000000000000000000000000000000000..8b7c0804784a42cf80c0297d110dcc68cc85b339 --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/patch_embed.py @@ -0,0 +1,88 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py + +from typing import Callable, Optional, Tuple, Union + +from torch import Tensor +import torch.nn as nn + + +def make_2tuple(x): + if isinstance(x, tuple): + assert len(x) == 2 + return x + + assert isinstance(x, int) + return (x, x) + + +class PatchEmbed(nn.Module): + """ + 2D image to patch embedding: (B,C,H,W) -> (B,N,D) + + Args: + img_size: Image size. + patch_size: Patch token size. + in_chans: Number of input image channels. + embed_dim: Number of linear projection output channels. + norm_layer: Normalization layer. + """ + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten_embedding: bool = True, + ) -> None: + super().__init__() + + image_HW = make_2tuple(img_size) + patch_HW = make_2tuple(patch_size) + patch_grid_size = ( + image_HW[0] // patch_HW[0], + image_HW[1] // patch_HW[1], + ) + + self.img_size = image_HW + self.patch_size = patch_HW + self.patches_resolution = patch_grid_size + self.num_patches = patch_grid_size[0] * patch_grid_size[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.flatten_embedding = flatten_embedding + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x: Tensor) -> Tensor: + _, _, H, W = x.shape + patch_H, patch_W = self.patch_size + + assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}" + assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}" + + x = self.proj(x) # B C H W + H, W = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) # B HW C + x = self.norm(x) + if not self.flatten_embedding: + x = x.reshape(-1, H, W, self.embed_dim) # B H W C + return x + + def flops(self) -> float: + Ho, Wo = self.patches_resolution + flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) + if self.norm is not None: + flops += Ho * Wo * self.embed_dim + return flops diff --git a/MoGe/moge/model/dinov2/layers/swiglu_ffn.py b/MoGe/moge/model/dinov2/layers/swiglu_ffn.py new file mode 100644 index 0000000000000000000000000000000000000000..5ce211515774d42e04c8b51003bae53b88f14b35 --- /dev/null +++ b/MoGe/moge/model/dinov2/layers/swiglu_ffn.py @@ -0,0 +1,72 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import os +from typing import Callable, Optional +import warnings + +from torch import Tensor, nn +import torch.nn.functional as F + + +class SwiGLUFFN(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.chunk(2, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None +try: + if XFORMERS_ENABLED: + from xformers.ops import SwiGLU + + XFORMERS_AVAILABLE = True + # warnings.warn("xFormers is available (SwiGLU)") + else: + # warnings.warn("xFormers is disabled (SwiGLU)") + raise ImportError +except ImportError: + SwiGLU = SwiGLUFFN + XFORMERS_AVAILABLE = False + + # warnings.warn("xFormers is not available (SwiGLU)") + + +class SwiGLUFFNFused(SwiGLU): + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + act_layer: Callable[..., nn.Module] = None, + drop: float = 0.0, + bias: bool = True, + ) -> None: + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + super().__init__( + in_features=in_features, + hidden_features=hidden_features, + out_features=out_features, + bias=bias, + ) diff --git a/MoGe/moge/model/dinov2/models/__init__.py b/MoGe/moge/model/dinov2/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3fdff20badbd5244bf79f16bf18dd2cb73982265 --- /dev/null +++ b/MoGe/moge/model/dinov2/models/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import logging + +from . import vision_transformer as vits + + +logger = logging.getLogger("dinov2") + + +def build_model(args, only_teacher=False, img_size=224): + args.arch = args.arch.removesuffix("_memeff") + if "vit" in args.arch: + vit_kwargs = dict( + img_size=img_size, + patch_size=args.patch_size, + init_values=args.layerscale, + ffn_layer=args.ffn_layer, + block_chunks=args.block_chunks, + qkv_bias=args.qkv_bias, + proj_bias=args.proj_bias, + ffn_bias=args.ffn_bias, + num_register_tokens=args.num_register_tokens, + interpolate_offset=args.interpolate_offset, + interpolate_antialias=args.interpolate_antialias, + ) + teacher = vits.__dict__[args.arch](**vit_kwargs) + if only_teacher: + return teacher, teacher.embed_dim + student = vits.__dict__[args.arch]( + **vit_kwargs, + drop_path_rate=args.drop_path_rate, + drop_path_uniform=args.drop_path_uniform, + ) + embed_dim = student.embed_dim + return student, teacher, embed_dim + + +def build_model_from_cfg(cfg, only_teacher=False): + return build_model(cfg.student, only_teacher=only_teacher, img_size=cfg.crops.global_crops_size) diff --git a/MoGe/moge/model/dinov2/models/vision_transformer.py b/MoGe/moge/model/dinov2/models/vision_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..1007ba57ddb35109c91716f1f5bf203db346e7be --- /dev/null +++ b/MoGe/moge/model/dinov2/models/vision_transformer.py @@ -0,0 +1,396 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +# References: +# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py + +from functools import partial +import math +import logging +from typing import Sequence, Tuple, Union, Callable + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from torch.nn.init import trunc_normal_ + +from ..layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block + + +logger = logging.getLogger("dinov2") + + +def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +class BlockChunk(nn.ModuleList): + def forward(self, x): + for b in self: + x = b(x) + return x + + +class DinoVisionTransformer(nn.Module): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4.0, + qkv_bias=True, + ffn_bias=True, + proj_bias=True, + drop_path_rate=0.0, + drop_path_uniform=False, + init_values=None, # for layerscale: None or 0 => no layerscale + embed_layer=PatchEmbed, + act_layer=nn.GELU, + block_fn=Block, + ffn_layer="mlp", + block_chunks=1, + num_register_tokens=0, + interpolate_antialias=False, + interpolate_offset=0.1, + ): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + proj_bias (bool): enable bias for proj in attn if True + ffn_bias (bool): enable bias for ffn if True + drop_path_rate (float): stochastic depth rate + drop_path_uniform (bool): apply uniform drop rate across blocks + weight_init (str): weight init scheme + init_values (float): layer-scale init values + embed_layer (nn.Module): patch embedding layer + act_layer (nn.Module): MLP activation layer + block_fn (nn.Module): transformer block class + ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" + block_chunks: (int) split block sequence into block_chunks units for FSDP wrap + num_register_tokens: (int) number of extra cls tokens (so-called "registers") + interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings + interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings + """ + super().__init__() + norm_layer = partial(nn.LayerNorm, eps=1e-6) + + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_antialias = interpolate_antialias + self.interpolate_offset = interpolate_offset + + self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim) + num_patches = self.patch_embed.num_patches + + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) + assert num_register_tokens >= 0 + self.register_tokens = ( + nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None + ) + + if drop_path_uniform is True: + dpr = [drop_path_rate] * depth + else: + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + if ffn_layer == "mlp": + logger.info("using MLP layer as FFN") + ffn_layer = Mlp + elif ffn_layer == "swiglufused" or ffn_layer == "swiglu": + logger.info("using SwiGLU layer as FFN") + ffn_layer = SwiGLUFFNFused + elif ffn_layer == "identity": + logger.info("using Identity layer as FFN") + + def f(*args, **kwargs): + return nn.Identity() + + ffn_layer = f + else: + raise NotImplementedError + + blocks_list = [ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + ffn_layer=ffn_layer, + init_values=init_values, + ) + for i in range(depth) + ] + if block_chunks > 0: + self.chunked_blocks = True + chunked_blocks = [] + chunksize = depth // block_chunks + for i in range(0, depth, chunksize): + # this is to keep the block index consistent if we chunk the block list + chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize]) + self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks]) + else: + self.chunked_blocks = False + self.blocks = nn.ModuleList(blocks_list) + + self.norm = norm_layer(embed_dim) + self.head = nn.Identity() + + self.mask_token = nn.Parameter(torch.zeros(1, embed_dim)) + + self.init_weights() + + def init_weights(self): + trunc_normal_(self.pos_embed, std=0.02) + nn.init.normal_(self.cls_token, std=1e-6) + if self.register_tokens is not None: + nn.init.normal_(self.register_tokens, std=1e-6) + named_apply(init_weights_vit_timm, self) + + def interpolate_pos_encoding(self, x, w, h): + previous_dtype = x.dtype + npatch = x.shape[1] - 1 + N = self.pos_embed.shape[1] - 1 + if npatch == N and w == h: + return self.pos_embed + pos_embed = self.pos_embed.float() + class_pos_embed = pos_embed[:, 0] + patch_pos_embed = pos_embed[:, 1:] + dim = x.shape[-1] + w0 = w // self.patch_size + h0 = h // self.patch_size + M = int(math.sqrt(N)) # Recover the number of patches in each dimension + assert N == M * M + kwargs = {} + if self.interpolate_offset: + # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8 + # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors + sx = float(w0 + self.interpolate_offset) / M + sy = float(h0 + self.interpolate_offset) / M + kwargs["scale_factor"] = (sx, sy) + else: + # Simply specify an output size instead of a scale factor + kwargs["size"] = (w0, h0) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2), + mode="bicubic", + antialias=self.interpolate_antialias, + **kwargs, + ) + assert (w0, h0) == patch_pos_embed.shape[-2:] + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype) + + def prepare_tokens_with_masks(self, x, masks=None): + B, nc, w, h = x.shape + x = self.patch_embed(x) + if masks is not None: + x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x) + + x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1) + x = x + self.interpolate_pos_encoding(x, w, h) + + if self.register_tokens is not None: + x = torch.cat( + ( + x[:, :1], + self.register_tokens.expand(x.shape[0], -1, -1), + x[:, 1:], + ), + dim=1, + ) + + return x + + def forward_features_list(self, x_list, masks_list): + x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)] + for blk in self.blocks: + x = blk(x) + + all_x = x + output = [] + for x, masks in zip(all_x, masks_list): + x_norm = self.norm(x) + output.append( + { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + ) + return output + + def forward_features(self, x, masks=None): + if isinstance(x, list): + return self.forward_features_list(x, masks) + + x = self.prepare_tokens_with_masks(x, masks) + + for blk in self.blocks: + x = blk(x) + + x_norm = self.norm(x) + return { + "x_norm_clstoken": x_norm[:, 0], + "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1], + "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :], + "x_prenorm": x, + "masks": masks, + } + + def _get_intermediate_layers_not_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + # If n is an int, take the n last blocks. If it's a list, take them + output, total_block_len = [], len(self.blocks) + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for i, blk in enumerate(self.blocks): + x = blk(x) + if i in blocks_to_take: + output.append(x) + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def _get_intermediate_layers_chunked(self, x, n=1): + x = self.prepare_tokens_with_masks(x) + output, i, total_block_len = [], 0, len(self.blocks[-1]) + # If n is an int, take the n last blocks. If it's a list, take them + blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n + for block_chunk in self.blocks: + for blk in block_chunk[i:]: # Passing the nn.Identity() + x = blk(x) + if i in blocks_to_take: + output.append(x) + i += 1 + assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found" + return output + + def get_intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, # Layers or n last layers to take + reshape: bool = False, + return_class_token: bool = False, + norm=True, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: + if self.chunked_blocks: + outputs = self._get_intermediate_layers_chunked(x, n) + else: + outputs = self._get_intermediate_layers_not_chunked(x, n) + if norm: + outputs = [self.norm(out) for out in outputs] + class_tokens = [out[:, 0] for out in outputs] + outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs] + if reshape: + B, _, w, h = x.shape + outputs = [ + out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous() + for out in outputs + ] + if return_class_token: + return tuple(zip(outputs, class_tokens)) + return tuple(outputs) + + def forward(self, *args, is_training=False, **kwargs): + ret = self.forward_features(*args, **kwargs) + if is_training: + return ret + else: + return self.head(ret["x_norm_clstoken"]) + + +def init_weights_vit_timm(module: nn.Module, name: str = ""): + """ViT weight initialization, original timm impl (for reproducibility)""" + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + + +def vit_small(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=384, + depth=12, + num_heads=6, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_base(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_large(patch_size=16, num_register_tokens=0, **kwargs): + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1024, + depth=24, + num_heads=16, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model + + +def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs): + """ + Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 + """ + model = DinoVisionTransformer( + patch_size=patch_size, + embed_dim=1536, + depth=40, + num_heads=24, + mlp_ratio=4, + block_fn=partial(Block, attn_class=MemEffAttention), + num_register_tokens=num_register_tokens, + **kwargs, + ) + return model diff --git a/MoGe/moge/model/dinov2/utils/__init__.py b/MoGe/moge/model/dinov2/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b88da6bf80be92af00b72dfdb0a806fa64a7a2d9 --- /dev/null +++ b/MoGe/moge/model/dinov2/utils/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. diff --git a/MoGe/moge/model/dinov2/utils/cluster.py b/MoGe/moge/model/dinov2/utils/cluster.py new file mode 100644 index 0000000000000000000000000000000000000000..3df87dc3e1eb4f0f8a280dc3137cfef031886314 --- /dev/null +++ b/MoGe/moge/model/dinov2/utils/cluster.py @@ -0,0 +1,95 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from enum import Enum +import os +from pathlib import Path +from typing import Any, Dict, Optional + + +class ClusterType(Enum): + AWS = "aws" + FAIR = "fair" + RSC = "rsc" + + +def _guess_cluster_type() -> ClusterType: + uname = os.uname() + if uname.sysname == "Linux": + if uname.release.endswith("-aws"): + # Linux kernel versions on AWS instances are of the form "5.4.0-1051-aws" + return ClusterType.AWS + elif uname.nodename.startswith("rsc"): + # Linux kernel versions on RSC instances are standard ones but hostnames start with "rsc" + return ClusterType.RSC + + return ClusterType.FAIR + + +def get_cluster_type(cluster_type: Optional[ClusterType] = None) -> Optional[ClusterType]: + if cluster_type is None: + return _guess_cluster_type() + + return cluster_type + + +def get_checkpoint_path(cluster_type: Optional[ClusterType] = None) -> Optional[Path]: + cluster_type = get_cluster_type(cluster_type) + if cluster_type is None: + return None + + CHECKPOINT_DIRNAMES = { + ClusterType.AWS: "checkpoints", + ClusterType.FAIR: "checkpoint", + ClusterType.RSC: "checkpoint/dino", + } + return Path("/") / CHECKPOINT_DIRNAMES[cluster_type] + + +def get_user_checkpoint_path(cluster_type: Optional[ClusterType] = None) -> Optional[Path]: + checkpoint_path = get_checkpoint_path(cluster_type) + if checkpoint_path is None: + return None + + username = os.environ.get("USER") + assert username is not None + return checkpoint_path / username + + +def get_slurm_partition(cluster_type: Optional[ClusterType] = None) -> Optional[str]: + cluster_type = get_cluster_type(cluster_type) + if cluster_type is None: + return None + + SLURM_PARTITIONS = { + ClusterType.AWS: "learnlab", + ClusterType.FAIR: "learnlab", + ClusterType.RSC: "learn", + } + return SLURM_PARTITIONS[cluster_type] + + +def get_slurm_executor_parameters( + nodes: int, num_gpus_per_node: int, cluster_type: Optional[ClusterType] = None, **kwargs +) -> Dict[str, Any]: + # create default parameters + params = { + "mem_gb": 0, # Requests all memory on a node, see https://slurm.schedmd.com/sbatch.html + "gpus_per_node": num_gpus_per_node, + "tasks_per_node": num_gpus_per_node, # one task per GPU + "cpus_per_task": 10, + "nodes": nodes, + "slurm_partition": get_slurm_partition(cluster_type), + } + # apply cluster-specific adjustments + cluster_type = get_cluster_type(cluster_type) + if cluster_type == ClusterType.AWS: + params["cpus_per_task"] = 12 + del params["mem_gb"] + elif cluster_type == ClusterType.RSC: + params["cpus_per_task"] = 12 + # set additional parameters / apply overrides + params.update(kwargs) + return params diff --git a/MoGe/moge/model/dinov2/utils/config.py b/MoGe/moge/model/dinov2/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..c9de578787bbcb376f8bd5a782206d0eb7ec1f52 --- /dev/null +++ b/MoGe/moge/model/dinov2/utils/config.py @@ -0,0 +1,72 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import math +import logging +import os + +from omegaconf import OmegaConf + +import dinov2.distributed as distributed +from dinov2.logging import setup_logging +from dinov2.utils import utils +from dinov2.configs import dinov2_default_config + + +logger = logging.getLogger("dinov2") + + +def apply_scaling_rules_to_cfg(cfg): # to fix + if cfg.optim.scaling_rule == "sqrt_wrt_1024": + base_lr = cfg.optim.base_lr + cfg.optim.lr = base_lr + cfg.optim.lr *= math.sqrt(cfg.train.batch_size_per_gpu * distributed.get_global_size() / 1024.0) + logger.info(f"sqrt scaling learning rate; base: {base_lr}, new: {cfg.optim.lr}") + else: + raise NotImplementedError + return cfg + + +def write_config(cfg, output_dir, name="config.yaml"): + logger.info(OmegaConf.to_yaml(cfg)) + saved_cfg_path = os.path.join(output_dir, name) + with open(saved_cfg_path, "w") as f: + OmegaConf.save(config=cfg, f=f) + return saved_cfg_path + + +def get_cfg_from_args(args): + args.output_dir = os.path.abspath(args.output_dir) + args.opts += [f"train.output_dir={args.output_dir}"] + default_cfg = OmegaConf.create(dinov2_default_config) + cfg = OmegaConf.load(args.config_file) + cfg = OmegaConf.merge(default_cfg, cfg, OmegaConf.from_cli(args.opts)) + return cfg + + +def default_setup(args): + distributed.enable(overwrite=True) + seed = getattr(args, "seed", 0) + rank = distributed.get_global_rank() + + global logger + setup_logging(output=args.output_dir, level=logging.INFO) + logger = logging.getLogger("dinov2") + + utils.fix_random_seeds(seed + rank) + logger.info("git:\n {}\n".format(utils.get_sha())) + logger.info("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items()))) + + +def setup(args): + """ + Create configs and perform basic setups. + """ + cfg = get_cfg_from_args(args) + os.makedirs(args.output_dir, exist_ok=True) + default_setup(args) + apply_scaling_rules_to_cfg(cfg) + write_config(cfg, args.output_dir) + return cfg diff --git a/MoGe/moge/model/dinov2/utils/dtype.py b/MoGe/moge/model/dinov2/utils/dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..80f4cd74d99faa2731dbe9f8d3a13d71b3f8e3a8 --- /dev/null +++ b/MoGe/moge/model/dinov2/utils/dtype.py @@ -0,0 +1,37 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + + +from typing import Dict, Union + +import numpy as np +import torch + + +TypeSpec = Union[str, np.dtype, torch.dtype] + + +_NUMPY_TO_TORCH_DTYPE: Dict[np.dtype, torch.dtype] = { + np.dtype("bool"): torch.bool, + np.dtype("uint8"): torch.uint8, + np.dtype("int8"): torch.int8, + np.dtype("int16"): torch.int16, + np.dtype("int32"): torch.int32, + np.dtype("int64"): torch.int64, + np.dtype("float16"): torch.float16, + np.dtype("float32"): torch.float32, + np.dtype("float64"): torch.float64, + np.dtype("complex64"): torch.complex64, + np.dtype("complex128"): torch.complex128, +} + + +def as_torch_dtype(dtype: TypeSpec) -> torch.dtype: + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str): + dtype = np.dtype(dtype) + assert isinstance(dtype, np.dtype), f"Expected an instance of nunpy dtype, got {type(dtype)}" + return _NUMPY_TO_TORCH_DTYPE[dtype] diff --git a/MoGe/moge/model/dinov2/utils/param_groups.py b/MoGe/moge/model/dinov2/utils/param_groups.py new file mode 100644 index 0000000000000000000000000000000000000000..9a5d2ff627cddadc222e5f836864ee39c865208f --- /dev/null +++ b/MoGe/moge/model/dinov2/utils/param_groups.py @@ -0,0 +1,103 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +from collections import defaultdict +import logging + + +logger = logging.getLogger("dinov2") + + +def get_vit_lr_decay_rate(name, lr_decay_rate=1.0, num_layers=12, force_is_backbone=False, chunked_blocks=False): + """ + Calculate lr decay rate for different ViT blocks. + Args: + name (string): parameter name. + lr_decay_rate (float): base lr decay rate. + num_layers (int): number of ViT blocks. + Returns: + lr decay rate for the given parameter. + """ + layer_id = num_layers + 1 + if name.startswith("backbone") or force_is_backbone: + if ( + ".pos_embed" in name + or ".patch_embed" in name + or ".mask_token" in name + or ".cls_token" in name + or ".register_tokens" in name + ): + layer_id = 0 + elif force_is_backbone and ( + "pos_embed" in name + or "patch_embed" in name + or "mask_token" in name + or "cls_token" in name + or "register_tokens" in name + ): + layer_id = 0 + elif ".blocks." in name and ".residual." not in name: + layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1 + elif chunked_blocks and "blocks." in name and "residual." not in name: + layer_id = int(name[name.find("blocks.") :].split(".")[2]) + 1 + elif "blocks." in name and "residual." not in name: + layer_id = int(name[name.find("blocks.") :].split(".")[1]) + 1 + + return lr_decay_rate ** (num_layers + 1 - layer_id) + + +def get_params_groups_with_decay(model, lr_decay_rate=1.0, patch_embed_lr_mult=1.0): + chunked_blocks = False + if hasattr(model, "n_blocks"): + logger.info("chunked fsdp") + n_blocks = model.n_blocks + chunked_blocks = model.chunked_blocks + elif hasattr(model, "blocks"): + logger.info("first code branch") + n_blocks = len(model.blocks) + elif hasattr(model, "backbone"): + logger.info("second code branch") + n_blocks = len(model.backbone.blocks) + else: + logger.info("else code branch") + n_blocks = 0 + all_param_groups = [] + + for name, param in model.named_parameters(): + name = name.replace("_fsdp_wrapped_module.", "") + if not param.requires_grad: + continue + decay_rate = get_vit_lr_decay_rate( + name, lr_decay_rate, num_layers=n_blocks, force_is_backbone=n_blocks > 0, chunked_blocks=chunked_blocks + ) + d = {"params": param, "is_last_layer": False, "lr_multiplier": decay_rate, "wd_multiplier": 1.0, "name": name} + + if "last_layer" in name: + d.update({"is_last_layer": True}) + + if name.endswith(".bias") or "norm" in name or "gamma" in name: + d.update({"wd_multiplier": 0.0}) + + if "patch_embed" in name: + d.update({"lr_multiplier": d["lr_multiplier"] * patch_embed_lr_mult}) + + all_param_groups.append(d) + logger.info(f"""{name}: lr_multiplier: {d["lr_multiplier"]}, wd_multiplier: {d["wd_multiplier"]}""") + + return all_param_groups + + +def fuse_params_groups(all_params_groups, keys=("lr_multiplier", "wd_multiplier", "is_last_layer")): + fused_params_groups = defaultdict(lambda: {"params": []}) + for d in all_params_groups: + identifier = "" + for k in keys: + identifier += k + str(d[k]) + "_" + + for k in keys: + fused_params_groups[identifier][k] = d[k] + fused_params_groups[identifier]["params"].append(d["params"]) + + return fused_params_groups.values() diff --git a/MoGe/moge/model/dinov2/utils/utils.py b/MoGe/moge/model/dinov2/utils/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..68f8e2c3be5f780bbb7e00359b5ac4fd0ba0785f --- /dev/null +++ b/MoGe/moge/model/dinov2/utils/utils.py @@ -0,0 +1,95 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +import logging +import os +import random +import subprocess +from urllib.parse import urlparse + +import numpy as np +import torch +from torch import nn + + +logger = logging.getLogger("dinov2") + + +def load_pretrained_weights(model, pretrained_weights, checkpoint_key): + if urlparse(pretrained_weights).scheme: # If it looks like an URL + state_dict = torch.hub.load_state_dict_from_url(pretrained_weights, map_location="cpu") + else: + state_dict = torch.load(pretrained_weights, map_location="cpu") + if checkpoint_key is not None and checkpoint_key in state_dict: + logger.info(f"Take key {checkpoint_key} in provided checkpoint dict") + state_dict = state_dict[checkpoint_key] + # remove `module.` prefix + state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} + # remove `backbone.` prefix induced by multicrop wrapper + state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()} + msg = model.load_state_dict(state_dict, strict=False) + logger.info("Pretrained weights found at {} and loaded with msg: {}".format(pretrained_weights, msg)) + + +def fix_random_seeds(seed=31): + """ + Fix random seeds. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + random.seed(seed) + + +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() + + sha = "N/A" + diff = "clean" + branch = "N/A" + try: + sha = _run(["git", "rev-parse", "HEAD"]) + subprocess.check_output(["git", "diff"], cwd=cwd) + diff = _run(["git", "diff-index", "HEAD"]) + diff = "has uncommitted changes" if diff else "clean" + branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +class CosineScheduler(object): + def __init__(self, base_value, final_value, total_iters, warmup_iters=0, start_warmup_value=0, freeze_iters=0): + super().__init__() + self.final_value = final_value + self.total_iters = total_iters + + freeze_schedule = np.zeros((freeze_iters)) + + warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters) + + iters = np.arange(total_iters - warmup_iters - freeze_iters) + schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters))) + self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule)) + + assert len(self.schedule) == self.total_iters + + def __getitem__(self, it): + if it >= self.total_iters: + return self.final_value + else: + return self.schedule[it] + + +def has_batchnorms(model): + bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm) + for name, module in model.named_modules(): + if isinstance(module, bn_types): + return True + return False diff --git a/MoGe/moge/model/utils.py b/MoGe/moge/model/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c50761d8740d9d0a0284e129503b8931c6fe08c4 --- /dev/null +++ b/MoGe/moge/model/utils.py @@ -0,0 +1,49 @@ +from typing import * + +import torch +import torch.nn as nn +import torch.nn.functional as F + +def wrap_module_with_gradient_checkpointing(module: nn.Module): + from torch.utils.checkpoint import checkpoint + class _CheckpointingWrapper(module.__class__): + _restore_cls = module.__class__ + def forward(self, *args, **kwargs): + return checkpoint(super().forward, *args, use_reentrant=False, **kwargs) + + module.__class__ = _CheckpointingWrapper + return module + + +def unwrap_module_with_gradient_checkpointing(module: nn.Module): + module.__class__ = module.__class__._restore_cls + + +def wrap_dinov2_attention_with_sdpa(module: nn.Module): + assert torch.__version__ >= '2.0', "SDPA requires PyTorch 2.0 or later" + class _AttentionWrapper(module.__class__): + def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor: + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) # (3, B, H, N, C // H) + + q, k, v = torch.unbind(qkv, 0) # (B, H, N, C // H) + + x = F.scaled_dot_product_attention(q, k, v, attn_bias) + x = x.permute(0, 2, 1, 3).reshape(B, N, C) + + x = self.proj(x) + x = self.proj_drop(x) + return x + module.__class__ = _AttentionWrapper + return module + + +def sync_ddp_hook(state, bucket: torch.distributed.GradBucket) -> torch.futures.Future[torch.Tensor]: + group_to_use = torch.distributed.group.WORLD + world_size = group_to_use.size() + grad = bucket.buffer() + grad.div_(world_size) + torch.distributed.all_reduce(grad, group=group_to_use) + fut = torch.futures.Future() + fut.set_result(grad) + return fut diff --git a/MoGe/moge/model/v1.py b/MoGe/moge/model/v1.py new file mode 100644 index 0000000000000000000000000000000000000000..e0955773b389af4271bb84ad045cd48a07bbc269 --- /dev/null +++ b/MoGe/moge/model/v1.py @@ -0,0 +1,391 @@ +from typing import * +from numbers import Number +from functools import partial +from pathlib import Path +import importlib +import warnings +import json + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils +import torch.utils.checkpoint +import torch.version +import utils3d +from huggingface_hub import hf_hub_download + +from ..utils.geometry_torch import normalized_view_plane_uv, recover_focal_shift, gaussian_blur_2d +from .utils import wrap_dinov2_attention_with_sdpa, wrap_module_with_gradient_checkpointing, unwrap_module_with_gradient_checkpointing +from ..utils.tools import timeit + + +class ResidualConvBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int = None, hidden_channels: int = None, padding_mode: str = 'replicate', activation: Literal['relu', 'leaky_relu', 'silu', 'elu'] = 'relu', norm: Literal['group_norm', 'layer_norm'] = 'group_norm'): + super(ResidualConvBlock, self).__init__() + if out_channels is None: + out_channels = in_channels + if hidden_channels is None: + hidden_channels = in_channels + + if activation =='relu': + activation_cls = lambda: nn.ReLU(inplace=True) + elif activation == 'leaky_relu': + activation_cls = lambda: nn.LeakyReLU(negative_slope=0.2, inplace=True) + elif activation =='silu': + activation_cls = lambda: nn.SiLU(inplace=True) + elif activation == 'elu': + activation_cls = lambda: nn.ELU(inplace=True) + else: + raise ValueError(f'Unsupported activation function: {activation}') + + self.layers = nn.Sequential( + nn.GroupNorm(1, in_channels), + activation_cls(), + nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1, padding_mode=padding_mode), + nn.GroupNorm(hidden_channels // 32 if norm == 'group_norm' else 1, hidden_channels), + activation_cls(), + nn.Conv2d(hidden_channels, out_channels, kernel_size=3, padding=1, padding_mode=padding_mode) + ) + + self.skip_connection = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0) if in_channels != out_channels else nn.Identity() + + def forward(self, x): + skip = self.skip_connection(x) + x = self.layers(x) + x = x + skip + return x + + +class Head(nn.Module): + def __init__( + self, + num_features: int, + dim_in: int, + dim_out: List[int], + dim_proj: int = 512, + dim_upsample: List[int] = [256, 128, 128], + dim_times_res_block_hidden: int = 1, + num_res_blocks: int = 1, + res_block_norm: Literal['group_norm', 'layer_norm'] = 'group_norm', + last_res_blocks: int = 0, + last_conv_channels: int = 32, + last_conv_size: int = 1 + ): + super().__init__() + + self.projects = nn.ModuleList([ + nn.Conv2d(in_channels=dim_in, out_channels=dim_proj, kernel_size=1, stride=1, padding=0,) for _ in range(num_features) + ]) + + self.upsample_blocks = nn.ModuleList([ + nn.Sequential( + self._make_upsampler(in_ch + 2, out_ch), + *(ResidualConvBlock(out_ch, out_ch, dim_times_res_block_hidden * out_ch, activation="relu", norm=res_block_norm) for _ in range(num_res_blocks)) + ) for in_ch, out_ch in zip([dim_proj] + dim_upsample[:-1], dim_upsample) + ]) + + self.output_block = nn.ModuleList([ + self._make_output_block( + dim_upsample[-1] + 2, dim_out_, dim_times_res_block_hidden, last_res_blocks, last_conv_channels, last_conv_size, res_block_norm, + ) for dim_out_ in dim_out + ]) + + def _make_upsampler(self, in_channels: int, out_channels: int): + upsampler = nn.Sequential( + nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2), + nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, padding_mode='replicate') + ) + upsampler[0].weight.data[:] = upsampler[0].weight.data[:, :, :1, :1] + return upsampler + + def _make_output_block(self, dim_in: int, dim_out: int, dim_times_res_block_hidden: int, last_res_blocks: int, last_conv_channels: int, last_conv_size: int, res_block_norm: Literal['group_norm', 'layer_norm']): + return nn.Sequential( + nn.Conv2d(dim_in, last_conv_channels, kernel_size=3, stride=1, padding=1, padding_mode='replicate'), + *(ResidualConvBlock(last_conv_channels, last_conv_channels, dim_times_res_block_hidden * last_conv_channels, activation='relu', norm=res_block_norm) for _ in range(last_res_blocks)), + nn.ReLU(inplace=True), + nn.Conv2d(last_conv_channels, dim_out, kernel_size=last_conv_size, stride=1, padding=last_conv_size // 2, padding_mode='replicate'), + ) + + def forward(self, hidden_states: torch.Tensor, image: torch.Tensor): + img_h, img_w = image.shape[-2:] + patch_h, patch_w = img_h // 14, img_w // 14 + + # Process the hidden states + x = torch.stack([ + proj(feat.permute(0, 2, 1).unflatten(2, (patch_h, patch_w)).contiguous()) + for proj, (feat, clstoken) in zip(self.projects, hidden_states) + ], dim=1).sum(dim=1) + + # Upsample stage + # (patch_h, patch_w) -> (patch_h * 2, patch_w * 2) -> (patch_h * 4, patch_w * 4) -> (patch_h * 8, patch_w * 8) + for i, block in enumerate(self.upsample_blocks): + # UV coordinates is for awareness of image aspect ratio + uv = normalized_view_plane_uv(width=x.shape[-1], height=x.shape[-2], aspect_ratio=img_w / img_h, dtype=x.dtype, device=x.device) + uv = uv.permute(2, 0, 1).unsqueeze(0).expand(x.shape[0], -1, -1, -1) + x = torch.cat([x, uv], dim=1) + for layer in block: + x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False) + + # (patch_h * 8, patch_w * 8) -> (img_h, img_w) + x = F.interpolate(x, (img_h, img_w), mode="bilinear", align_corners=False) + uv = normalized_view_plane_uv(width=x.shape[-1], height=x.shape[-2], aspect_ratio=img_w / img_h, dtype=x.dtype, device=x.device) + uv = uv.permute(2, 0, 1).unsqueeze(0).expand(x.shape[0], -1, -1, -1) + x = torch.cat([x, uv], dim=1) + + if isinstance(self.output_block, nn.ModuleList): + output = [torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False) for block in self.output_block] + else: + output = torch.utils.checkpoint.checkpoint(self.output_block, x, use_reentrant=False) + + return output + + +class MoGeModel(nn.Module): + image_mean: torch.Tensor + image_std: torch.Tensor + + def __init__(self, + encoder: str = 'dinov2_vitb14', + intermediate_layers: Union[int, List[int]] = 4, + dim_proj: int = 512, + dim_upsample: List[int] = [256, 128, 128], + dim_times_res_block_hidden: int = 1, + num_res_blocks: int = 1, + remap_output: Literal[False, True, 'linear', 'sinh', 'exp', 'sinh_exp'] = 'linear', + res_block_norm: Literal['group_norm', 'layer_norm'] = 'group_norm', + num_tokens_range: Tuple[Number, Number] = [1200, 2500], + last_res_blocks: int = 0, + last_conv_channels: int = 32, + last_conv_size: int = 1, + mask_threshold: float = 0.5, + **deprecated_kwargs + ): + super(MoGeModel, self).__init__() + + if deprecated_kwargs: + # Process legacy arguments + if 'trained_area_range' in deprecated_kwargs: + num_tokens_range = [deprecated_kwargs['trained_area_range'][0] // 14 ** 2, deprecated_kwargs['trained_area_range'][1] // 14 ** 2] + del deprecated_kwargs['trained_area_range'] + warnings.warn(f"The following deprecated/invalid arguments are ignored: {deprecated_kwargs}") + + self.encoder = encoder + self.remap_output = remap_output + self.intermediate_layers = intermediate_layers + self.num_tokens_range = num_tokens_range + self.mask_threshold = mask_threshold + + # NOTE: We have copied the DINOv2 code in torchhub to this repository. + # Minimal modifications have been made: removing irrelevant code, unnecessary warnings and fixing importing issues. + hub_loader = getattr(importlib.import_module(".dinov2.hub.backbones", __package__), encoder) + self.backbone = hub_loader(pretrained=False) + dim_feature = self.backbone.blocks[0].attn.qkv.in_features + + self.head = Head( + num_features=intermediate_layers if isinstance(intermediate_layers, int) else len(intermediate_layers), + dim_in=dim_feature, + dim_out=[3, 1], + dim_proj=dim_proj, + dim_upsample=dim_upsample, + dim_times_res_block_hidden=dim_times_res_block_hidden, + num_res_blocks=num_res_blocks, + res_block_norm=res_block_norm, + last_res_blocks=last_res_blocks, + last_conv_channels=last_conv_channels, + last_conv_size=last_conv_size + ) + + image_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + image_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + + self.register_buffer("image_mean", image_mean) + self.register_buffer("image_std", image_std) + + if torch.__version__ >= '2.0': + self.enable_pytorch_native_sdpa() + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Union[str, Path, IO[bytes]], model_kwargs: Optional[Dict[str, Any]] = None, **hf_kwargs) -> 'MoGeModel': + """ + Load a model from a checkpoint file. + + ### Parameters: + - `pretrained_model_name_or_path`: path to the checkpoint file or repo id. + - `model_kwargs`: additional keyword arguments to override the parameters in the checkpoint. + - `hf_kwargs`: additional keyword arguments to pass to the `hf_hub_download` function. Ignored if `pretrained_model_name_or_path` is a local path. + + ### Returns: + - A new instance of `MoGe` with the parameters loaded from the checkpoint. + """ + if Path(pretrained_model_name_or_path).exists(): + checkpoint = torch.load(pretrained_model_name_or_path, map_location='cpu', weights_only=True) + else: + cached_checkpoint_path = hf_hub_download( + repo_id=pretrained_model_name_or_path, + repo_type="model", + filename="model.pt", + **hf_kwargs + ) + checkpoint = torch.load(cached_checkpoint_path, map_location='cpu', weights_only=True) + model_config = checkpoint['model_config'] + if model_kwargs is not None: + model_config.update(model_kwargs) + model = cls(**model_config) + model.load_state_dict(checkpoint['model']) + return model + + def init_weights(self): + "Load the backbone with pretrained dinov2 weights from torch hub" + state_dict = torch.hub.load('facebookresearch/dinov2', self.encoder, pretrained=True).state_dict() + self.backbone.load_state_dict(state_dict) + + def enable_gradient_checkpointing(self): + for i in range(len(self.backbone.blocks)): + self.backbone.blocks[i] = wrap_module_with_gradient_checkpointing(self.backbone.blocks[i]) + + def enable_pytorch_native_sdpa(self): + for i in range(len(self.backbone.blocks)): + self.backbone.blocks[i].attn = wrap_dinov2_attention_with_sdpa(self.backbone.blocks[i].attn) + + def _remap_points(self, points: torch.Tensor) -> torch.Tensor: + if self.remap_output == 'linear': + pass + elif self.remap_output =='sinh': + points = torch.sinh(points) + elif self.remap_output == 'exp': + xy, z = points.split([2, 1], dim=-1) + z = torch.exp(z) + points = torch.cat([xy * z, z], dim=-1) + elif self.remap_output =='sinh_exp': + xy, z = points.split([2, 1], dim=-1) + points = torch.cat([torch.sinh(xy), torch.exp(z)], dim=-1) + else: + raise ValueError(f"Invalid remap output type: {self.remap_output}") + return points + + def forward(self, image: torch.Tensor, num_tokens: int) -> Dict[str, torch.Tensor]: + original_height, original_width = image.shape[-2:] + + # Resize to expected resolution defined by num_tokens + resize_factor = ((num_tokens * 14 ** 2) / (original_height * original_width)) ** 0.5 + resized_width, resized_height = int(original_width * resize_factor), int(original_height * resize_factor) + image = F.interpolate(image, (resized_height, resized_width), mode="bicubic", align_corners=False, antialias=True) + + # Apply image transformation for DINOv2 + image = (image - self.image_mean) / self.image_std + image_14 = F.interpolate(image, (resized_height // 14 * 14, resized_width // 14 * 14), mode="bilinear", align_corners=False, antialias=True) + + # Get intermediate layers from the backbone + features = self.backbone.get_intermediate_layers(image_14, self.intermediate_layers, return_class_token=True) + + # Predict points (and mask) + output = self.head(features, image) + points, mask = output + + # Make sure fp32 precision for output + with torch.autocast(device_type=image.device.type, dtype=torch.float32): + # Resize to original resolution + points = F.interpolate(points, (original_height, original_width), mode='bilinear', align_corners=False, antialias=False) + mask = F.interpolate(mask, (original_height, original_width), mode='bilinear', align_corners=False, antialias=False) + + # Post-process points and mask + points, mask = points.permute(0, 2, 3, 1), mask.squeeze(1) + points = self._remap_points(points) # slightly improves the performance in case of very large output values + + return_dict = {'points': points, 'mask': mask} + return return_dict + + @torch.inference_mode() + def infer( + self, + image: torch.Tensor, + fov_x: Union[Number, torch.Tensor] = None, + resolution_level: int = 9, + num_tokens: int = None, + apply_mask: bool = True, + force_projection: bool = True, + use_fp16: bool = True, + ) -> Dict[str, torch.Tensor]: + """ + User-friendly inference function + + ### Parameters + - `image`: input image tensor of shape (B, 3, H, W) or (3, H, W)\ + - `fov_x`: the horizontal camera FoV in degrees. If None, it will be inferred from the predicted point map. Default: None + - `resolution_level`: An integer [0-9] for the resolution level for inference. + The higher, the finer details will be captured, but slower. Defaults to 9. Note that it is irrelevant to the output size, which is always the same as the input size. + `resolution_level` actually controls `num_tokens`. See `num_tokens` for more details. + - `num_tokens`: number of tokens used for inference. A integer in the (suggested) range of `[1200, 2500]`. + `resolution_level` will be ignored if `num_tokens` is provided. Default: None + - `apply_mask`: if True, the output point map will be masked using the predicted mask. Default: True + - `force_projection`: if True, the output point map will be recomputed to match the projection constraint. Default: True + - `use_fp16`: if True, use mixed precision to speed up inference. Default: True + + ### Returns + + A dictionary containing the following keys: + - `points`: output tensor of shape (B, H, W, 3) or (H, W, 3). + - `depth`: tensor of shape (B, H, W) or (H, W) containing the depth map. + - `intrinsics`: tensor of shape (B, 3, 3) or (3, 3) containing the camera intrinsics. + """ + if image.dim() == 3: + omit_batch_dim = True + image = image.unsqueeze(0) + else: + omit_batch_dim = False + + original_height, original_width = image.shape[-2:] + area = original_height * original_width + aspect_ratio = original_width / original_height + + if num_tokens is None: + min_tokens, max_tokens = self.num_tokens_range + num_tokens = int(min_tokens + (resolution_level / 9) * (max_tokens - min_tokens)) + + with torch.autocast(device_type=image.device.type, dtype=torch.float16, enabled=use_fp16): + output = self.forward(image, num_tokens) + points, mask = output['points'], output['mask'] + + mask_binary = mask > self.mask_threshold + + # Get camera-space point map. (Focal here is the focal length relative to half the image diagonal) + if fov_x is None: + focal, shift = recover_focal_shift(points, mask_binary) + else: + focal = aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5 / torch.tan(torch.deg2rad(torch.as_tensor(fov_x, device=points.device, dtype=points.dtype) / 2)) + if focal.ndim == 0: + focal = focal[None].expand(points.shape[0]) + _, shift = recover_focal_shift(points, mask_binary, focal=focal) + fx = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5 / aspect_ratio + fy = focal / 2 * (1 + aspect_ratio ** 2) ** 0.5 + intrinsics = utils3d.torch.intrinsics_from_focal_center(fx, fy, 0.5, 0.5) + depth = points[..., 2] + shift[..., None, None] + + # If projection constraint is forced, recompute the point map using the actual depth map + if force_projection: + points = utils3d.torch.depth_to_points(depth, intrinsics=intrinsics) + else: + points = points + torch.stack([torch.zeros_like(shift), torch.zeros_like(shift), shift], dim=-1)[..., None, None, :] + + # Apply mask if needed + if apply_mask: + points = torch.where(mask_binary[..., None], points, torch.inf) + depth = torch.where(mask_binary, depth, torch.inf) + + if omit_batch_dim: + points = points.squeeze(0) + intrinsics = intrinsics.squeeze(0) + depth = depth.squeeze(0) + mask_binary = mask_binary.squeeze(0) + mask = mask.squeeze(0) + + return_dict = { + 'points': points, + 'intrinsics': intrinsics, + 'depth': depth, + 'mask': mask_binary, + 'mask_prob': torch.sigmoid(mask) + } + + return return_dict \ No newline at end of file diff --git a/MoGe/moge/scripts/__init__.py b/MoGe/moge/scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MoGe/moge/scripts/app.py b/MoGe/moge/scripts/app.py new file mode 100644 index 0000000000000000000000000000000000000000..ccfd6686d57d8a4e825b3f29f99b2da813861e15 --- /dev/null +++ b/MoGe/moge/scripts/app.py @@ -0,0 +1,156 @@ +import os +import sys +from pathlib import Path +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) +import time +import uuid +import tempfile +from typing import * +import atexit +from concurrent.futures import ThreadPoolExecutor + +import click + + +@click.command(help='Web demo') +@click.option('--share', is_flag=True, help='Whether to run the app in shared mode.') +@click.option('--max_size', default=800, type=int, help='The maximum size of the input image.') +@click.option('--pretrained', 'pretrained_model_name_or_path', default='Ruicheng/moge-vitl', help='The name or path of the pre-trained model.') +def main(share: bool, max_size: int, pretrained_model_name_or_path: str): + # Lazy import + import cv2 + import torch + import numpy as np + import trimesh + import trimesh.visual + from PIL import Image + import gradio as gr + try: + import spaces # This is for deployment at huggingface.co/spaces + HUGGINFACE_SPACES_INSTALLED = True + except ImportError: + HUGGINFACE_SPACES_INSTALLED = False + + import utils3d + from moge.utils.vis import colorize_depth + from moge.model.v1 import MoGeModel + + + model = MoGeModel.from_pretrained(pretrained_model_name_or_path).cuda().eval() + thread_pool_executor = ThreadPoolExecutor(max_workers=1) + + def delete_later(path: Union[str, os.PathLike], delay: int = 300): + def _delete(): + try: + os.remove(path) + except: + pass + def _wait_and_delete(): + time.sleep(delay) + _delete(path) + thread_pool_executor.submit(_wait_and_delete) + atexit.register(_delete) + + # Inference on GPU. + @(spaces.GPU if HUGGINFACE_SPACES_INSTALLED else lambda x: x) + def run_with_gpu(image: np.ndarray) -> Dict[str, np.ndarray]: + image_tensor = torch.tensor(image, dtype=torch.float32, device=torch.device('cuda')).permute(2, 0, 1) / 255 + output = model.infer(image_tensor, apply_mask=True, resolution_level=9) + output = {k: v.cpu().numpy() for k, v in output.items()} + return output + + # Full inference pipeline + def run(image: np.ndarray, remove_edge: bool = True): + run_id = str(uuid.uuid4()) + + larger_size = max(image.shape[:2]) + if larger_size > max_size: + scale = max_size / larger_size + image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_AREA) + + height, width = image.shape[:2] + + output = run_with_gpu(image) + points, depth, mask = output['points'], output['depth'], output['mask'] + normals, normals_mask = utils3d.numpy.points_to_normals(points, mask=mask) + fov_x, fov_y = utils3d.numpy.intrinsics_to_fov(output['intrinsics']) + fov_x, fov_y = np.rad2deg([fov_x, fov_y]) + + faces, vertices, vertex_colors, vertex_uvs = utils3d.numpy.image_mesh( + points, + image.astype(np.float32) / 255, + utils3d.numpy.image_uv(width=width, height=height), + mask=mask & ~(utils3d.numpy.depth_edge(depth, rtol=0.03, mask=mask) & utils3d.numpy.normals_edge(normals, tol=5, mask=normals_mask)), + tri=True + ) + vertices, vertex_uvs = vertices * [1, -1, -1], vertex_uvs * [1, -1] + [0, 1] + + tempdir = Path(tempfile.gettempdir(), 'moge') + tempdir.mkdir(exist_ok=True) + + output_glb_path = Path(tempdir, f'{run_id}.glb') + output_glb_path.parent.mkdir(exist_ok=True) + trimesh.Trimesh( + vertices=vertices * [-1, 1, -1], # No idea why Gradio 3D Viewer' default camera is flipped + faces=faces, + visual = trimesh.visual.texture.TextureVisuals( + uv=vertex_uvs, + material=trimesh.visual.material.PBRMaterial( + baseColorTexture=Image.fromarray(image), + metallicFactor=0.5, + roughnessFactor=1.0 + ) + ), + process=False + ).export(output_glb_path) + + output_ply_path = Path(tempdir, f'{run_id}.ply') + output_ply_path.parent.mkdir(exist_ok=True) + trimesh.Trimesh( + vertices=vertices, + faces=faces, + vertex_colors=vertex_colors, + process=False + ).export(output_ply_path) + + colorized_depth = colorize_depth(depth) + + delete_later(output_glb_path, delay=300) + delete_later(output_ply_path, delay=300) + + return ( + colorized_depth, + output_glb_path, + output_ply_path.as_posix(), + f'Horizontal FOV: {fov_x:.2f}, Vertical FOV: {fov_y:.2f}' + ) + + gr.Interface( + fn=run, + inputs=[ + gr.Image(type="numpy", image_mode="RGB"), + gr.Checkbox(True, label="Remove edges"), + ], + outputs=[ + gr.Image(type="numpy", label="Depth map (colorized)", format='png'), + gr.Model3D(display_mode="solid", clear_color=[1.0, 1.0, 1.0, 1.0], label="3D Viewer"), + gr.File(type="filepath", label="Download the model as .ply file"), + gr.Textbox('--', label="FOV (Horizontal, Vertical)") + ], + title=None, + description=f""" +## Turn a 2D image into a 3D point map with [MoGe](https://wangrc.site/MoGePage/) + +NOTE: +* The maximum size is set to {max_size:d}px for efficiency purpose. Oversized images will be downsampled. +* The color in the 3D viewer may look dark due to rendering of 3D viewer. You may download the 3D model as .glb or .ply file to view it in other 3D viewers. +""", + clear_btn=None, + allow_flagging="never", + theme=gr.themes.Soft() + ).launch(share=share) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/MoGe/moge/scripts/cli.py b/MoGe/moge/scripts/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..45c3b9006bf56306e403f8da5b6d5068215221ee --- /dev/null +++ b/MoGe/moge/scripts/cli.py @@ -0,0 +1,27 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +from pathlib import Path +import sys +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) + +import click + + +@click.group(help='MoGe command line interface.') +def cli(): + pass + +def main(): + from moge.scripts import app, infer, infer_baseline, infer_panorama, eval_baseline, vis_data + cli.add_command(app.main, name='app') + cli.add_command(infer.main, name='infer') + cli.add_command(infer_baseline.main, name='infer_baseline') + cli.add_command(infer_panorama.main, name='infer_panorama') + cli.add_command(eval_baseline.main, name='eval_baseline') + cli.add_command(vis_data.main, name='vis_data') + cli() + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/MoGe/moge/scripts/eval_baseline.py b/MoGe/moge/scripts/eval_baseline.py new file mode 100644 index 0000000000000000000000000000000000000000..8217d9e6500b1d72a00e1a0a225ba4c2134b892e --- /dev/null +++ b/MoGe/moge/scripts/eval_baseline.py @@ -0,0 +1,165 @@ +import os +import sys +from pathlib import Path +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) +import json +from typing import * +import importlib +import importlib.util + +import click + + +@click.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, help='Evaluation script.') +@click.option('--baseline', 'baseline_code_path', type=click.Path(), required=True, help='Path to the baseline model python code.') +@click.option('--config', 'config_path', type=click.Path(), default='configs/eval/all_benchmarks.json', help='Path to the evaluation configurations. ' + 'Defaults to "configs/eval/all_benchmarks.json".') +@click.option('--output', '-o', 'output_path', type=click.Path(), required=True, help='Path to the output json file.') +@click.option('--oracle', 'oracle_mode', is_flag=True, help='Use oracle mode for evaluation, i.e., use the GT intrinsics input.') +@click.option('--dump_pred', is_flag=True, help='Dump predition results.') +@click.option('--dump_gt', is_flag=True, help='Dump ground truth.') +@click.pass_context +def main(ctx: click.Context, baseline_code_path: str, config_path: str, oracle_mode: bool, output_path: Union[str, Path], dump_pred: bool, dump_gt: bool): + # Lazy import + import cv2 + import numpy as np + from tqdm import tqdm + import torch + import torch.nn.functional as F + import utils3d + + from moge.test.baseline import MGEBaselineInterface + from moge.test.dataloader import EvalDataLoaderPipeline + from moge.test.metrics import compute_metrics + from moge.utils.geometry_torch import intrinsics_to_fov + from moge.utils.vis import colorize_depth, colorize_normal + from moge.utils.tools import key_average, flatten_nested_dict, timeit, import_file_as_module + + # Load the baseline model + module = import_file_as_module(baseline_code_path, Path(baseline_code_path).stem) + baseline_cls: Type[MGEBaselineInterface] = getattr(module, 'Baseline') + baseline : MGEBaselineInterface = baseline_cls.load.main(ctx.args, standalone_mode=False) + + # Load the evaluation configurations + with open(config_path, 'r') as f: + config = json.load(f) + + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + all_metrics = {} + # Iterate over the dataset + for benchmark_name, benchmark_config in tqdm(list(config.items()), desc='Benchmarks'): + filenames, metrics_list = [], [] + with ( + EvalDataLoaderPipeline(**benchmark_config) as eval_data_pipe, + tqdm(total=len(eval_data_pipe), desc=benchmark_name, leave=False) as pbar + ): + # Iterate over the samples in the dataset + for i in range(len(eval_data_pipe)): + sample = eval_data_pipe.get() + sample = {k: v.to(baseline.device) if isinstance(v, torch.Tensor) else v for k, v in sample.items()} + image = sample['image'] + gt_intrinsics = sample['intrinsics'] + + # Inference + torch.cuda.synchronize() + with torch.inference_mode(), timeit('_inference_timer', verbose=False) as timer: + if oracle_mode: + pred = baseline.infer_for_evaluation(image, gt_intrinsics) + else: + pred = baseline.infer_for_evaluation(image) + torch.cuda.synchronize() + + # Compute metrics + metrics, misc = compute_metrics(pred, sample, vis=dump_pred or dump_gt) + metrics['inference_time'] = timer.time + metrics_list.append(metrics) + + # Dump results + dump_path = Path(output_path.replace(".json", f"_dump"), f'{benchmark_name}', sample['filename'].replace('.zip', '')) + if dump_pred: + dump_path.joinpath('pred').mkdir(parents=True, exist_ok=True) + cv2.imwrite(str(dump_path / 'pred' / 'image.jpg'), cv2.cvtColor((image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) + + with Path(dump_path, 'pred', 'metrics.json').open('w') as f: + json.dump(metrics, f, indent=4) + + if 'pred_points' in misc: + points = misc['pred_points'].cpu().numpy() + cv2.imwrite(str(dump_path / 'pred' / 'points.exr'), cv2.cvtColor(points.astype(np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + + if 'pred_depth' in misc: + depth = misc['pred_depth'].cpu().numpy() + if 'mask' in pred: + mask = pred['mask'].cpu().numpy() + depth = np.where(mask, depth, np.inf) + cv2.imwrite(str(dump_path / 'pred' / 'depth.png'), cv2.cvtColor(colorize_depth(depth), cv2.COLOR_RGB2BGR)) + + if 'mask' in pred: + mask = pred['mask'].cpu().numpy() + cv2.imwrite(str(dump_path / 'pred' / 'mask.png'), (mask * 255).astype(np.uint8)) + + if 'normal' in pred: + normal = pred['normal'].cpu().numpy() + cv2.imwrite(str(dump_path / 'pred' / 'normal.png'), cv2.cvtColor(colorize_normal(normal), cv2.COLOR_RGB2BGR)) + + if 'intrinsics' in pred: + intrinsics = pred['intrinsics'] + fov_x, fov_y = intrinsics_to_fov(intrinsics) + with open(dump_path / 'pred' / 'fov.json', 'w') as f: + json.dump({ + 'fov_x': np.rad2deg(fov_x.item()), + 'fov_y': np.rad2deg(fov_y.item()), + 'intrinsics': intrinsics.cpu().numpy().tolist(), + }, f) + + if dump_gt: + dump_path.joinpath('gt').mkdir(parents=True, exist_ok=True) + cv2.imwrite(str(dump_path / 'gt' / 'image.jpg'), cv2.cvtColor((image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) + + if 'points' in sample: + points = sample['points'] + cv2.imwrite(str(dump_path / 'gt' / 'points.exr'), cv2.cvtColor(points.cpu().numpy().astype(np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + + if 'depth' in sample: + depth = sample['depth'] + mask = sample['depth_mask'] + cv2.imwrite(str(dump_path / 'gt' / 'depth.png'), cv2.cvtColor(colorize_depth(depth.cpu().numpy(), mask=mask.cpu().numpy()), cv2.COLOR_RGB2BGR)) + + if 'normal' in sample: + normal = sample['normal'] + cv2.imwrite(str(dump_path / 'gt' / 'normal.png'), cv2.cvtColor(colorize_normal(normal.cpu().numpy()), cv2.COLOR_RGB2BGR)) + + if 'depth_mask' in sample: + mask = sample['depth_mask'] + cv2.imwrite(str(dump_path / 'gt' /'mask.png'), (mask.cpu().numpy() * 255).astype(np.uint8)) + + if 'intrinsics' in sample: + intrinsics = sample['intrinsics'] + fov_x, fov_y = intrinsics_to_fov(intrinsics) + with open(dump_path / 'gt' / 'info.json', 'w') as f: + json.dump({ + 'fov_x': np.rad2deg(fov_x.item()), + 'fov_y': np.rad2deg(fov_y.item()), + 'intrinsics': intrinsics.cpu().numpy().tolist(), + }, f) + + # Save intermediate results + if i % 100 == 0 or i == len(eval_data_pipe) - 1: + Path(output_path).write_text( + json.dumps({ + **all_metrics, + benchmark_name: key_average(metrics_list) + }, indent=4) + ) + pbar.update(1) + + all_metrics[benchmark_name] = key_average(metrics_list) + + # Save final results + all_metrics['mean'] = key_average(list(all_metrics.values())) + Path(output_path).write_text(json.dumps(all_metrics, indent=4)) + + +if __name__ == '__main__': + main() diff --git a/MoGe/moge/scripts/infer.py b/MoGe/moge/scripts/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..f58defce98b6ebfeb9070646e745abeec48b799a --- /dev/null +++ b/MoGe/moge/scripts/infer.py @@ -0,0 +1,140 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +from pathlib import Path +import sys +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) +from typing import * +import itertools +import json +import warnings + +import cv2 +import numpy as np +import torch +from PIL import Image +from tqdm import tqdm +import trimesh +import trimesh.visual +import click + +from moge.model.v1 import MoGeModel +from moge.utils.io import save_glb, save_ply +from moge.utils.vis import colorize_depth, colorize_normal +import utils3d + + +@click.command(help='Inference script') +@click.option('--input', '-i', 'input_path', type=click.Path(exists=True), help='Input image or folder path. "jpg" and "png" are supported.') +@click.option('--fov_x', 'fov_x_', type=float, default=None, help='If camera parameters are known, set the horizontal field of view in degrees. Otherwise, MoGe will estimate it.') +@click.option('--output', '-o', 'output_path', default='./output', type=click.Path(), help='Output folder path') +@click.option('--pretrained', 'pretrained_model_name_or_path', type=str, default='Ruicheng/moge-vitl', help='Pretrained model name or path. Defaults to "Ruicheng/moge-vitl"') +@click.option('--device', 'device_name', type=str, default='cuda', help='Device name (e.g. "cuda", "cuda:0", "cpu"). Defaults to "cuda"') +@click.option('--fp16', 'use_fp16', is_flag=True, help='Use fp16 precision for 2x faster inference.') +@click.option('--resize', 'resize_to', type=int, default=None, help='Resize the image(s) & output maps to a specific size. Defaults to None (no resizing).') +@click.option('--resolution_level', type=int, default=9, help='An integer [0-9] for the resolution level for inference. \ +Higher value means more tokens and the finer details will be captured, but inference can be slower. \ +Defaults to 9. Note that it is irrelevant to the output size, which is always the same as the input size. \ +`resolution_level` actually controls `num_tokens`. See `num_tokens` for more details.') +@click.option('--num_tokens', type=int, default=None, help='number of tokens used for inference. A integer in the (suggested) range of `[1200, 2500]`. \ +`resolution_level` will be ignored if `num_tokens` is provided. Default: None') +@click.option('--threshold', type=float, default=0.03, help='Threshold for removing edges. Defaults to 0.03. Smaller value removes more edges. "inf" means no thresholding.') +@click.option('--maps', 'save_maps_', is_flag=True, help='Whether to save the output maps and fov(image, depth, mask, points, fov).') +@click.option('--glb', 'save_glb_', is_flag=True, help='Whether to save the output as a.glb file. The color will be saved as a texture.') +@click.option('--ply', 'save_ply_', is_flag=True, help='Whether to save the output as a.ply file. The color will be saved as vertex colors.') +@click.option('--show', 'show', is_flag=True, help='Whether show the output in a window. Note that this requires pyglet<2 installed as required by trimesh.') +def main( + input_path: str, + fov_x_: float, + output_path: str, + pretrained_model_name_or_path: str, + device_name: str, + use_fp16: bool, + resize_to: int, + resolution_level: int, + num_tokens: int, + threshold: float, + save_maps_: bool, + save_glb_: bool, + save_ply_: bool, + show: bool, +): + device = torch.device(device_name) + + include_suffices = ['jpg', 'png', 'jpeg', 'JPG', 'PNG', 'JPEG'] + if Path(input_path).is_dir(): + image_paths = sorted(itertools.chain(*(Path(input_path).rglob(f'*.{suffix}') for suffix in include_suffices))) + else: + image_paths = [Path(input_path)] + + if len(image_paths) == 0: + raise FileNotFoundError(f'No image files found in {input_path}') + + model = MoGeModel.from_pretrained(pretrained_model_name_or_path).to(device).eval() + + + if not any([save_maps_, save_glb_, save_ply_]): + warnings.warn('No output format specified. Defaults to saving all. Please use "--maps", "--glb", or "--ply" to specify the output.') + save_maps_ = save_glb_ = save_ply_ = True + + for image_path in (pbar := tqdm(image_paths, desc='Inference', disable=len(image_paths) <= 1)): + image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB) + height, width = image.shape[:2] + if resize_to is not None: + height, width = min(resize_to, int(resize_to * height / width)), min(resize_to, int(resize_to * width / height)) + image = cv2.resize(image, (width, height), cv2.INTER_AREA) + image_tensor = torch.tensor(image / 255, dtype=torch.float32, device=device).permute(2, 0, 1) + + # Inference + output = model.infer(image_tensor, fov_x=fov_x_, resolution_level=resolution_level, num_tokens=num_tokens, use_fp16=use_fp16) + points, depth, mask, intrinsics = output['points'].cpu().numpy(), output['depth'].cpu().numpy(), output['mask'].cpu().numpy(), output['intrinsics'].cpu().numpy() + normals, normals_mask = utils3d.numpy.points_to_normals(points, mask=mask) + + save_path = Path(output_path, image_path.relative_to(input_path).parent, image_path.stem) + save_path.mkdir(exist_ok=True, parents=True) + + # Save images / maps + if save_maps_: + cv2.imwrite(str(save_path / 'image.jpg'), cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_path / 'depth_vis.png'), cv2.cvtColor(colorize_depth(depth), cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_path / 'depth.exr'), depth, [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + cv2.imwrite(str(save_path / 'mask.png'), (mask * 255).astype(np.uint8)) + cv2.imwrite(str(save_path / 'points.exr'), cv2.cvtColor(points, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + fov_x, fov_y = utils3d.numpy.intrinsics_to_fov(intrinsics) + with open(save_path / 'fov.json', 'w') as f: + json.dump({ + 'fov_x': round(float(np.rad2deg(fov_x)), 2), + 'fov_y': round(float(np.rad2deg(fov_y)), 2), + }, f) + + # Export mesh & visulization + if save_glb_ or save_ply_ or show: + faces, vertices, vertex_colors, vertex_uvs = utils3d.numpy.image_mesh( + points, + image.astype(np.float32) / 255, + utils3d.numpy.image_uv(width=width, height=height), + mask=mask & ~(utils3d.numpy.depth_edge(depth, rtol=threshold, mask=mask) & utils3d.numpy.normals_edge(normals, tol=5, mask=normals_mask)), + tri=True + ) + # When exporting the model, follow the OpenGL coordinate conventions: + # - world coordinate system: x right, y up, z backward. + # - texture coordinate system: (0, 0) for left-bottom, (1, 1) for right-top. + vertices, vertex_uvs = vertices * [1, -1, -1], vertex_uvs * [1, -1] + [0, 1] + + if save_glb_: + save_glb(save_path / 'mesh.glb', vertices, faces, vertex_uvs, image) + + if save_ply_: + save_ply(save_path / 'mesh.ply', vertices, faces, vertex_colors) + + if show: + trimesh.Trimesh( + vertices=vertices, + vertex_colors=vertex_colors, + faces=faces, + process=False + ).show() + + +if __name__ == '__main__': + main() diff --git a/MoGe/moge/scripts/infer_baseline.py b/MoGe/moge/scripts/infer_baseline.py new file mode 100644 index 0000000000000000000000000000000000000000..5409674f7cd5ce21de9200fd9038cb7d71c99e0f --- /dev/null +++ b/MoGe/moge/scripts/infer_baseline.py @@ -0,0 +1,140 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +from pathlib import Path +import sys +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) +import json +from pathlib import Path +from typing import * +import itertools +import warnings + +import click + + +@click.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, help='Inference script for wrapped baselines methods') +@click.option('--baseline', 'baseline_code_path', required=True, type=click.Path(), help='Path to the baseline model python code.') +@click.option('--input', '-i', 'input_path', type=str, required=True, help='Input image or folder') +@click.option('--output', '-o', 'output_path', type=str, default='./output', help='Output folder') +@click.option('--size', 'image_size', type=int, default=None, help='Resize input image') +@click.option('--skip', is_flag=True, help='Skip existing output') +@click.option('--maps', 'save_maps_', is_flag=True, help='Save output point / depth maps') +@click.option('--ply', 'save_ply_', is_flag=True, help='Save mesh in PLY format') +@click.option('--glb', 'save_glb_', is_flag=True, help='Save mesh in GLB format') +@click.option('--threshold', type=float, default=0.03, help='Depth edge detection threshold for saving mesh') +@click.pass_context +def main(ctx: click.Context, baseline_code_path: str, input_path: str, output_path: str, image_size: int, skip: bool, save_maps_, save_ply_: bool, save_glb_: bool, threshold: float): + # Lazy import + import cv2 + import numpy as np + from tqdm import tqdm + import torch + import utils3d + + from moge.utils.io import save_ply, save_glb + from moge.utils.geometry_numpy import intrinsics_to_fov_numpy + from moge.utils.vis import colorize_depth, colorize_depth_affine, colorize_disparity + from moge.utils.tools import key_average, flatten_nested_dict, timeit, import_file_as_module + from moge.test.baseline import MGEBaselineInterface + + # Load the baseline model + module = import_file_as_module(baseline_code_path, Path(baseline_code_path).stem) + baseline_cls: Type[MGEBaselineInterface] = getattr(module, 'Baseline') + baseline : MGEBaselineInterface = baseline_cls.load.main(ctx.args, standalone_mode=False) + + # Input images list + include_suffices = ['jpg', 'png', 'jpeg', 'JPG', 'PNG', 'JPEG'] + if Path(input_path).is_dir(): + image_paths = sorted(itertools.chain(*(Path(input_path).rglob(f'*.{suffix}') for suffix in include_suffices))) + else: + image_paths = [Path(input_path)] + + if not any([save_maps_, save_glb_, save_ply_]): + warnings.warn('No output format specified. Defaults to saving maps only. Please use "--maps", "--glb", or "--ply" to specify the output.') + save_maps_ = True + + for image_path in (pbar := tqdm(image_paths, desc='Inference', disable=len(image_paths) <= 1)): + # Load one image at a time + image_np = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB) + height, width = image_np.shape[:2] + if image_size is not None and max(image_np.shape[:2]) > image_size: + height, width = min(image_size, int(image_size * height / width)), min(image_size, int(image_size * width / height)) + image_np = cv2.resize(image_np, (width, height), cv2.INTER_AREA) + image = torch.from_numpy(image_np.astype(np.float32) / 255.0).permute(2, 0, 1).to(baseline.device) + + # Inference + torch.cuda.synchronize() + with torch.inference_mode(), (timer := timeit('Inference', verbose=False, average=True)): + output = baseline.infer(image) + torch.cuda.synchronize() + + inference_time = timer.average_time + pbar.set_postfix({'average inference time': f'{inference_time:.3f}s'}) + + # Save the output + save_path = Path(output_path, image_path.relative_to(input_path).parent, image_path.stem) + if skip and save_path.exists(): + continue + save_path.mkdir(parents=True, exist_ok=True) + + if save_maps_: + cv2.imwrite(str(save_path / 'image.jpg'), cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)) + + if 'mask' in output: + mask = output['mask'].cpu().numpy() + cv2.imwrite(str(save_path /'mask.png'), (mask * 255).astype(np.uint8)) + + for k in ['points_metric', 'points_scale_invariant', 'points_affine_invariant']: + if k in output: + points = output[k].cpu().numpy() + cv2.imwrite(str(save_path / f'{k}.exr'), cv2.cvtColor(points, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + + for k in ['depth_metric', 'depth_scale_invariant', 'depth_affine_invariant', 'disparity_affine_invariant']: + if k in output: + depth = output[k].cpu().numpy() + cv2.imwrite(str(save_path / f'{k}.exr'), depth, [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + if k in ['depth_metric', 'depth_scale_invariant']: + depth_vis = colorize_depth(depth) + elif k == 'depth_affine_invariant': + depth_vis = colorize_depth_affine(depth) + elif k == 'disparity_affine_invariant': + depth_vis = colorize_disparity(depth) + cv2.imwrite(str(save_path / f'{k}_vis.png'), cv2.cvtColor(depth_vis, cv2.COLOR_RGB2BGR)) + + if 'intrinsics' in output: + intrinsics = output['intrinsics'].cpu().numpy() + fov_x, fov_y = intrinsics_to_fov_numpy(intrinsics) + with open(save_path / 'fov.json', 'w') as f: + json.dump({ + 'fov_x': float(np.rad2deg(fov_x)), + 'fov_y': float(np.rad2deg(fov_y)), + 'intrinsics': intrinsics.tolist() + }, f, indent=4) + + # Export mesh & visulization + if save_ply_ or save_glb_: + assert any(k in output for k in ['points_metric', 'points_scale_invariant', 'points_affine_invariant']), 'No point map found in output' + points = next(output[k] for k in ['points_metric', 'points_scale_invariant', 'points_affine_invariant'] if k in output).cpu().numpy() + mask = output['mask'] if 'mask' in output else np.ones_like(points[..., 0], dtype=bool) + normals, normals_mask = utils3d.numpy.points_to_normals(points, mask=mask) + faces, vertices, vertex_colors, vertex_uvs = utils3d.numpy.image_mesh( + points, + image_np.astype(np.float32) / 255, + utils3d.numpy.image_uv(width=width, height=height), + mask=mask & ~(utils3d.numpy.depth_edge(depth, rtol=threshold, mask=mask) & utils3d.numpy.normals_edge(normals, tol=5, mask=normals_mask)), + tri=True + ) + # When exporting the model, follow the OpenGL coordinate conventions: + # - world coordinate system: x right, y up, z backward. + # - texture coordinate system: (0, 0) for left-bottom, (1, 1) for right-top. + vertices, vertex_uvs = vertices * [1, -1, -1], vertex_uvs * [1, -1] + [0, 1] + + if save_glb_: + save_glb(save_path / 'mesh.glb', vertices, faces, vertex_uvs, image_np) + + if save_ply_: + save_ply(save_path / 'mesh.ply', vertices, faces, vertex_colors) + +if __name__ == '__main__': + main() diff --git a/MoGe/moge/scripts/infer_panorama.py b/MoGe/moge/scripts/infer_panorama.py new file mode 100644 index 0000000000000000000000000000000000000000..cce65cb90cd1c6750d42cdda4e72d4ce3a2c0549 --- /dev/null +++ b/MoGe/moge/scripts/infer_panorama.py @@ -0,0 +1,162 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +from pathlib import Path +import sys +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) +from typing import * +import itertools +import json +import warnings + +import click + + +@click.command(help='Inference script for panorama images') +@click.option('--input', '-i', 'input_path', type=click.Path(exists=True), required=True, help='Input image or folder path. "jpg" and "png" are supported.') +@click.option('--output', '-o', 'output_path', type=click.Path(), default='./output', help='Output folder path') +@click.option('--pretrained', 'pretrained_model_name_or_path', type=str, default='Ruicheng/moge-vitl', help='Pretrained model name or path. Defaults to "Ruicheng/moge-vitl"') +@click.option('--device', 'device_name', type=str, default='cuda', help='Device name (e.g. "cuda", "cuda:0", "cpu"). Defaults to "cuda"') +@click.option('--resize', 'resize_to', type=int, default=None, help='Resize the image(s) & output maps to a specific size. Defaults to None (no resizing).') +@click.option('--resolution_level', type=int, default=9, help='An integer [0-9] for the resolution level of inference. The higher, the better but slower. Defaults to 9. Note that it is irrelevant to the output resolution.') +@click.option('--threshold', type=float, default=0.03, help='Threshold for removing edges. Defaults to 0.03. Smaller value removes more edges. "inf" means no thresholding.') +@click.option('--batch_size', type=int, default=4, help='Batch size for inference. Defaults to 4.') +@click.option('--splitted', 'save_splitted', is_flag=True, help='Whether to save the splitted images. Defaults to False.') +@click.option('--maps', 'save_maps_', is_flag=True, help='Whether to save the output maps and fov(image, depth, mask, points, fov).') +@click.option('--glb', 'save_glb_', is_flag=True, help='Whether to save the output as a.glb file. The color will be saved as a texture.') +@click.option('--ply', 'save_ply_', is_flag=True, help='Whether to save the output as a.ply file. The color will be saved as vertex colors.') +@click.option('--show', 'show', is_flag=True, help='Whether show the output in a window. Note that this requires pyglet<2 installed as required by trimesh.') +def main( + input_path: str, + output_path: str, + pretrained_model_name_or_path: str, + device_name: str, + resize_to: int, + resolution_level: int, + threshold: float, + batch_size: int, + save_splitted: bool, + save_maps_: bool, + save_glb_: bool, + save_ply_: bool, + show: bool, +): + # Lazy import + import cv2 + import numpy as np + from numpy import ndarray + import torch + from PIL import Image + from tqdm import tqdm, trange + import trimesh + import trimesh.visual + from scipy.sparse import csr_array, hstack, vstack + from scipy.ndimage import convolve + from scipy.sparse.linalg import lsmr + + import utils3d + from moge.model.v1 import MoGeModel + from moge.utils.io import save_glb, save_ply + from moge.utils.vis import colorize_depth + from moge.utils.panorama import spherical_uv_to_directions, get_panorama_cameras, split_panorama_image, merge_panorama_depth + + + device = torch.device(device_name) + + include_suffices = ['jpg', 'png', 'jpeg', 'JPG', 'PNG', 'JPEG'] + if Path(input_path).is_dir(): + image_paths = sorted(itertools.chain(*(Path(input_path).rglob(f'*.{suffix}') for suffix in include_suffices))) + else: + image_paths = [Path(input_path)] + + if len(image_paths) == 0: + raise FileNotFoundError(f'No image files found in {input_path}') + + # Write outputs + if not any([save_maps_, save_glb_, save_ply_]): + warnings.warn('No output format specified. Defaults to saving all. Please use "--maps", "--glb", or "--ply" to specify the output.') + save_maps_ = save_glb_ = save_ply_ = True + + model = MoGeModel.from_pretrained(pretrained_model_name_or_path).to(device).eval() + + for image_path in (pbar := tqdm(image_paths, desc='Total images', disable=len(image_paths) <= 1)): + image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB) + height, width = image.shape[:2] + if resize_to is not None: + height, width = min(resize_to, int(resize_to * height / width)), min(resize_to, int(resize_to * width / height)) + image = cv2.resize(image, (width, height), cv2.INTER_AREA) + + splitted_extrinsics, splitted_intriniscs = get_panorama_cameras() + splitted_resolution = 512 + splitted_images = split_panorama_image(image, splitted_extrinsics, splitted_intriniscs, splitted_resolution) + + # Infer each view + print('Inferring...') if pbar.disable else pbar.set_postfix_str(f'Inferring') + + splitted_distance_maps, splitted_masks = [], [] + for i in trange(0, len(splitted_images), batch_size, desc='Inferring splitted views', disable=len(splitted_images) <= batch_size, leave=False): + image_tensor = torch.tensor(np.stack(splitted_images[i:i + batch_size]) / 255, dtype=torch.float32, device=device).permute(0, 3, 1, 2) + fov_x, fov_y = np.rad2deg(utils3d.numpy.intrinsics_to_fov(np.array(splitted_intriniscs[i:i + batch_size]))) + fov_x = torch.tensor(fov_x, dtype=torch.float32, device=device) + output = model.infer(image_tensor, fov_x=fov_x, apply_mask=False) + distance_map, mask = output['points'].norm(dim=-1).cpu().numpy(), output['mask'].cpu().numpy() + splitted_distance_maps.extend(list(distance_map)) + splitted_masks.extend(list(mask)) + + # Save splitted + if save_splitted: + splitted_save_path = Path(output_path, image_path.stem, 'splitted') + splitted_save_path.mkdir(exist_ok=True, parents=True) + for i in range(len(splitted_images)): + cv2.imwrite(str(splitted_save_path / f'{i:02d}.jpg'), cv2.cvtColor(splitted_images[i], cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(splitted_save_path / f'{i:02d}_distance_vis.png'), cv2.cvtColor(colorize_depth(splitted_distance_maps[i], splitted_masks[i]), cv2.COLOR_RGB2BGR)) + + # Merge + print('Merging...') if pbar.disable else pbar.set_postfix_str(f'Merging') + + merging_width, merging_height = min(1920, width), min(960, height) + panorama_depth, panorama_mask = merge_panorama_depth(merging_width, merging_height, splitted_distance_maps, splitted_masks, splitted_extrinsics, splitted_intriniscs) + panorama_depth = panorama_depth.astype(np.float32) + panorama_depth = cv2.resize(panorama_depth, (width, height), cv2.INTER_LINEAR) + panorama_mask = cv2.resize(panorama_mask.astype(np.uint8), (width, height), cv2.INTER_NEAREST) > 0 + points = panorama_depth[:, :, None] * spherical_uv_to_directions(utils3d.numpy.image_uv(width=width, height=height)) + + # Write outputs + print('Writing outputs...') if pbar.disable else pbar.set_postfix_str(f'Inferring') + save_path = Path(output_path, image_path.relative_to(input_path).parent, image_path.stem) + save_path.mkdir(exist_ok=True, parents=True) + if save_maps_: + cv2.imwrite(str(save_path / 'image.jpg'), cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_path / 'depth_vis.png'), cv2.cvtColor(colorize_depth(panorama_depth, mask=panorama_mask), cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_path / 'depth.exr'), panorama_depth, [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + cv2.imwrite(str(save_path / 'points.exr'), points, [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + cv2.imwrite(str(save_path /'mask.png'), (panorama_mask * 255).astype(np.uint8)) + + # Export mesh & visulization + if save_glb_ or save_ply_ or show: + normals, normals_mask = utils3d.numpy.points_to_normals(points, panorama_mask) + faces, vertices, vertex_colors, vertex_uvs = utils3d.numpy.image_mesh( + points, + image.astype(np.float32) / 255, + utils3d.numpy.image_uv(width=width, height=height), + mask=panorama_mask & ~(utils3d.numpy.depth_edge(panorama_depth, rtol=threshold) & utils3d.numpy.normals_edge(normals, tol=5, mask=normals_mask)), + tri=True + ) + + if save_glb_: + save_glb(save_path / 'mesh.glb', vertices, faces, vertex_uvs, image) + + if save_ply_: + save_ply(save_path / 'mesh.ply', vertices, faces, vertex_colors) + + if show: + trimesh.Trimesh( + vertices=vertices, + vertex_colors=vertex_colors, + faces=faces, + process=False + ).show() + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/MoGe/moge/scripts/train.py b/MoGe/moge/scripts/train.py new file mode 100644 index 0000000000000000000000000000000000000000..d96d3ad4d31ed0b6c30bbbbcd83033b227e90829 --- /dev/null +++ b/MoGe/moge/scripts/train.py @@ -0,0 +1,452 @@ +import os +from pathlib import Path +import sys +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) +import json +import time +import random +from typing import * +import itertools +from contextlib import nullcontext +from concurrent.futures import ThreadPoolExecutor +import io + +import numpy as np +import cv2 +from PIL import Image +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.version +import accelerate +from accelerate import Accelerator, DistributedDataParallelKwargs +from accelerate.utils import set_seed +import utils3d +import click +from tqdm import tqdm, trange +import mlflow +torch.backends.cudnn.benchmark = False # Varying input size, make sure cudnn benchmark is disabled + +from moge.train.dataloader import TrainDataLoaderPipeline +from moge.train.losses import ( + affine_invariant_global_loss, + affine_invariant_local_loss, + edge_loss, + normal_loss, + mask_l2_loss, + mask_bce_loss, + monitoring, +) +from moge.train.utils import build_optimizer, build_lr_scheduler +from moge.utils.geometry_torch import intrinsics_to_fov +from moge.utils.vis import colorize_depth, colorize_normal +from moge.utils.tools import key_average, recursive_replace, CallbackOnException, flatten_nested_dict +from moge.test.metrics import compute_metrics + + +@click.command() +@click.option('--config', 'config_path', type=str, default='configs/debug.json') +@click.option('--workspace', type=str, default='workspace/debug', help='Path to the workspace') +@click.option('--checkpoint', 'checkpoint_path', type=str, default=None, help='Path to the checkpoint to load') +@click.option('--batch_size_forward', type=int, default=8, help='Batch size for each forward pass on each device') +@click.option('--gradient_accumulation_steps', type=int, default=1, help='Number of steps to accumulate gradients') +@click.option('--enable_gradient_checkpointing', type=bool, default=True, help='Use gradient checkpointing in backbone') +@click.option('--enable_mixed_precision', type=bool, default=False, help='Use mixed precision training. Backbone is converted to FP16') +@click.option('--enable_ema', type=bool, default=True, help='Maintain an exponential moving average of the model weights') +@click.option('--num_iterations', type=int, default=1000000, help='Number of iterations to train the model') +@click.option('--save_every', type=int, default=10000, help='Save checkpoint every n iterations') +@click.option('--log_every', type=int, default=1000, help='Log metrics every n iterations') +@click.option('--vis_every', type=int, default=0, help='Visualize every n iterations') +@click.option('--num_vis_images', type=int, default=32, help='Number of images to visualize, must be a multiple of divided batch size') +@click.option('--enable_mlflow', type=bool, default=True, help='Log metrics to MLFlow') +@click.option('--seed', type=int, default=0, help='Random seed') +def main( + config_path: str, + workspace: str, + checkpoint_path: str, + batch_size_forward: int, + gradient_accumulation_steps: int, + enable_gradient_checkpointing: bool, + enable_mixed_precision: bool, + enable_ema: bool, + num_iterations: int, + save_every: int, + log_every: int, + vis_every: int, + num_vis_images: int, + enable_mlflow: bool, + seed: Optional[int], +): + # Load config + with open(config_path, 'r') as f: + config = json.load(f) + + accelerator = Accelerator( + gradient_accumulation_steps=gradient_accumulation_steps, + mixed_precision='fp16' if enable_mixed_precision else None, + kwargs_handlers=[ + DistributedDataParallelKwargs(find_unused_parameters=True) + ] + ) + device = accelerator.device + batch_size_total = batch_size_forward * gradient_accumulation_steps * accelerator.num_processes + + # Log config + if accelerator.is_main_process: + if enable_mlflow: + try: + mlflow.log_params({ + **click.get_current_context().params, + 'batch_size_total': batch_size_total, + }) + except: + print('Failed to log config to MLFlow') + Path(workspace).mkdir(parents=True, exist_ok=True) + with Path(workspace).joinpath('config.json').open('w') as f: + json.dump(config, f, indent=4) + + # Set seed + if seed is not None: + set_seed(seed, device_specific=True) + + # Initialize model + print('Initialize model') + with accelerator.local_main_process_first(): + from moge.model import import_model_class_by_version + MoGeModel = import_model_class_by_version(config['model_version']) + model = MoGeModel(**config['model']) + count_total_parameters = sum(p.numel() for p in model.parameters()) + print(f'Total parameters: {count_total_parameters}') + + # Set up EMA model + if enable_ema and accelerator.is_main_process: + ema_avg_fn = lambda averaged_model_parameter, model_parameter, num_averaged: 0.999 * averaged_model_parameter + 0.001 * model_parameter + ema_model = torch.optim.swa_utils.AveragedModel(model, device=accelerator.device, avg_fn=ema_avg_fn) + + # Set gradient checkpointing + if enable_gradient_checkpointing: + model.enable_gradient_checkpointing() + import warnings + warnings.filterwarnings("ignore", category=FutureWarning, module="torch.utils.checkpoint") + + # Initalize optimizer & lr scheduler + optimizer = build_optimizer(model, config['optimizer']) + lr_scheduler = build_lr_scheduler(optimizer, config['lr_scheduler']) + + count_grouped_parameters = [sum(p.numel() for p in param_group['params'] if p.requires_grad) for param_group in optimizer.param_groups] + for i, count in enumerate(count_grouped_parameters): + print(f'- Group {i}: {count} parameters') + + # Attempt to load checkpoint + checkpoint: Dict[str, Any] + with accelerator.local_main_process_first(): + if checkpoint_path.endswith('.pt'): + # - Load specific checkpoint file + print(f'Load checkpoint: {checkpoint_path}') + checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True) + elif checkpoint_path == "latest": + # - Load latest + checkpoint_path = Path(workspace, 'checkpoint', 'latest.pt') + if checkpoint_path.exists(): + print(f'Load checkpoint: {checkpoint_path}') + checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True) + i_step = checkpoint['step'] + if 'model' not in checkpoint and (checkpoint_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}.pt')).exists(): + print(f'Load model checkpoint: {checkpoint_model_path}') + checkpoint['model'] = torch.load(checkpoint_model_path, map_location='cpu', weights_only=True)['model'] + if 'optimizer' not in checkpoint and (checkpoint_optimizer_path := Path(workspace, 'checkpoint', f'{i_step:08d}_optimizer.pt')).exists(): + print(f'Load optimizer checkpoint: {checkpoint_optimizer_path}') + checkpoint.update(torch.load(checkpoint_optimizer_path, map_location='cpu', weights_only=True)) + if enable_ema and accelerator.is_main_process: + if 'ema_model' not in checkpoint and (checkpoint_ema_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}_ema.pt')).exists(): + print(f'Load EMA model checkpoint: {checkpoint_ema_model_path}') + checkpoint['ema_model'] = torch.load(checkpoint_ema_model_path, map_location='cpu', weights_only=True)['model'] + else: + checkpoint = None + elif checkpoint_path is not None: + # - Load by step number + i_step = int(checkpoint_path) + checkpoint = {'step': i_step} + if (checkpoint_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}.pt')).exists(): + print(f'Load model checkpoint: {checkpoint_model_path}') + checkpoint['model'] = torch.load(checkpoint_model_path, map_location='cpu', weights_only=True)['model'] + if (checkpoint_optimizer_path := Path(workspace, 'checkpoint', f'{i_step:08d}_optimizer.pt')).exists(): + print(f'Load optimizer checkpoint: {checkpoint_optimizer_path}') + checkpoint.update(torch.load(checkpoint_optimizer_path, map_location='cpu', weights_only=True)) + if enable_ema and accelerator.is_main_process: + if (checkpoint_ema_model_path := Path(workspace, 'checkpoint', f'{i_step:08d}_ema.pt')).exists(): + print(f'Load EMA model checkpoint: {checkpoint_ema_model_path}') + checkpoint['ema_model'] = torch.load(checkpoint_ema_model_path, map_location='cpu', weights_only=True)['model'] + else: + checkpoint = None + + if checkpoint is None: + # Initialize model weights + print('Initialize model weights') + with accelerator.local_main_process_first(): + model.init_weights() + initial_step = 0 + else: + model.load_state_dict(checkpoint['model'], strict=False) + if 'step' in checkpoint: + initial_step = checkpoint['step'] + 1 + else: + initial_step = 0 + if 'optimizer' in checkpoint: + optimizer.load_state_dict(checkpoint['optimizer']) + if enable_ema and accelerator.is_main_process and 'ema_model' in checkpoint: + ema_model.module.load_state_dict(checkpoint['ema_model'], strict=False) + if 'lr_scheduler' in checkpoint: + lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) + + del checkpoint + + model, optimizer = accelerator.prepare(model, optimizer) + if torch.version.hip and isinstance(model, torch.nn.parallel.DistributedDataParallel): + # Hacking potential gradient synchronization issue in ROCm backend + from moge.model.utils import sync_ddp_hook + model.register_comm_hook(None, sync_ddp_hook) + + # Initialize training data pipeline + with accelerator.local_main_process_first(): + train_data_pipe = TrainDataLoaderPipeline(config['data'], batch_size_forward) + + def _write_bytes_retry_loop(save_path: Path, data: bytes): + while True: + try: + save_path.write_bytes(data) + break + except Exception as e: + print('Error while saving checkpoint, retrying in 1 minute: ', e) + time.sleep(60) + + # Ready to train + records = [] + model.train() + with ( + train_data_pipe, + tqdm(initial=initial_step, total=num_iterations, desc='Training', disable=not accelerator.is_main_process) as pbar, + ThreadPoolExecutor(max_workers=1) as save_checkpoint_executor, + ): + # Get some batches for visualization + if accelerator.is_main_process: + batches_for_vis: List[Dict[str, torch.Tensor]] = [] + num_vis_images = num_vis_images // batch_size_forward * batch_size_forward + for _ in range(num_vis_images // batch_size_forward): + batch = train_data_pipe.get() + batches_for_vis.append(batch) + + # Visualize GT + if vis_every > 0 and accelerator.is_main_process and initial_step == 0: + save_dir = Path(workspace).joinpath('vis/gt') + for i_batch, batch in enumerate(tqdm(batches_for_vis, desc='Visualize GT', leave=False)): + image, gt_depth, gt_mask, gt_mask_inf, gt_intrinsics, info = batch['image'], batch['depth'], batch['depth_mask'], batch['depth_mask_inf'], batch['intrinsics'], batch['info'] + gt_points = utils3d.torch.depth_to_points(gt_depth, intrinsics=gt_intrinsics) + gt_normal, gt_normal_mask = utils3d.torch.points_to_normals(gt_points, gt_mask) + for i_instance in range(batch['image'].shape[0]): + idx = i_batch * batch_size_forward + i_instance + image_i = (image[i_instance].numpy().transpose(1, 2, 0) * 255).astype(np.uint8) + gt_depth_i = gt_depth[i_instance].numpy() + gt_mask_i = gt_mask[i_instance].numpy() + gt_mask_inf_i = gt_mask_inf[i_instance].numpy() + gt_points_i = gt_points[i_instance].numpy() + gt_normal_i = gt_normal[i_instance].numpy() + save_dir.joinpath(f'{idx:04d}').mkdir(parents=True, exist_ok=True) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/image.jpg')), cv2.cvtColor(image_i, cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/points.exr')), cv2.cvtColor(gt_points_i, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/mask.png')), gt_mask_i * 255) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/depth_vis.png')), cv2.cvtColor(colorize_depth(gt_depth_i, gt_mask_i), cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/normal.png')), cv2.cvtColor(colorize_normal(gt_normal_i), cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/mask_inf.png')), gt_mask_inf_i * 255) + with save_dir.joinpath(f'{idx:04d}/info.json').open('w') as f: + json.dump(info[i_instance], f) + + # Reset seed to avoid training on the same data when resuming training + if seed is not None: + set_seed(seed + initial_step, device_specific=True) + + # Training loop + for i_step in range(initial_step, num_iterations): + + i_accumulate, weight_accumulate = 0, 0 + while i_accumulate < gradient_accumulation_steps: + # Load batch + batch = train_data_pipe.get() + image, gt_depth, gt_mask, gt_mask_fin, gt_mask_inf, gt_intrinsics, label_type, is_metric = batch['image'], batch['depth'], batch['depth_mask'], batch['depth_mask_fin'], batch['depth_mask_inf'], batch['intrinsics'], batch['label_type'], batch['is_metric'] + image, gt_depth, gt_mask, gt_mask_fin, gt_mask_inf, gt_intrinsics = image.to(device), gt_depth.to(device), gt_mask.to(device), gt_mask_fin.to(device), gt_mask_inf.to(device), gt_intrinsics.to(device) + current_batch_size = image.shape[0] + if all(label == 'invalid' for label in label_type): + continue # NOTE: Skip all-invalid batches to avoid messing up the optimizer. + + gt_points = utils3d.torch.depth_to_points(gt_depth, intrinsics=gt_intrinsics) + gt_focal = 1 / (1 / gt_intrinsics[..., 0, 0] ** 2 + 1 / gt_intrinsics[..., 1, 1] ** 2) ** 0.5 + + with accelerator.accumulate(model): + # Forward + if i_step <= config.get('low_resolution_training_steps', 0): + num_tokens = config['model']['num_tokens_range'][0] + else: + num_tokens = accelerate.utils.broadcast_object_list([random.randint(*config['model']['num_tokens_range'])])[0] + with torch.autocast(device_type=accelerator.device.type, dtype=torch.float16, enabled=enable_mixed_precision): + output = model(image, num_tokens=num_tokens) + pred_points, pred_mask, pred_metric_scale = output['points'], output['mask'], output.get('metric_scale', None) + + # Compute loss (per instance) + loss_list, weight_list = [], [] + for i in range(current_batch_size): + gt_metric_scale = None + loss_dict, weight_dict, misc_dict = {}, {}, {} + misc_dict['monitoring'] = monitoring(pred_points[i]) + for k, v in config['loss'][label_type[i]].items(): + weight_dict[k] = v['weight'] + if v['function'] == 'affine_invariant_global_loss': + loss_dict[k], misc_dict[k], gt_metric_scale = affine_invariant_global_loss(pred_points[i], gt_points[i], gt_mask[i], **v['params']) + elif v['function'] == 'affine_invariant_local_loss': + loss_dict[k], misc_dict[k] = affine_invariant_local_loss(pred_points[i], gt_points[i], gt_mask[i], gt_focal[i], gt_metric_scale, **v['params']) + elif v['function'] == 'normal_loss': + loss_dict[k], misc_dict[k] = normal_loss(pred_points[i], gt_points[i], gt_mask[i]) + elif v['function'] == 'edge_loss': + loss_dict[k], misc_dict[k] = edge_loss(pred_points[i], gt_points[i], gt_mask[i]) + elif v['function'] == 'mask_bce_loss': + loss_dict[k], misc_dict[k] = mask_bce_loss(pred_mask[i], gt_mask_fin[i], gt_mask_inf[i]) + elif v['function'] == 'mask_l2_loss': + loss_dict[k], misc_dict[k] = mask_l2_loss(pred_mask[i], gt_mask_fin[i], gt_mask_inf[i]) + else: + raise ValueError(f'Undefined loss function: {v["function"]}') + weight_dict = {'.'.join(k): v for k, v in flatten_nested_dict(weight_dict).items()} + loss_dict = {'.'.join(k): v for k, v in flatten_nested_dict(loss_dict).items()} + loss_ = sum([weight_dict[k] * loss_dict[k] for k in loss_dict], start=torch.tensor(0.0, device=device)) + loss_list.append(loss_) + + if torch.isnan(loss_).item(): + pbar.write(f'NaN loss in process {accelerator.process_index}') + pbar.write(str(loss_dict)) + + misc_dict = {'.'.join(k): v for k, v in flatten_nested_dict(misc_dict).items()} + records.append({ + **{k: v.item() for k, v in loss_dict.items()}, + **misc_dict, + }) + + loss = sum(loss_list) / len(loss_list) + + # Backward & update + accelerator.backward(loss) + if accelerator.sync_gradients: + if not enable_mixed_precision and any(torch.isnan(p.grad).any() for p in model.parameters() if p.grad is not None): + if accelerator.is_main_process: + pbar.write(f'NaN gradients, skip update') + optimizer.zero_grad() + continue + accelerator.clip_grad_norm_(model.parameters(), 1.0) + + optimizer.step() + optimizer.zero_grad() + + i_accumulate += 1 + + lr_scheduler.step() + + # EMA update + if enable_ema and accelerator.is_main_process and accelerator.sync_gradients: + ema_model.update_parameters(model) + + # Log metrics + if i_step == initial_step or i_step % log_every == 0: + records = [key_average(records)] + records = accelerator.gather_for_metrics(records, use_gather_object=True) + if accelerator.is_main_process: + records = key_average(records) + if enable_mlflow: + try: + mlflow.log_metrics(records, step=i_step) + except Exception as e: + print(f'Error while logging metrics to mlflow: {e}') + records = [] + + # Save model weight checkpoint + if accelerator.is_main_process and (i_step % save_every == 0): + # NOTE: Writing checkpoint is done in a separate thread to avoid blocking the main process + pbar.write(f'Save checkpoint: {i_step:08d}') + Path(workspace, 'checkpoint').mkdir(parents=True, exist_ok=True) + + # Model checkpoint + with io.BytesIO() as f: + torch.save({ + 'model_config': config['model'], + 'model': accelerator.unwrap_model(model).state_dict(), + }, f) + checkpoint_bytes = f.getvalue() + save_checkpoint_executor.submit( + _write_bytes_retry_loop, Path(workspace, 'checkpoint', f'{i_step:08d}.pt'), checkpoint_bytes + ) + + # Optimizer checkpoint + with io.BytesIO() as f: + torch.save({ + 'model_config': config['model'], + 'step': i_step, + 'optimizer': optimizer.state_dict(), + 'lr_scheduler': lr_scheduler.state_dict(), + }, f) + checkpoint_bytes = f.getvalue() + save_checkpoint_executor.submit( + _write_bytes_retry_loop, Path(workspace, 'checkpoint', f'{i_step:08d}_optimizer.pt'), checkpoint_bytes + ) + + # EMA model checkpoint + if enable_ema: + with io.BytesIO() as f: + torch.save({ + 'model_config': config['model'], + 'model': ema_model.module.state_dict(), + }, f) + checkpoint_bytes = f.getvalue() + save_checkpoint_executor.submit( + _write_bytes_retry_loop, Path(workspace, 'checkpoint', f'{i_step:08d}_ema.pt'), checkpoint_bytes + ) + + # Latest checkpoint + with io.BytesIO() as f: + torch.save({ + 'model_config': config['model'], + 'step': i_step, + }, f) + checkpoint_bytes = f.getvalue() + save_checkpoint_executor.submit( + _write_bytes_retry_loop, Path(workspace, 'checkpoint', 'latest.pt'), checkpoint_bytes + ) + + # Visualize + if vis_every > 0 and accelerator.is_main_process and (i_step == initial_step or i_step % vis_every == 0): + unwrapped_model = accelerator.unwrap_model(model) + save_dir = Path(workspace).joinpath(f'vis/step_{i_step:08d}') + save_dir.mkdir(parents=True, exist_ok=True) + with torch.inference_mode(): + for i_batch, batch in enumerate(tqdm(batches_for_vis, desc=f'Visualize: {i_step:08d}', leave=False)): + image, gt_depth, gt_mask, gt_intrinsics = batch['image'], batch['depth'], batch['depth_mask'], batch['intrinsics'] + image, gt_depth, gt_mask, gt_intrinsics = image.to(device), gt_depth.to(device), gt_mask.to(device), gt_intrinsics.to(device) + + output = unwrapped_model.infer(image) + pred_points, pred_depth, pred_mask = output['points'].cpu().numpy(), output['depth'].cpu().numpy(), output['mask'].cpu().numpy() + image = image.cpu().numpy() + + for i_instance in range(image.shape[0]): + idx = i_batch * batch_size_forward + i_instance + image_i = (image[i_instance].transpose(1, 2, 0) * 255).astype(np.uint8) + pred_points_i = pred_points[i_instance] + pred_mask_i = pred_mask[i_instance] + pred_depth_i = pred_depth[i_instance] + save_dir.joinpath(f'{idx:04d}').mkdir(parents=True, exist_ok=True) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/image.jpg')), cv2.cvtColor(image_i, cv2.COLOR_RGB2BGR)) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/points.exr')), cv2.cvtColor(pred_points_i, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/mask.png')), pred_mask_i * 255) + cv2.imwrite(str(save_dir.joinpath(f'{idx:04d}/depth_vis.png')), cv2.cvtColor(colorize_depth(pred_depth_i, pred_mask_i), cv2.COLOR_RGB2BGR)) + + pbar.set_postfix({'loss': loss.item()}, refresh=False) + pbar.update(1) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/MoGe/moge/scripts/vis_data.py b/MoGe/moge/scripts/vis_data.py new file mode 100644 index 0000000000000000000000000000000000000000..bdb21766a67e4370578acbdf7bd17d1feb46b937 --- /dev/null +++ b/MoGe/moge/scripts/vis_data.py @@ -0,0 +1,84 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +import sys +from pathlib import Path +if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: + sys.path.insert(0, _package_root) + +import click + + +@click.command() +@click.argument('folder_or_path', type=click.Path(exists=True)) +@click.option('--output', '-o', 'output_folder', type=click.Path(), help='Path to output folder') +@click.option('--max_depth', '-m', type=float, default=float('inf'), help='max depth') +@click.option('--fov', type=float, default=None, help='field of view in degrees') +@click.option('--show', 'show', is_flag=True, help='show point cloud') +@click.option('--depth', 'depth_filename', type=str, default='depth.png', help='depth image file name') +@click.option('--ply', 'save_ply', is_flag=True, help='save point cloud as PLY file') +@click.option('--depth_vis', 'save_depth_vis', is_flag=True, help='save depth image') +@click.option('--inf', 'inf_mask', is_flag=True, help='use infinity mask') +@click.option('--version', 'version', type=str, default='v3', help='version of rgbd data') +def main( + folder_or_path: str, + output_folder: str, + max_depth: float, + fov: float, + depth_filename: str, + show: bool, + save_ply: bool, + save_depth_vis: bool, + inf_mask: bool, + version: str +): + # Lazy import + import cv2 + import numpy as np + import utils3d + from tqdm import tqdm + import trimesh + + from moge.utils.io import read_image, read_depth, read_meta + from moge.utils.vis import colorize_depth, colorize_normal + + filepaths = sorted(p.parent for p in Path(folder_or_path).rglob('meta.json')) + + for filepath in tqdm(filepaths): + image = read_image(Path(filepath, 'image.jpg')) + depth, unit = read_depth(Path(filepath, depth_filename)) + meta = read_meta(Path(filepath,'meta.json')) + depth_mask = np.isfinite(depth) + depth_mask_inf = (depth == np.inf) + intrinsics = np.array(meta['intrinsics']) + + extrinsics = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]], dtype=float) # OpenGL's identity camera + verts = utils3d.numpy.unproject_cv(utils3d.numpy.image_uv(*image.shape[:2]), depth, extrinsics=extrinsics, intrinsics=intrinsics) + + depth_mask_ply = depth_mask & (depth < depth[depth_mask].min() * max_depth) + point_cloud = trimesh.PointCloud(verts[depth_mask_ply], image[depth_mask_ply] / 255) + + if show: + point_cloud.show() + + if output_folder is None: + output_path = filepath + else: + output_path = Path(output_folder, filepath.name) + output_path.mkdir(exist_ok=True, parents=True) + + if inf_mask: + depth = np.where(depth_mask_inf, np.inf, depth) + depth_mask = depth_mask | depth_mask_inf + + if save_depth_vis: + p = output_path.joinpath('depth_vis.png') + cv2.imwrite(str(p), cv2.cvtColor(colorize_depth(depth, depth_mask), cv2.COLOR_RGB2BGR)) + print(f"{p}") + + if save_ply: + p = output_path.joinpath('pointcloud.ply') + point_cloud.export(p) + print(f"{p}") + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/MoGe/moge/test/__init__.py b/MoGe/moge/test/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MoGe/moge/test/baseline.py b/MoGe/moge/test/baseline.py new file mode 100644 index 0000000000000000000000000000000000000000..05980aaf96870304534fcec6532225e870351a66 --- /dev/null +++ b/MoGe/moge/test/baseline.py @@ -0,0 +1,43 @@ +from typing import * + +import click +import torch + + +class MGEBaselineInterface: + """ + Abstract class for model wrapper to uniformize the interface of loading and inference across different models. + """ + device: torch.device + + @click.command() + @staticmethod + def load(*args, **kwargs) -> "MGEBaselineInterface": + """ + Customized static method to create an instance of the model wrapper from command line arguments. Decorated by `click.command()` + """ + raise NotImplementedError(f"{type(self).__name__} has not implemented the load method.") + + def infer(self, image: torch.FloatTensor, intrinsics: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + """ + ### Parameters + `image`: [B, 3, H, W] or [3, H, W], RGB values in range [0, 1] + `intrinsics`: [B, 3, 3] or [3, 3], camera intrinsics. Optional. + + ### Returns + A dictionary containing: + - `points_*`. point map output in OpenCV identity camera space. + Supported suffixes: `metric`, `scale_invariant`, `affine_invariant`. + - `depth_*`. depth map output + Supported suffixes: `metric` (in meters), `scale_invariant`, `affine_invariant`. + - `disparity_affine_invariant`. affine disparity map output + """ + raise NotImplementedError(f"{type(self).__name__} has not implemented the infer method.") + + def infer_for_evaluation(self, image: torch.FloatTensor, intrinsics: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: + """ + If the model has a special evaluation mode, override this method to provide the evaluation mode inference. + + By default, this method simply calls `infer()`. + """ + return self.infer(image, intrinsics) \ No newline at end of file diff --git a/MoGe/moge/test/dataloader.py b/MoGe/moge/test/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..76679829afdf385938b604fa8bb5ef07b2560e7b --- /dev/null +++ b/MoGe/moge/test/dataloader.py @@ -0,0 +1,221 @@ +import os +from typing import * +from pathlib import Path +import math + +import numpy as np +import torch +from PIL import Image +import cv2 +import utils3d + +from ..utils import pipeline +from ..utils.geometry_numpy import focal_to_fov_numpy, mask_aware_nearest_resize_numpy, norm3d +from ..utils.io import * +from ..utils.tools import timeit + + +class EvalDataLoaderPipeline: + + def __init__( + self, + path: str, + width: int, + height: int, + split: int = '.index.txt', + drop_max_depth: float = 1000., + num_load_workers: int = 4, + num_process_workers: int = 8, + include_segmentation: bool = False, + include_normal: bool = False, + depth_to_normal: bool = False, + max_segments: int = 100, + min_seg_area: int = 1000, + depth_unit: str = None, + has_sharp_boundary = False, + subset: int = None, + ): + filenames = Path(path).joinpath(split).read_text(encoding='utf-8').splitlines() + filenames = filenames[::subset] + self.width = width + self.height = height + self.drop_max_depth = drop_max_depth + self.path = Path(path) + self.filenames = filenames + self.include_segmentation = include_segmentation + self.include_normal = include_normal + self.max_segments = max_segments + self.min_seg_area = min_seg_area + self.depth_to_normal = depth_to_normal + self.depth_unit = depth_unit + self.has_sharp_boundary = has_sharp_boundary + + self.rng = np.random.default_rng(seed=0) + + self.pipeline = pipeline.Sequential([ + self._generator, + pipeline.Parallel([self._load_instance] * num_load_workers), + pipeline.Parallel([self._process_instance] * num_process_workers), + pipeline.Buffer(4) + ]) + + def __len__(self): + return math.ceil(len(self.filenames)) + + def _generator(self): + for idx in range(len(self)): + yield idx + + def _load_instance(self, idx): + if idx >= len(self.filenames): + return None + + path = self.path.joinpath(self.filenames[idx]) + + instance = { + 'filename': self.filenames[idx], + 'width': self.width, + 'height': self.height, + } + instance['image'] = read_image(Path(path, 'image.jpg')) + + depth, _ = read_depth(Path(path, 'depth.png')) # ignore depth unit from depth file, use config instead + instance.update({ + 'depth': np.nan_to_num(depth, nan=1, posinf=1, neginf=1), + 'depth_mask': np.isfinite(depth), + 'depth_mask_inf': np.isinf(depth), + }) + + if self.include_segmentation: + segmentation_mask, segmentation_labels = read_segmentation(Path(path,'segmentation.png')) + instance.update({ + 'segmentation_mask': segmentation_mask, + 'segmentation_labels': segmentation_labels, + }) + + meta = read_meta(Path(path, 'meta.json')) + instance['intrinsics'] = np.array(meta['intrinsics'], dtype=np.float32) + + return instance + + def _process_instance(self, instance: dict): + if instance is None: + return None + + image, depth, depth_mask, intrinsics = instance['image'], instance['depth'], instance['depth_mask'], instance['intrinsics'] + segmentation_mask, segmentation_labels = instance.get('segmentation_mask', None), instance.get('segmentation_labels', None) + + raw_height, raw_width = image.shape[:2] + raw_horizontal, raw_vertical = abs(1.0 / intrinsics[0, 0]), abs(1.0 / intrinsics[1, 1]) + raw_pixel_w, raw_pixel_h = raw_horizontal / raw_width, raw_vertical / raw_height + tgt_width, tgt_height = instance['width'], instance['height'] + tgt_aspect = tgt_width / tgt_height + + # set expected target view field + tgt_horizontal = min(raw_horizontal, raw_vertical * tgt_aspect) + tgt_vertical = tgt_horizontal / tgt_aspect + + # set target view direction + cu, cv = 0.5, 0.5 + direction = utils3d.numpy.unproject_cv(np.array([[cu, cv]], dtype=np.float32), np.array([1.0], dtype=np.float32), intrinsics=intrinsics)[0] + R = utils3d.numpy.rotation_matrix_from_vectors(direction, np.array([0, 0, 1], dtype=np.float32)) + + # restrict target view field within the raw view + corners = np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=np.float32) + corners = np.concatenate([corners, np.ones((4, 1), dtype=np.float32)], axis=1) @ (np.linalg.inv(intrinsics).T @ R.T) # corners in viewport's camera plane + corners = corners[:, :2] / corners[:, 2:3] + + warp_horizontal, warp_vertical = abs(1.0 / intrinsics[0, 0]), abs(1.0 / intrinsics[1, 1]) + for i in range(4): + intersection, _ = utils3d.numpy.ray_intersection( + np.array([0., 0.]), np.array([[tgt_aspect, 1.0], [tgt_aspect, -1.0]]), + corners[i - 1], corners[i] - corners[i - 1], + ) + warp_horizontal, warp_vertical = min(warp_horizontal, 2 * np.abs(intersection[:, 0]).min()), min(warp_vertical, 2 * np.abs(intersection[:, 1]).min()) + tgt_horizontal, tgt_vertical = min(tgt_horizontal, warp_horizontal), min(tgt_vertical, warp_vertical) + + # get target view intrinsics + fx, fy = 1.0 / tgt_horizontal, 1.0 / tgt_vertical + tgt_intrinsics = utils3d.numpy.intrinsics_from_focal_center(fx, fy, 0.5, 0.5).astype(np.float32) + + # do homogeneous transformation with the rotation and intrinsics + # 4.1 The image and depth is resized first to approximately the same pixel size as the target image with PIL's antialiasing resampling + tgt_pixel_w, tgt_pixel_h = tgt_horizontal / tgt_width, tgt_vertical / tgt_height # (should be exactly the same for x and y axes) + rescaled_w, rescaled_h = int(raw_width * raw_pixel_w / tgt_pixel_w), int(raw_height * raw_pixel_h / tgt_pixel_h) + image = np.array(Image.fromarray(image).resize((rescaled_w, rescaled_h), Image.Resampling.LANCZOS)) + + depth, depth_mask = mask_aware_nearest_resize_numpy(depth, depth_mask, (rescaled_w, rescaled_h)) + distance = norm3d(utils3d.numpy.depth_to_points(depth, intrinsics=intrinsics)) + segmentation_mask = cv2.resize(segmentation_mask, (rescaled_w, rescaled_h), interpolation=cv2.INTER_NEAREST) if segmentation_mask is not None else None + + # 4.2 calculate homography warping + transform = intrinsics @ np.linalg.inv(R) @ np.linalg.inv(tgt_intrinsics) + uv_tgt = utils3d.numpy.image_uv(width=tgt_width, height=tgt_height) + pts = np.concatenate([uv_tgt, np.ones((tgt_height, tgt_width, 1), dtype=np.float32)], axis=-1) @ transform.T + uv_remap = pts[:, :, :2] / (pts[:, :, 2:3] + 1e-12) + pixel_remap = utils3d.numpy.uv_to_pixel(uv_remap, width=rescaled_w, height=rescaled_h).astype(np.float32) + + tgt_image = cv2.remap(image, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_LINEAR) + tgt_distance = cv2.remap(distance, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) + tgt_ray_length = utils3d.numpy.unproject_cv(uv_tgt, np.ones_like(uv_tgt[:, :, 0]), intrinsics=tgt_intrinsics) + tgt_ray_length = (tgt_ray_length[:, :, 0] ** 2 + tgt_ray_length[:, :, 1] ** 2 + tgt_ray_length[:, :, 2] ** 2) ** 0.5 + tgt_depth = tgt_distance / (tgt_ray_length + 1e-12) + tgt_depth_mask = cv2.remap(depth_mask.astype(np.uint8), pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) > 0 + tgt_segmentation_mask = cv2.remap(segmentation_mask, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) if segmentation_mask is not None else None + + # drop depth greater than drop_max_depth + max_depth = np.nanquantile(np.where(tgt_depth_mask, tgt_depth, np.nan), 0.01) * self.drop_max_depth + tgt_depth_mask &= tgt_depth <= max_depth + tgt_depth = np.nan_to_num(tgt_depth, nan=0.0) + + if self.depth_unit is not None: + tgt_depth *= self.depth_unit + + if not np.any(tgt_depth_mask): + # always make sure that mask is not empty, otherwise the loss calculation will crash + tgt_depth_mask = np.ones_like(tgt_depth_mask) + tgt_depth = np.ones_like(tgt_depth) + instance['label_type'] = 'invalid' + + tgt_pts = utils3d.numpy.unproject_cv(uv_tgt, tgt_depth, intrinsics=tgt_intrinsics) + + # Process segmentation labels + if self.include_segmentation and segmentation_mask is not None: + for k in ['undefined', 'unannotated', 'background', 'sky']: + if k in segmentation_labels: + del segmentation_labels[k] + seg_id2count = dict(zip(*np.unique(tgt_segmentation_mask, return_counts=True))) + sorted_labels = sorted(segmentation_labels.keys(), key=lambda x: seg_id2count.get(segmentation_labels[x], 0), reverse=True) + segmentation_labels = {k: segmentation_labels[k] for k in sorted_labels[:self.max_segments] if seg_id2count.get(segmentation_labels[k], 0) >= self.min_seg_area} + + instance.update({ + 'image': torch.from_numpy(tgt_image.astype(np.float32) / 255.0).permute(2, 0, 1), + 'depth': torch.from_numpy(tgt_depth).float(), + 'depth_mask': torch.from_numpy(tgt_depth_mask).bool(), + 'intrinsics': torch.from_numpy(tgt_intrinsics).float(), + 'points': torch.from_numpy(tgt_pts).float(), + 'segmentation_mask': torch.from_numpy(tgt_segmentation_mask).long() if tgt_segmentation_mask is not None else None, + 'segmentation_labels': segmentation_labels, + 'is_metric': self.depth_unit is not None, + 'has_sharp_boundary': self.has_sharp_boundary, + }) + + instance = {k: v for k, v in instance.items() if v is not None} + + return instance + + def start(self): + self.pipeline.start() + + def stop(self): + self.pipeline.stop() + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.stop() + + def get(self): + return self.pipeline.get() \ No newline at end of file diff --git a/MoGe/moge/test/metrics.py b/MoGe/moge/test/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..904064f2a30d05dca3a53db7ecc076a0c2aaa0ad --- /dev/null +++ b/MoGe/moge/test/metrics.py @@ -0,0 +1,343 @@ +from typing import * +from numbers import Number + +import torch +import torch.nn.functional as F +import numpy as np +import utils3d + +from ..utils.geometry_torch import ( + weighted_mean, + mask_aware_nearest_resize, + intrinsics_to_fov +) +from ..utils.alignment import ( + align_points_scale_z_shift, + align_points_scale_xyz_shift, + align_points_xyz_shift, + align_affine_lstsq, + align_depth_scale, + align_depth_affine, + align_points_scale, +) +from ..utils.tools import key_average, timeit + + +def rel_depth(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6): + rel = (torch.abs(pred - gt) / (gt + eps)).mean() + return rel.item() + + +def delta1_depth(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6): + delta1 = (torch.maximum(gt / pred, pred / gt) < 1.25).float().mean() + return delta1.item() + + +def rel_point(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6): + dist_gt = torch.norm(gt, dim=-1) + dist_err = torch.norm(pred - gt, dim=-1) + rel = (dist_err / (dist_gt + eps)).mean() + return rel.item() + + +def delta1_point(pred: torch.Tensor, gt: torch.Tensor, eps: float = 1e-6): + dist_pred = torch.norm(pred, dim=-1) + dist_gt = torch.norm(gt, dim=-1) + dist_err = torch.norm(pred - gt, dim=-1) + + delta1 = (dist_err < 0.25 * torch.minimum(dist_gt, dist_pred)).float().mean() + return delta1.item() + + +def rel_point_local(pred: torch.Tensor, gt: torch.Tensor, diameter: torch.Tensor): + dist_err = torch.norm(pred - gt, dim=-1) + rel = (dist_err / diameter).mean() + return rel.item() + + +def delta1_point_local(pred: torch.Tensor, gt: torch.Tensor, diameter: torch.Tensor): + dist_err = torch.norm(pred - gt, dim=-1) + delta1 = (dist_err < 0.25 * diameter).float().mean() + return delta1.item() + + +def boundary_f1(pred: torch.Tensor, gt: torch.Tensor, mask: torch.Tensor, radius: int = 1): + neighbor_x, neight_y = torch.meshgrid( + torch.linspace(-radius, radius, 2 * radius + 1, device=pred.device), + torch.linspace(-radius, radius, 2 * radius + 1, device=pred.device), + indexing='xy' + ) + neighbor_mask = (neighbor_x ** 2 + neight_y ** 2) <= radius ** 2 + 1e-5 + + pred_window = utils3d.torch.sliding_window_2d(pred, window_size=2 * radius + 1, stride=1, dim=(-2, -1)) # [H, W, 2*R+1, 2*R+1] + gt_window = utils3d.torch.sliding_window_2d(gt, window_size=2 * radius + 1, stride=1, dim=(-2, -1)) # [H, W, 2*R+1, 2*R+1] + mask_window = neighbor_mask & utils3d.torch.sliding_window_2d(mask, window_size=2 * radius + 1, stride=1, dim=(-2, -1)) # [H, W, 2*R+1, 2*R+1] + + pred_rel = pred_window / pred[radius:-radius, radius:-radius, None, None] + gt_rel = gt_window / gt[radius:-radius, radius:-radius, None, None] + valid = mask[radius:-radius, radius:-radius, None, None] & mask_window + + f1_list = [] + w_list = t_list = torch.linspace(0.05, 0.25, 10).tolist() + + for t in t_list: + pred_label = pred_rel > 1 + t + gt_label = gt_rel > 1 + t + TP = (pred_label & gt_label & valid).float().sum() + precision = TP / (gt_label & valid).float().sum().clamp_min(1e-12) + recall = TP / (pred_label & valid).float().sum().clamp_min(1e-12) + f1 = 2 * precision * recall / (precision + recall).clamp_min(1e-12) + f1_list.append(f1.item()) + + f1_avg = sum(w * f1 for w, f1 in zip(w_list, f1_list)) / sum(w_list) + return f1_avg + + +def compute_metrics( + pred: Dict[str, torch.Tensor], + gt: Dict[str, torch.Tensor], + vis: bool = False +) -> Tuple[Dict[str, Dict[str, Number]], Dict[str, torch.Tensor]]: + """ + A unified function to compute metrics for different types of predictions and ground truths. + + #### Supported keys in pred: + - `disparity_affine_invariant`: disparity map predicted by a depth estimator with scale and shift invariant. + - `depth_scale_invariant`: depth map predicted by a depth estimator with scale invariant. + - `depth_affine_invariant`: depth map predicted by a depth estimator with scale and shift invariant. + - `depth_metric`: depth map predicted by a depth estimator with no scale or shift. + - `points_scale_invariant`: point map predicted by a point estimator with scale invariant. + - `points_affine_invariant`: point map predicted by a point estimator with scale and xyz shift invariant. + - `points_metric`: point map predicted by a point estimator with no scale or shift. + - `intrinsics`: normalized camera intrinsics matrix. + + #### Required keys in gt: + - `depth`: depth map ground truth (in metric units if `depth_metric` is used) + - `points`: point map ground truth in camera coordinates. + - `mask`: mask indicating valid pixels in the ground truth. + - `intrinsics`: normalized ground-truth camera intrinsics matrix. + - `is_metric`: whether the depth is in metric units. + """ + metrics = {} + misc = {} + + mask = gt['depth_mask'] + gt_depth = gt['depth'] + gt_points = gt['points'] + + height, width = mask.shape[-2:] + _, lr_mask, lr_index = mask_aware_nearest_resize(None, mask, (64, 64), return_index=True) + + only_depth = not any('point' in k for k in pred) + pred_depth_aligned, pred_points_aligned = None, None + + # Metric depth + if 'depth_metric' in pred and gt['is_metric']: + pred_depth, gt_depth = pred['depth_metric'], gt['depth'] + metrics['depth_metric'] = { + 'rel': rel_depth(pred_depth[mask], gt_depth[mask]), + 'delta1': delta1_depth(pred_depth[mask], gt_depth[mask]) + } + + if pred_depth_aligned is None: + pred_depth_aligned = pred_depth + + # Scale-invariant depth + if 'depth_scale_invariant' in pred: + pred_depth_scale_invariant = pred['depth_scale_invariant'] + elif 'depth_metric' in pred: + pred_depth_scale_invariant = pred['depth_metric'] + else: + pred_depth_scale_invariant = None + + if pred_depth_scale_invariant is not None: + pred_depth = pred_depth_scale_invariant + + pred_depth_lr_masked, gt_depth_lr_masked = pred_depth[lr_index][lr_mask], gt_depth[lr_index][lr_mask] + scale = align_depth_scale(pred_depth_lr_masked, gt_depth_lr_masked, 1 / gt_depth_lr_masked) + pred_depth = pred_depth * scale + + metrics['depth_scale_invariant'] = { + 'rel': rel_depth(pred_depth[mask], gt_depth[mask]), + 'delta1': delta1_depth(pred_depth[mask], gt_depth[mask]) + } + + if pred_depth_aligned is None: + pred_depth_aligned = pred_depth + + # Affine-invariant depth + if 'depth_affine_invariant' in pred: + pred_depth_affine_invariant = pred['depth_affine_invariant'] + elif 'depth_scale_invariant' in pred: + pred_depth_affine_invariant = pred['depth_scale_invariant'] + elif 'depth_metric' in pred: + pred_depth_affine_invariant = pred['depth_metric'] + else: + pred_depth_affine_invariant = None + + if pred_depth_affine_invariant is not None: + pred_depth = pred_depth_affine_invariant + + pred_depth_lr_masked, gt_depth_lr_masked = pred_depth[lr_index][lr_mask], gt_depth[lr_index][lr_mask] + scale, shift = align_depth_affine(pred_depth_lr_masked, gt_depth_lr_masked, 1 / gt_depth_lr_masked) + pred_depth = pred_depth * scale + shift + + metrics['depth_affine_invariant'] = { + 'rel': rel_depth(pred_depth[mask], gt_depth[mask]), + 'delta1': delta1_depth(pred_depth[mask], gt_depth[mask]) + } + + if pred_depth_aligned is None: + pred_depth_aligned = pred_depth + + # Affine-invariant disparity + if 'disparity_affine_invariant' in pred: + pred_disparity_affine_invariant = pred['disparity_affine_invariant'] + elif 'depth_scale_invariant' in pred: + pred_disparity_affine_invariant = 1 / pred['depth_scale_invariant'] + elif 'depth_metric' in pred: + pred_disparity_affine_invariant = 1 / pred['depth_metric'] + else: + pred_disparity_affine_invariant = None + + if pred_disparity_affine_invariant is not None: + pred_disp = pred_disparity_affine_invariant + + scale, shift = align_affine_lstsq(pred_disp[mask], 1 / gt_depth[mask]) + pred_disp = pred_disp * scale + shift + + # NOTE: The alignment is done on the disparity map could introduce extreme outliers at disparities close to 0. + # Therefore we clamp the disparities by minimum ground truth disparity. + pred_depth = 1 / pred_disp.clamp_min(1 / gt_depth[mask].max().item()) + + metrics['disparity_affine_invariant'] = { + 'rel': rel_depth(pred_depth[mask], gt_depth[mask]), + 'delta1': delta1_depth(pred_depth[mask], gt_depth[mask]) + } + + if pred_depth_aligned is None: + pred_depth_aligned = 1 / pred_disp.clamp_min(1e-6) + + # Metric points + if 'points_metric' in pred and gt['is_metric']: + pred_points = pred['points_metric'] + + pred_points_lr_masked, gt_points_lr_masked = pred_points[lr_index][lr_mask], gt_points[lr_index][lr_mask] + shift = align_points_xyz_shift(pred_points_lr_masked, gt_points_lr_masked, 1 / gt_points_lr_masked.norm(dim=-1)) + pred_points = pred_points + shift + + metrics['points_metric'] = { + 'rel': rel_point(pred_points[mask], gt_points[mask]), + 'delta1': delta1_point(pred_points[mask], gt_points[mask]) + } + + if pred_points_aligned is None: + pred_points_aligned = pred['points_metric'] + + # Scale-invariant points (in camera space) + if 'points_scale_invariant' in pred: + pred_points_scale_invariant = pred['points_scale_invariant'] + elif 'points_metric' in pred: + pred_points_scale_invariant = pred['points_metric'] + else: + pred_points_scale_invariant = None + + if pred_points_scale_invariant is not None: + pred_points = pred_points_scale_invariant + + pred_points_lr_masked, gt_points_lr_masked = pred_points_scale_invariant[lr_index][lr_mask], gt_points[lr_index][lr_mask] + scale = align_points_scale(pred_points_lr_masked, gt_points_lr_masked, 1 / gt_points_lr_masked.norm(dim=-1)) + pred_points = pred_points * scale + + metrics['points_scale_invariant'] = { + 'rel': rel_point(pred_points[mask], gt_points[mask]), + 'delta1': delta1_point(pred_points[mask], gt_points[mask]) + } + + if vis and pred_points_aligned is None: + pred_points_aligned = pred['points_scale_invariant'] * scale + + # Affine-invariant points + if 'points_affine_invariant' in pred: + pred_points_affine_invariant = pred['points_affine_invariant'] + elif 'points_scale_invariant' in pred: + pred_points_affine_invariant = pred['points_scale_invariant'] + elif 'points_metric' in pred: + pred_points_affine_invariant = pred['points_metric'] + else: + pred_points_affine_invariant = None + + if pred_points_affine_invariant is not None: + pred_points = pred_points_affine_invariant + + pred_points_lr_masked, gt_points_lr_masked = pred_points[lr_index][lr_mask], gt_points[lr_index][lr_mask] + scale, shift = align_points_scale_xyz_shift(pred_points_lr_masked, gt_points_lr_masked, 1 / gt_points_lr_masked.norm(dim=-1)) + pred_points = pred_points * scale + shift + + metrics['points_affine_invariant'] = { + 'rel': rel_point(pred_points[mask], gt_points[mask]), + 'delta1': delta1_point(pred_points[mask], gt_points[mask]) + } + + if vis and pred_points_aligned is None: + pred_points_aligned = pred['points_affine_invariant'] * scale + shift + + # Local points + if 'segmentation_mask' in gt and 'points' in gt and any('points' in k for k in pred.keys()): + pred_points = next(pred[k] for k in pred.keys() if 'points' in k) + gt_points = gt['points'] + segmentation_mask = gt['segmentation_mask'] + segmentation_labels = gt['segmentation_labels'] + segmentation_mask_lr = segmentation_mask[lr_index] + local_points_metrics = [] + for _, seg_id in segmentation_labels.items(): + valid_mask = (segmentation_mask == seg_id) & mask + + pred_points_masked = pred_points[valid_mask] + gt_points_masked = gt_points[valid_mask] + + valid_mask_lr = (segmentation_mask_lr == seg_id) & lr_mask + if valid_mask_lr.sum().item() < 10: + continue + pred_points_masked_lr = pred_points[lr_index][valid_mask_lr] + gt_points_masked_lr = gt_points[lr_index][valid_mask_lr] + diameter = (gt_points_masked.max(dim=0).values - gt_points_masked.min(dim=0).values).max() + scale, shift = align_points_scale_xyz_shift(pred_points_masked_lr, gt_points_masked_lr, 1 / diameter.expand(gt_points_masked_lr.shape[0])) + pred_points_masked = pred_points_masked * scale + shift + + local_points_metrics.append({ + 'rel': rel_point_local(pred_points_masked, gt_points_masked, diameter), + 'delta1': delta1_point_local(pred_points_masked, gt_points_masked, diameter), + }) + + metrics['local_points'] = key_average(local_points_metrics) + + # FOV. NOTE: If there is no random augmentation applied to the input images, all GT FOV are generallly the same. + # Fair evaluation of FOV requires random augmentation. + if 'intrinsics' in pred and 'intrinsics' in gt: + pred_intrinsics = pred['intrinsics'] + gt_intrinsics = gt['intrinsics'] + pred_fov_x, pred_fov_y = intrinsics_to_fov(pred_intrinsics) + gt_fov_x, gt_fov_y = intrinsics_to_fov(gt_intrinsics) + metrics['fov_x'] = { + 'mae': torch.rad2deg(pred_fov_x - gt_fov_x).abs().mean().item(), + 'deviation': torch.rad2deg(pred_fov_x - gt_fov_x).item(), + } + + # Boundary F1 + if pred_depth_aligned is not None and gt['has_sharp_boundary']: + metrics['boundary'] = { + 'radius1_f1': boundary_f1(pred_depth_aligned, gt_depth, mask, radius=1), + 'radius2_f1': boundary_f1(pred_depth_aligned, gt_depth, mask, radius=2), + 'radius3_f1': boundary_f1(pred_depth_aligned, gt_depth, mask, radius=3), + } + + if vis: + if pred_points_aligned is not None: + misc['pred_points'] = pred_points_aligned + if only_depth: + misc['pred_points'] = utils3d.torch.depth_to_points(pred_depth_aligned, intrinsics=gt['intrinsics']) + if pred_depth_aligned is not None: + misc['pred_depth'] = pred_depth_aligned + + return metrics, misc \ No newline at end of file diff --git a/MoGe/moge/train/__init__.py b/MoGe/moge/train/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MoGe/moge/train/dataloader.py b/MoGe/moge/train/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..b82294f310408485fa4c2694274e57e539ffef0d --- /dev/null +++ b/MoGe/moge/train/dataloader.py @@ -0,0 +1,339 @@ +import os +from pathlib import Path +import json +import time +import random +from typing import * +import traceback +import itertools +from numbers import Number +import io + +import numpy as np +import cv2 +from PIL import Image +import torch +import torchvision.transforms.v2.functional as TF +import utils3d +from tqdm import tqdm + +from ..utils import pipeline +from ..utils.io import * +from ..utils.geometry_numpy import mask_aware_nearest_resize_numpy, harmonic_mean_numpy, norm3d, depth_occlusion_edge_numpy, depth_of_field + + +class TrainDataLoaderPipeline: + def __init__(self, config: dict, batch_size: int, num_load_workers: int = 4, num_process_workers: int = 8, buffer_size: int = 8): + self.config = config + + self.batch_size = batch_size + self.clamp_max_depth = config['clamp_max_depth'] + self.fov_range_absolute = config.get('fov_range_absolute', 0.0) + self.fov_range_relative = config.get('fov_range_relative', 0.0) + self.center_augmentation = config.get('center_augmentation', 0.0) + self.image_augmentation = config.get('image_augmentation', []) + self.depth_interpolation = config.get('depth_interpolation', 'bilinear') + + if 'image_sizes' in config: + self.image_size_strategy = 'fixed' + self.image_sizes = config['image_sizes'] + elif 'aspect_ratio_range' in config and 'area_range' in config: + self.image_size_strategy = 'aspect_area' + self.aspect_ratio_range = config['aspect_ratio_range'] + self.area_range = config['area_range'] + else: + raise ValueError('Invalid image size configuration') + + # Load datasets + self.datasets = {} + for dataset in tqdm(config['datasets'], desc='Loading datasets'): + name = dataset['name'] + content = Path(dataset['path'], dataset.get('index', '.index.txt')).joinpath().read_text() + filenames = content.splitlines() + self.datasets[name] = { + **dataset, + 'path': dataset['path'], + 'filenames': filenames, + } + self.dataset_names = [dataset['name'] for dataset in config['datasets']] + self.dataset_weights = [dataset['weight'] for dataset in config['datasets']] + + # Build pipeline + self.pipeline = pipeline.Sequential([ + self._sample_batch, + pipeline.Unbatch(), + pipeline.Parallel([self._load_instance] * num_load_workers), + pipeline.Parallel([self._process_instance] * num_process_workers), + pipeline.Batch(self.batch_size), + self._collate_batch, + pipeline.Buffer(buffer_size), + ]) + + self.invalid_instance = { + 'intrinsics': np.array([[1.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.0, 0.0, 1.0]], dtype=np.float32), + 'image': np.zeros((256, 256, 3), dtype=np.uint8), + 'depth': np.ones((256, 256), dtype=np.float32), + 'depth_mask': np.ones((256, 256), dtype=bool), + 'depth_mask_inf': np.zeros((256, 256), dtype=bool), + 'label_type': 'invalid', + } + + def _sample_batch(self): + batch_id = 0 + last_area = None + while True: + # Depending on the sample strategy, choose a dataset and a filename + batch_id += 1 + batch = [] + + # Sample instances + for _ in range(self.batch_size): + dataset_name = random.choices(self.dataset_names, weights=self.dataset_weights)[0] + filename = random.choice(self.datasets[dataset_name]['filenames']) + + path = Path(self.datasets[dataset_name]['path'], filename) + + instance = { + 'batch_id': batch_id, + 'seed': random.randint(0, 2 ** 32 - 1), + 'dataset': dataset_name, + 'filename': filename, + 'path': path, + 'label_type': self.datasets[dataset_name]['label_type'], + } + batch.append(instance) + + # Decide the image size for this batch + if self.image_size_strategy == 'fixed': + width, height = random.choice(self.config['image_sizes']) + elif self.image_size_strategy == 'aspect_area': + area = random.uniform(*self.area_range) + aspect_ratio_ranges = [self.datasets[instance['dataset']].get('aspect_ratio_range', self.aspect_ratio_range) for instance in batch] + aspect_ratio_range = (min(r[0] for r in aspect_ratio_ranges), max(r[1] for r in aspect_ratio_ranges)) + aspect_ratio = random.uniform(*aspect_ratio_range) + width, height = int((area * aspect_ratio) ** 0.5), int((area / aspect_ratio) ** 0.5) + else: + raise ValueError('Invalid image size strategy') + + for instance in batch: + instance['width'], instance['height'] = width, height + + yield batch + + def _load_instance(self, instance: dict): + try: + image = read_image(Path(instance['path'], 'image.jpg')) + depth, _ = read_depth(Path(instance['path'], self.datasets[instance['dataset']].get('depth', 'depth.png'))) + + meta = read_meta(Path(instance['path'], 'meta.json')) + intrinsics = np.array(meta['intrinsics'], dtype=np.float32) + depth_mask = np.isfinite(depth) + depth_mask_inf = np.isinf(depth) + depth = np.nan_to_num(depth, nan=1, posinf=1, neginf=1) + data = { + 'image': image, + 'depth': depth, + 'depth_mask': depth_mask, + 'depth_mask_inf': depth_mask_inf, + 'intrinsics': intrinsics + } + instance.update({ + **data, + }) + except Exception as e: + print(f"Failed to load instance {instance['dataset']}/{instance['filename']} because of exception:", e) + instance.update(self.invalid_instance) + return instance + + def _process_instance(self, instance: Dict[str, Union[np.ndarray, str, float, bool]]): + image, depth, depth_mask, depth_mask_inf, intrinsics, label_type = instance['image'], instance['depth'], instance['depth_mask'], instance['depth_mask_inf'], instance['intrinsics'], instance['label_type'] + depth_unit = self.datasets[instance['dataset']].get('depth_unit', None) + + raw_height, raw_width = image.shape[:2] + raw_horizontal, raw_vertical = abs(1.0 / intrinsics[0, 0]), abs(1.0 / intrinsics[1, 1]) + raw_fov_x, raw_fov_y = utils3d.numpy.intrinsics_to_fov(intrinsics) + raw_pixel_w, raw_pixel_h = raw_horizontal / raw_width, raw_vertical / raw_height + tgt_width, tgt_height = instance['width'], instance['height'] + tgt_aspect = tgt_width / tgt_height + + rng = np.random.default_rng(instance['seed']) + + # 1. set target fov + center_augmentation = self.datasets[instance['dataset']].get('center_augmentation', self.center_augmentation) + fov_range_absolute_min, fov_range_absolute_max = self.datasets[instance['dataset']].get('fov_range_absolute', self.fov_range_absolute) + fov_range_relative_min, fov_range_relative_max = self.datasets[instance['dataset']].get('fov_range_relative', self.fov_range_relative) + tgt_fov_x_min = min(fov_range_relative_min * raw_fov_x, fov_range_relative_min * utils3d.focal_to_fov(utils3d.fov_to_focal(raw_fov_y) / tgt_aspect)) + tgt_fov_x_max = min(fov_range_relative_max * raw_fov_x, fov_range_relative_max * utils3d.focal_to_fov(utils3d.fov_to_focal(raw_fov_y) / tgt_aspect)) + tgt_fov_x_min, tgt_fov_max = max(np.deg2rad(fov_range_absolute_min), tgt_fov_x_min), min(np.deg2rad(fov_range_absolute_max), tgt_fov_x_max) + tgt_fov_x = rng.uniform(min(tgt_fov_x_min, tgt_fov_x_max), tgt_fov_x_max) + tgt_fov_y = utils3d.focal_to_fov(utils3d.numpy.fov_to_focal(tgt_fov_x) * tgt_aspect) + + # 2. set target image center (principal point) and the corresponding z-direction in raw camera space + center_dtheta = center_augmentation * rng.uniform(-0.5, 0.5) * (raw_fov_x - tgt_fov_x) + center_dphi = center_augmentation * rng.uniform(-0.5, 0.5) * (raw_fov_y - tgt_fov_y) + cu, cv = 0.5 + 0.5 * np.tan(center_dtheta) / np.tan(raw_fov_x / 2), 0.5 + 0.5 * np.tan(center_dphi) / np.tan(raw_fov_y / 2) + direction = utils3d.unproject_cv(np.array([[cu, cv]], dtype=np.float32), np.array([1.0], dtype=np.float32), intrinsics=intrinsics)[0] + + # 3. obtain the rotation matrix for homography warping + R = utils3d.rotation_matrix_from_vectors(direction, np.array([0, 0, 1], dtype=np.float32)) + + # 4. shrink the target view to fit into the warped image + corners = np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=np.float32) + corners = np.concatenate([corners, np.ones((4, 1), dtype=np.float32)], axis=1) @ (np.linalg.inv(intrinsics).T @ R.T) # corners in viewport's camera plane + corners = corners[:, :2] / corners[:, 2:3] + tgt_horizontal, tgt_vertical = np.tan(tgt_fov_x / 2) * 2, np.tan(tgt_fov_y / 2) * 2 + warp_horizontal, warp_vertical = float('inf'), float('inf') + for i in range(4): + intersection, _ = utils3d.numpy.ray_intersection( + np.array([0., 0.]), np.array([[tgt_aspect, 1.0], [tgt_aspect, -1.0]]), + corners[i - 1], corners[i] - corners[i - 1], + ) + warp_horizontal, warp_vertical = min(warp_horizontal, 2 * np.abs(intersection[:, 0]).min()), min(warp_vertical, 2 * np.abs(intersection[:, 1]).min()) + tgt_horizontal, tgt_vertical = min(tgt_horizontal, warp_horizontal), min(tgt_vertical, warp_vertical) + + # 5. obtain the target intrinsics + fx, fy = 1 / tgt_horizontal, 1 / tgt_vertical + tgt_intrinsics = utils3d.numpy.intrinsics_from_focal_center(fx, fy, 0.5, 0.5).astype(np.float32) + + # 6. do homogeneous transformation + # 6.1 The image and depth are resized first to approximately the same pixel size as the target image with PIL's antialiasing resampling + tgt_pixel_w, tgt_pixel_h = tgt_horizontal / tgt_width, tgt_vertical / tgt_height # (should be exactly the same for x and y axes) + rescaled_w, rescaled_h = int(raw_width * raw_pixel_w / tgt_pixel_w), int(raw_height * raw_pixel_h / tgt_pixel_h) + image = np.array(Image.fromarray(image).resize((rescaled_w, rescaled_h), Image.Resampling.LANCZOS)) + + fg_edge_mask, bg_edge_mask = depth_occlusion_edge_numpy(depth, mask=depth_mask, kernel_size=5, tol=0.01) + edge_mask = fg_edge_mask | bg_edge_mask + _, depth_mask_nearest, resize_index = mask_aware_nearest_resize_numpy(None, depth_mask, (rescaled_w, rescaled_h), return_index=True) + depth_nearest = depth[resize_index] + distance_nearest = norm3d(utils3d.numpy.depth_to_points(depth_nearest, intrinsics=intrinsics)) + edge_mask = edge_mask[resize_index] + + if self.depth_interpolation == 'bilinear': + depth_mask_bilinear = cv2.resize(depth_mask.astype(np.float32), (rescaled_w, rescaled_h), interpolation=cv2.INTER_LINEAR) + depth_bilinear = 1 / cv2.resize(1 / depth, (rescaled_w, rescaled_h), interpolation=cv2.INTER_LINEAR) + distance_bilinear = norm3d(utils3d.numpy.depth_to_points(depth_bilinear, intrinsics=intrinsics)) + + depth_mask_inf = cv2.resize(depth_mask_inf.astype(np.uint8), (rescaled_w, rescaled_h), interpolation=cv2.INTER_NEAREST) > 0 + + # 6.2 calculate homography warping + transform = intrinsics @ np.linalg.inv(R) @ np.linalg.inv(tgt_intrinsics) + uv_tgt = utils3d.numpy.image_uv(width=tgt_width, height=tgt_height) + pts = np.concatenate([uv_tgt, np.ones((tgt_height, tgt_width, 1), dtype=np.float32)], axis=-1) @ transform.T + uv_remap = pts[:, :, :2] / (pts[:, :, 2:3] + 1e-12) + pixel_remap = utils3d.numpy.uv_to_pixel(uv_remap, width=rescaled_w, height=rescaled_h).astype(np.float32) + + tgt_image = cv2.remap(image, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_LANCZOS4) + tgt_ray_length = norm3d(utils3d.numpy.unproject_cv(uv_tgt, np.ones_like(uv_tgt[:, :, 0]), intrinsics=tgt_intrinsics)) + tgt_depth_mask_nearest = cv2.remap(depth_mask_nearest.astype(np.uint8), pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) > 0 + tgt_depth_nearest = cv2.remap(distance_nearest, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) / tgt_ray_length + tgt_edge_mask = cv2.remap(edge_mask.astype(np.uint8), pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) > 0 + if self.depth_interpolation == 'bilinear': + tgt_depth_mask_bilinear = cv2.remap(depth_mask_bilinear, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_LINEAR) + tgt_depth_bilinear = cv2.remap(distance_bilinear, pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_LINEAR) / tgt_ray_length + tgt_depth = np.where((tgt_depth_mask_bilinear == 1) & ~tgt_edge_mask, tgt_depth_bilinear, tgt_depth_nearest) + else: + tgt_depth = tgt_depth_nearest + tgt_depth_mask = tgt_depth_mask_nearest + + tgt_depth_mask_inf = cv2.remap(depth_mask_inf.astype(np.uint8), pixel_remap[:, :, 0], pixel_remap[:, :, 1], cv2.INTER_NEAREST) > 0 + + # always make sure that mask is not empty + if tgt_depth_mask.sum() / tgt_depth_mask.size < 0.001: + tgt_depth_mask = np.ones_like(tgt_depth_mask) + tgt_depth = np.ones_like(tgt_depth) + instance['label_type'] = 'invalid' + + # Flip augmentation + if rng.choice([True, False]): + tgt_image = np.flip(tgt_image, axis=1).copy() + tgt_depth = np.flip(tgt_depth, axis=1).copy() + tgt_depth_mask = np.flip(tgt_depth_mask, axis=1).copy() + tgt_depth_mask_inf = np.flip(tgt_depth_mask_inf, axis=1).copy() + + # Color augmentation + image_augmentation = self.datasets[instance['dataset']].get('image_augmentation', self.image_augmentation) + if 'jittering' in image_augmentation: + tgt_image = torch.from_numpy(tgt_image).permute(2, 0, 1) + tgt_image = TF.adjust_brightness(tgt_image, rng.uniform(0.7, 1.3)) + tgt_image = TF.adjust_contrast(tgt_image, rng.uniform(0.7, 1.3)) + tgt_image = TF.adjust_saturation(tgt_image, rng.uniform(0.7, 1.3)) + tgt_image = TF.adjust_hue(tgt_image, rng.uniform(-0.1, 0.1)) + tgt_image = TF.adjust_gamma(tgt_image, rng.uniform(0.7, 1.3)) + tgt_image = tgt_image.permute(1, 2, 0).numpy() + if 'dof' in image_augmentation: + if rng.uniform() < 0.5: + dof_strength = rng.integers(12) + tgt_disp = np.where(tgt_depth_mask_inf, 0, 1 / tgt_depth) + disp_min, disp_max = tgt_disp[tgt_depth_mask].min(), tgt_disp[tgt_depth_mask].max() + tgt_disp = cv2.inpaint(tgt_disp, (~tgt_depth_mask & ~tgt_depth_mask_inf).astype(np.uint8), 3, cv2.INPAINT_TELEA).clip(disp_min, disp_max) + dof_focus = rng.uniform(disp_min, disp_max) + tgt_image = depth_of_field(tgt_image, tgt_disp, dof_focus, dof_strength) + if 'shot_noise' in image_augmentation: + if rng.uniform() < 0.5: + k = np.exp(rng.uniform(np.log(100), np.log(10000))) / 255 + tgt_image = (rng.poisson(tgt_image * k) / k).clip(0, 255).astype(np.uint8) + if 'jpeg_loss' in image_augmentation: + if rng.uniform() < 0.5: + tgt_image = cv2.imdecode(cv2.imencode('.jpg', tgt_image, [cv2.IMWRITE_JPEG_QUALITY, rng.integers(20, 100)])[1], cv2.IMREAD_COLOR) + if 'blurring' in image_augmentation: + if rng.uniform() < 0.5: + ratio = rng.uniform(0.25, 1) + tgt_image = cv2.resize(cv2.resize(tgt_image, (int(tgt_width * ratio), int(tgt_height * ratio)), interpolation=cv2.INTER_AREA), (tgt_width, tgt_height), interpolation=rng.choice([cv2.INTER_LINEAR_EXACT, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4])) + + # convert depth to metric if necessary + if depth_unit is not None: + tgt_depth *= depth_unit + instance['is_metric'] = True + else: + instance['is_metric'] = False + + # clamp depth maximum values + max_depth = np.nanquantile(np.where(tgt_depth_mask, tgt_depth, np.nan), 0.01) * self.clamp_max_depth + tgt_depth = np.clip(tgt_depth, 0, max_depth) + tgt_depth = np.nan_to_num(tgt_depth, nan=1.0) + + if self.datasets[instance['dataset']].get('finite_depth_mask', None) == "only_known": + tgt_depth_mask_fin = tgt_depth_mask + else: + tgt_depth_mask_fin = ~tgt_depth_mask_inf + + instance.update({ + 'image': torch.from_numpy(tgt_image.astype(np.float32) / 255.0).permute(2, 0, 1), + 'depth': torch.from_numpy(tgt_depth).float(), + 'depth_mask': torch.from_numpy(tgt_depth_mask).bool(), + 'depth_mask_fin': torch.from_numpy(tgt_depth_mask_fin).bool(), + 'depth_mask_inf': torch.from_numpy(tgt_depth_mask_inf).bool(), + 'intrinsics': torch.from_numpy(tgt_intrinsics).float(), + }) + + return instance + + def _collate_batch(self, instances: List[Dict[str, Any]]): + batch = {k: torch.stack([instance[k] for instance in instances], dim=0) for k in ['image', 'depth', 'depth_mask', 'depth_mask_fin', 'depth_mask_inf', 'intrinsics']} + batch = { + 'label_type': [instance['label_type'] for instance in instances], + 'is_metric': [instance['is_metric'] for instance in instances], + 'info': [{'dataset': instance['dataset'], 'filename': instance['filename']} for instance in instances], + **batch, + } + return batch + + def get(self) -> Dict[str, Union[torch.Tensor, str]]: + return self.pipeline.get() + + def start(self): + self.pipeline.start() + + def stop(self): + self.pipeline.stop() + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.pipeline.terminate() + self.pipeline.join() + return False + + diff --git a/MoGe/moge/train/losses.py b/MoGe/moge/train/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..2b251b230f4cc86d8358f613acf483badfb49e14 --- /dev/null +++ b/MoGe/moge/train/losses.py @@ -0,0 +1,270 @@ +from typing import * +import math + +import torch +import torch.nn.functional as F +import utils3d + +from ..utils.geometry_torch import ( + weighted_mean, + harmonic_mean, + geometric_mean, + mask_aware_nearest_resize, + normalized_view_plane_uv, + angle_diff_vec3 +) +from ..utils.alignment import ( + align_points_scale_z_shift, + align_points_scale, + align_points_scale_xyz_shift, + align_points_z_shift, +) + + +def _smooth(err: torch.FloatTensor, beta: float = 0.0) -> torch.FloatTensor: + if beta == 0: + return err + else: + return torch.where(err < beta, 0.5 * err.square() / beta, err - 0.5 * beta) + + +def affine_invariant_global_loss( + pred_points: torch.Tensor, + gt_points: torch.Tensor, + mask: torch.Tensor, + align_resolution: int = 64, + beta: float = 0.0, + trunc: float = 1.0, + sparsity_aware: bool = False +): + device = pred_points.device + + # Align + (pred_points_lr, gt_points_lr), lr_mask = mask_aware_nearest_resize((pred_points, gt_points), mask=mask, size=(align_resolution, align_resolution)) + scale, shift = align_points_scale_z_shift(pred_points_lr.flatten(-3, -2), gt_points_lr.flatten(-3, -2), lr_mask.flatten(-2, -1) / gt_points_lr[..., 2].flatten(-2, -1).clamp_min(1e-2), trunc=trunc) + valid = scale > 0 + scale, shift = torch.where(valid, scale, 0), torch.where(valid[..., None], shift, 0) + + pred_points = scale[..., None, None, None] * pred_points + shift[..., None, None, :] + + # Compute loss + weight = (valid[..., None, None] & mask).float() / gt_points[..., 2].clamp_min(1e-5) + weight = weight.clamp_max(10.0 * weighted_mean(weight, mask, dim=(-2, -1), keepdim=True)) # In case your data contains extremely small depth values + loss = _smooth((pred_points - gt_points).abs() * weight[..., None], beta=beta).mean(dim=(-3, -2, -1)) + + if sparsity_aware: + # Reweighting improves performance on sparse depth data. NOTE: this is not used in MoGe-1. + sparsity = mask.float().mean(dim=(-2, -1)) / lr_mask.float().mean(dim=(-2, -1)) + loss = loss / (sparsity + 1e-7) + + err = (pred_points.detach() - gt_points).norm(dim=-1) / gt_points[..., 2] + + # Record any scalar metric + misc = { + 'truncated_error': weighted_mean(err.clamp_max(1.0), mask).item(), + 'delta': weighted_mean((err < 1).float(), mask).item() + } + + return loss, misc, scale.detach() + + +def monitoring(points: torch.Tensor): + return { + 'std': points.std().item(), + } + + +def compute_anchor_sampling_weight( + points: torch.Tensor, + mask: torch.Tensor, + radius_2d: torch.Tensor, + radius_3d: torch.Tensor, + num_test: int = 64 +) -> torch.Tensor: + # Importance sampling to balance the sampled probability of fine strutures. + # NOTE: MoGe-1 uses uniform random sampling instead of importance sampling. + # This is an incremental trick introduced later than the publication of MoGe-1 paper. + + height, width = points.shape[-3:-1] + + pixel_i, pixel_j = torch.meshgrid( + torch.arange(height, device=points.device), + torch.arange(width, device=points.device), + indexing='ij' + ) + + test_delta_i = torch.randint(-radius_2d, radius_2d + 1, (height, width, num_test,), device=points.device) # [num_test] + test_delta_j = torch.randint(-radius_2d, radius_2d + 1, (height, width, num_test,), device=points.device) # [num_test] + test_i, test_j = pixel_i[..., None] + test_delta_i, pixel_j[..., None] + test_delta_j # [height, width, num_test] + test_mask = (test_i >= 0) & (test_i < height) & (test_j >= 0) & (test_j < width) # [height, width, num_test] + test_i, test_j = test_i.clamp(0, height - 1), test_j.clamp(0, width - 1) # [height, width, num_test] + test_mask = test_mask & mask[..., test_i, test_j] # [..., height, width, num_test] + test_points = points[..., test_i, test_j, :] # [..., height, width, num_test, 3] + test_dist = (test_points - points[..., None, :]).norm(dim=-1) # [..., height, width, num_test] + + weight = 1 / ((test_dist <= radius_3d[..., None]) & test_mask).float().sum(dim=-1).clamp_min(1) + weight = torch.where(mask, weight, 0) + weight = weight / weight.sum(dim=(-2, -1), keepdim=True).add(1e-7) # [..., height, width] + return weight + + +def affine_invariant_local_loss( + pred_points: torch.Tensor, + gt_points: torch.Tensor, + gt_mask: torch.Tensor, + focal: torch.Tensor, + global_scale: torch.Tensor, + level: Literal[4, 16, 64], + align_resolution: int = 32, + num_patches: int = 16, + beta: float = 0.0, + trunc: float = 1.0, + sparsity_aware: bool = False +): + device, dtype = pred_points.device, pred_points.dtype + *batch_shape, height, width, _ = pred_points.shape + batch_size = math.prod(batch_shape) + pred_points, gt_points, gt_mask, focal, global_scale = pred_points.reshape(-1, height, width, 3), gt_points.reshape(-1, height, width, 3), gt_mask.reshape(-1, height, width), focal.reshape(-1), global_scale.reshape(-1) if global_scale is not None else None + + # Sample patch anchor points indices [num_total_patches] + radius_2d = math.ceil(0.5 / level * (height ** 2 + width ** 2) ** 0.5) + radius_3d = 0.5 / level / focal * gt_points[..., 2] + anchor_sampling_weights = compute_anchor_sampling_weight(gt_points, gt_mask, radius_2d, radius_3d, num_test=64) + where_mask = torch.where(gt_mask) + random_selection = torch.multinomial(anchor_sampling_weights[where_mask], num_patches * batch_size, replacement=True) + patch_batch_idx, patch_anchor_i, patch_anchor_j = [indices[random_selection] for indices in where_mask] # [num_total_patches] + + # Get patch indices [num_total_patches, patch_h, patch_w] + patch_i, patch_j = torch.meshgrid( + torch.arange(-radius_2d, radius_2d + 1, device=device), + torch.arange(-radius_2d, radius_2d + 1, device=device), + indexing='ij' + ) + patch_i, patch_j = patch_i + patch_anchor_i[:, None, None], patch_j + patch_anchor_j[:, None, None] + patch_mask = (patch_i >= 0) & (patch_i < height) & (patch_j >= 0) & (patch_j < width) + patch_i, patch_j = patch_i.clamp(0, height - 1), patch_j.clamp(0, width - 1) + + # Get patch mask and gt patch points + gt_patch_anchor_points = gt_points[patch_batch_idx, patch_anchor_i, patch_anchor_j] + gt_patch_radius_3d = 0.5 / level / focal[patch_batch_idx] * gt_patch_anchor_points[:, 2] + gt_patch_points = gt_points[patch_batch_idx[:, None, None], patch_i, patch_j] + gt_patch_dist = (gt_patch_points - gt_patch_anchor_points[:, None, None, :]).norm(dim=-1) + patch_mask &= gt_mask[patch_batch_idx[:, None, None], patch_i, patch_j] + patch_mask &= gt_patch_dist <= gt_patch_radius_3d[:, None, None] + + # Pick only non-empty patches + MINIMUM_POINTS_PER_PATCH = 32 + nonempty = torch.where(patch_mask.sum(dim=(-2, -1)) >= MINIMUM_POINTS_PER_PATCH) + num_nonempty_patches = nonempty[0].shape[0] + if num_nonempty_patches == 0: + return torch.tensor(0.0, dtype=dtype, device=device), {} + + # Finalize all patch variables + patch_batch_idx, patch_i, patch_j = patch_batch_idx[nonempty], patch_i[nonempty], patch_j[nonempty] + patch_mask = patch_mask[nonempty] # [num_nonempty_patches, patch_h, patch_w] + gt_patch_points = gt_patch_points[nonempty] # [num_nonempty_patches, patch_h, patch_w, 3] + gt_patch_radius_3d = gt_patch_radius_3d[nonempty] # [num_nonempty_patches] + gt_patch_anchor_points = gt_patch_anchor_points[nonempty] # [num_nonempty_patches, 3] + pred_patch_points = pred_points[patch_batch_idx[:, None, None], patch_i, patch_j] + + # Align patch points + (pred_patch_points_lr, gt_patch_points_lr), patch_lr_mask = mask_aware_nearest_resize((pred_patch_points, gt_patch_points), mask=patch_mask, size=(align_resolution, align_resolution)) + local_scale, local_shift = align_points_scale_xyz_shift(pred_patch_points_lr.flatten(-3, -2), gt_patch_points_lr.flatten(-3, -2), patch_lr_mask.flatten(-2) / gt_patch_radius_3d[:, None].add(1e-7), trunc=trunc) + if global_scale is not None: + scale_differ = local_scale / global_scale[patch_batch_idx] + patch_valid = (scale_differ > 0.1) & (scale_differ < 10.0) & (global_scale > 0) + else: + patch_valid = local_scale > 0 + local_scale, local_shift = torch.where(patch_valid, local_scale, 0), torch.where(patch_valid[:, None], local_shift, 0) + patch_mask &= patch_valid[:, None, None] + + pred_patch_points = local_scale[:, None, None, None] * pred_patch_points + local_shift[:, None, None, :] # [num_patches_nonempty, patch_h, patch_w, 3] + + # Compute loss + gt_mean = harmonic_mean(gt_points[..., 2], gt_mask, dim=(-2, -1)) + patch_weight = patch_mask.float() / gt_patch_points[..., 2].clamp_min(0.1 * gt_mean[patch_batch_idx, None, None]) # [num_patches_nonempty, patch_h, patch_w] + loss = _smooth((pred_patch_points - gt_patch_points).abs() * patch_weight[..., None], beta=beta).mean(dim=(-3, -2, -1)) # [num_patches_nonempty] + + if sparsity_aware: + # Reweighting improves performance on sparse depth data. NOTE: this is not used in MoGe-1. + sparsity = patch_mask.float().mean(dim=(-2, -1)) / patch_lr_mask.float().mean(dim=(-2, -1)) + loss = loss / (sparsity + 1e-7) + loss = torch.scatter_reduce(torch.zeros(batch_size, dtype=dtype, device=device), dim=0, index=patch_batch_idx, src=loss, reduce='sum') / num_patches + loss = loss.reshape(batch_shape) + + err = (pred_patch_points.detach() - gt_patch_points).norm(dim=-1) / gt_patch_radius_3d[..., None, None] + + # Record any scalar metric + misc = { + 'truncated_error': weighted_mean(err.clamp_max(1), patch_mask).item(), + 'delta': weighted_mean((err < 1).float(), patch_mask).item() + } + + return loss, misc + +def normal_loss(points: torch.Tensor, gt_points: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + device, dtype = points.device, points.dtype + height, width = points.shape[-3:-1] + + leftup, rightup, leftdown, rightdown = points[..., :-1, :-1, :], points[..., :-1, 1:, :], points[..., 1:, :-1, :], points[..., 1:, 1:, :] + upxleft = torch.cross(rightup - rightdown, leftdown - rightdown, dim=-1) + leftxdown = torch.cross(leftup - rightup, rightdown - rightup, dim=-1) + downxright = torch.cross(leftdown - leftup, rightup - leftup, dim=-1) + rightxup = torch.cross(rightdown - leftdown, leftup - leftdown, dim=-1) + + gt_leftup, gt_rightup, gt_leftdown, gt_rightdown = gt_points[..., :-1, :-1, :], gt_points[..., :-1, 1:, :], gt_points[..., 1:, :-1, :], gt_points[..., 1:, 1:, :] + gt_upxleft = torch.cross(gt_rightup - gt_rightdown, gt_leftdown - gt_rightdown, dim=-1) + gt_leftxdown = torch.cross(gt_leftup - gt_rightup, gt_rightdown - gt_rightup, dim=-1) + gt_downxright = torch.cross(gt_leftdown - gt_leftup, gt_rightup - gt_leftup, dim=-1) + gt_rightxup = torch.cross(gt_rightdown - gt_leftdown, gt_leftup - gt_leftdown, dim=-1) + + mask_leftup, mask_rightup, mask_leftdown, mask_rightdown = mask[..., :-1, :-1], mask[..., :-1, 1:], mask[..., 1:, :-1], mask[..., 1:, 1:] + mask_upxleft = mask_rightup & mask_leftdown & mask_rightdown + mask_leftxdown = mask_leftup & mask_rightdown & mask_rightup + mask_downxright = mask_leftdown & mask_rightup & mask_leftup + mask_rightxup = mask_rightdown & mask_leftup & mask_leftdown + + MIN_ANGLE, MAX_ANGLE, BETA_RAD = math.radians(1), math.radians(90), math.radians(3) + + loss = mask_upxleft * _smooth(angle_diff_vec3(upxleft, gt_upxleft).clamp(MIN_ANGLE, MAX_ANGLE), beta=BETA_RAD) \ + + mask_leftxdown * _smooth(angle_diff_vec3(leftxdown, gt_leftxdown).clamp(MIN_ANGLE, MAX_ANGLE), beta=BETA_RAD) \ + + mask_downxright * _smooth(angle_diff_vec3(downxright, gt_downxright).clamp(MIN_ANGLE, MAX_ANGLE), beta=BETA_RAD) \ + + mask_rightxup * _smooth(angle_diff_vec3(rightxup, gt_rightxup).clamp(MIN_ANGLE, MAX_ANGLE), beta=BETA_RAD) + + loss = loss.mean() / (4 * max(points.shape[-3:-1])) + + return loss, {} + + +def edge_loss(points: torch.Tensor, gt_points: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + device, dtype = points.device, points.dtype + height, width = points.shape[-3:-1] + + dx = points[..., :-1, :, :] - points[..., 1:, :, :] + dy = points[..., :, :-1, :] - points[..., :, 1:, :] + + gt_dx = gt_points[..., :-1, :, :] - gt_points[..., 1:, :, :] + gt_dy = gt_points[..., :, :-1, :] - gt_points[..., :, 1:, :] + + mask_dx = mask[..., :-1, :] & mask[..., 1:, :] + mask_dy = mask[..., :, :-1] & mask[..., :, 1:] + + MIN_ANGLE, MAX_ANGLE, BETA_RAD = math.radians(0.1), math.radians(90), math.radians(3) + + loss_dx = mask_dx * _smooth(angle_diff_vec3(dx, gt_dx).clamp(MIN_ANGLE, MAX_ANGLE), beta=BETA_RAD) + loss_dy = mask_dy * _smooth(angle_diff_vec3(dy, gt_dy).clamp(MIN_ANGLE, MAX_ANGLE), beta=BETA_RAD) + loss = (loss_dx.mean(dim=(-2, -1)) + loss_dy.mean(dim=(-2, -1))) / (2 * max(points.shape[-3:-1])) + + return loss, {} + + +def mask_l2_loss(pred_mask: torch.Tensor, gt_mask_pos: torch.Tensor, gt_mask_neg: torch.Tensor) -> torch.Tensor: + loss = gt_mask_neg.float() * pred_mask.square() + gt_mask_pos.float() * (1 - pred_mask).square() + loss = loss.mean(dim=(-2, -1)) + return loss, {} + + +def mask_bce_loss(pred_mask_prob: torch.Tensor, gt_mask_pos: torch.Tensor, gt_mask_neg: torch.Tensor) -> torch.Tensor: + loss = (gt_mask_pos | gt_mask_neg) * F.binary_cross_entropy(pred_mask_prob, gt_mask_pos.float(), reduction='none') + loss = loss.mean(dim=(-2, -1)) + return loss, {} diff --git a/MoGe/moge/train/utils.py b/MoGe/moge/train/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5f21e00876b927991381bf2f777a68b02c5b38cc --- /dev/null +++ b/MoGe/moge/train/utils.py @@ -0,0 +1,57 @@ +from typing import * +import fnmatch + +import sympy +import torch +import torch.nn as nn + + +def any_match(s: str, patterns: List[str]) -> bool: + return any(fnmatch.fnmatch(s, pat) for pat in patterns) + + +def build_optimizer(model: nn.Module, optimizer_config: Dict[str, Any]) -> torch.optim.Optimizer: + named_param_groups = [ + { + k: p for k, p in model.named_parameters() if any_match(k, param_group_config['params']['include']) and not any_match(k, param_group_config['params'].get('exclude', [])) + } for param_group_config in optimizer_config['params'] + ] + excluded_params = [k for k, p in model.named_parameters() if p.requires_grad and not any(k in named_params for named_params in named_param_groups)] + assert len(excluded_params) == 0, f'The following parameters require grad but are excluded from the optimizer: {excluded_params}' + optimizer_cls = getattr(torch.optim, optimizer_config['type']) + optimizer = optimizer_cls([ + { + **param_group_config, + 'params': list(params.values()), + } for param_group_config, params in zip(optimizer_config['params'], named_param_groups) + ]) + return optimizer + + +def parse_lr_lambda(s: str) -> Callable[[int], float]: + epoch = sympy.symbols('epoch') + lr_lambda = sympy.sympify(s) + return sympy.lambdify(epoch, lr_lambda, 'math') + + +def build_lr_scheduler(optimizer: torch.optim.Optimizer, scheduler_config: Dict[str, Any]) -> torch.optim.lr_scheduler._LRScheduler: + if scheduler_config['type'] == "SequentialLR": + child_schedulers = [ + build_lr_scheduler(optimizer, child_scheduler_config) + for child_scheduler_config in scheduler_config['params']['schedulers'] + ] + return torch.optim.lr_scheduler.SequentialLR(optimizer, schedulers=child_schedulers, milestones=scheduler_config['params']['milestones']) + elif scheduler_config['type'] == "LambdaLR": + lr_lambda = scheduler_config['params']['lr_lambda'] + if isinstance(lr_lambda, str): + lr_lambda = parse_lr_lambda(lr_lambda) + elif isinstance(lr_lambda, list): + lr_lambda = [parse_lr_lambda(l) for l in lr_lambda] + return torch.optim.lr_scheduler.LambdaLR( + optimizer, + lr_lambda=lr_lambda, + ) + else: + scheduler_cls = getattr(torch.optim.lr_scheduler, scheduler_config['type']) + scheduler = scheduler_cls(optimizer, **scheduler_config.get('params', {})) + return scheduler \ No newline at end of file diff --git a/MoGe/moge/utils/__init__.py b/MoGe/moge/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MoGe/moge/utils/alignment.py b/MoGe/moge/utils/alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..3d6bb78766ec1a43a89a4fc931b64f70c5201e2d --- /dev/null +++ b/MoGe/moge/utils/alignment.py @@ -0,0 +1,416 @@ +from typing import * +import math +from collections import namedtuple + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.types +import utils3d + + +def scatter_min(size: int, dim: int, index: torch.LongTensor, src: torch.Tensor) -> torch.return_types.min: + "Scatter the minimum value along the given dimension of `input` into `src` at the indices specified in `index`." + shape = src.shape[:dim] + (size,) + src.shape[dim + 1:] + minimum = torch.full(shape, float('inf'), dtype=src.dtype, device=src.device).scatter_reduce(dim=dim, index=index, src=src, reduce='amin', include_self=False) + minimum_where = torch.where(src == torch.gather(minimum, dim=dim, index=index)) + indices = torch.full(shape, -1, dtype=torch.long, device=src.device) + indices[(*minimum_where[:dim], index[minimum_where], *minimum_where[dim + 1:])] = minimum_where[dim] + return torch.return_types.min((minimum, indices)) + + +def split_batch_fwd(fn: Callable, chunk_size: int, *args, **kwargs): + batch_size = next(x for x in (*args, *kwargs.values()) if isinstance(x, torch.Tensor)).shape[0] + n_chunks = batch_size // chunk_size + (batch_size % chunk_size > 0) + splited_args = tuple(arg.split(chunk_size, dim=0) if isinstance(arg, torch.Tensor) else [arg] * n_chunks for arg in args) + splited_kwargs = {k: [v.split(chunk_size, dim=0) if isinstance(v, torch.Tensor) else [v] * n_chunks] for k, v in kwargs.items()} + results = [] + for i in range(n_chunks): + chunk_args = tuple(arg[i] for arg in splited_args) + chunk_kwargs = {k: v[i] for k, v in splited_kwargs.items()} + results.append(fn(*chunk_args, **chunk_kwargs)) + + if isinstance(results[0], tuple): + return tuple(torch.cat(r, dim=0) for r in zip(*results)) + else: + return torch.cat(results, dim=0) + + +def _pad_inf(x_: torch.Tensor): + return torch.cat([torch.full_like(x_[..., :1], -torch.inf), x_, torch.full_like(x_[..., :1], torch.inf)], dim=-1) + + +def _pad_cumsum(cumsum: torch.Tensor): + return torch.cat([torch.zeros_like(cumsum[..., :1]), cumsum, cumsum[..., -1:]], dim=-1) + + +def _compute_residual(a: torch.Tensor, xyw: torch.Tensor, trunc: float): + return a.mul(xyw[..., 0]).sub_(xyw[..., 1]).abs_().mul_(xyw[..., 2]).clamp_max_(trunc).sum(dim=-1) + + +def align(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor, trunc: Optional[Union[float, torch.Tensor]] = None, eps: float = 1e-7) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: + """ + If trunc is None, solve `min sum_i w_i * |a * x_i - y_i|`, otherwise solve `min sum_i min(trunc, w_i * |a * x_i - y_i|)`. + + w_i must be >= 0. + + ### Parameters: + - `x`: tensor of shape (..., n) + - `y`: tensor of shape (..., n) + - `w`: tensor of shape (..., n) + - `trunc`: optional, float or tensor of shape (..., n) or None + + ### Returns: + - `a`: tensor of shape (...), differentiable + - `loss`: tensor of shape (...), value of loss function at `a`, detached + - `index`: tensor of shape (...), where a = y[idx] / x[idx] + """ + if trunc is None: + x, y, w = torch.broadcast_tensors(x, y, w) + sign = torch.sign(x) + x, y = x * sign, y * sign + y_div_x = y / x.clamp_min(eps) + y_div_x, argsort = y_div_x.sort(dim=-1) + + wx = torch.gather(x * w, dim=-1, index=argsort) + derivatives = 2 * wx.cumsum(dim=-1) - wx.sum(dim=-1, keepdim=True) + search = torch.searchsorted(derivatives, torch.zeros_like(derivatives[..., :1]), side='left').clamp_max(derivatives.shape[-1] - 1) + + a = y_div_x.gather(dim=-1, index=search).squeeze(-1) + index = argsort.gather(dim=-1, index=search).squeeze(-1) + loss = (w * (a[..., None] * x - y).abs()).sum(dim=-1) + + else: + # Reshape to (batch_size, n) for simplicity + x, y, w = torch.broadcast_tensors(x, y, w) + batch_shape = x.shape[:-1] + batch_size = math.prod(batch_shape) + x, y, w = x.reshape(-1, x.shape[-1]), y.reshape(-1, y.shape[-1]), w.reshape(-1, w.shape[-1]) + + sign = torch.sign(x) + x, y = x * sign, y * sign + wx, wy = w * x, w * y + xyw = torch.stack([x, y, w], dim=-1) # Stacked for convenient gathering + + y_div_x = A = y / x.clamp_min(eps) + B = (wy - trunc) / wx.clamp_min(eps) + C = (wy + trunc) / wx.clamp_min(eps) + with torch.no_grad(): + # Caculate prefix sum by orders of A, B, C + A, A_argsort = A.sort(dim=-1) + Q_A = torch.cumsum(torch.gather(wx, dim=-1, index=A_argsort), dim=-1) + A, Q_A = _pad_inf(A), _pad_cumsum(Q_A) # Pad [-inf, A1, ..., An, inf] and [0, Q1, ..., Qn, Qn] to handle edge cases. + + B, B_argsort = B.sort(dim=-1) + Q_B = torch.cumsum(torch.gather(wx, dim=-1, index=B_argsort), dim=-1) + B, Q_B = _pad_inf(B), _pad_cumsum(Q_B) + + C, C_argsort = C.sort(dim=-1) + Q_C = torch.cumsum(torch.gather(wx, dim=-1, index=C_argsort), dim=-1) + C, Q_C = _pad_inf(C), _pad_cumsum(Q_C) + + # Caculate left and right derivative of A + j_A = torch.searchsorted(A, y_div_x, side='left').sub_(1) + j_B = torch.searchsorted(B, y_div_x, side='left').sub_(1) + j_C = torch.searchsorted(C, y_div_x, side='left').sub_(1) + left_derivative = 2 * torch.gather(Q_A, dim=-1, index=j_A) - torch.gather(Q_B, dim=-1, index=j_B) - torch.gather(Q_C, dim=-1, index=j_C) + j_A = torch.searchsorted(A, y_div_x, side='right').sub_(1) + j_B = torch.searchsorted(B, y_div_x, side='right').sub_(1) + j_C = torch.searchsorted(C, y_div_x, side='right').sub_(1) + right_derivative = 2 * torch.gather(Q_A, dim=-1, index=j_A) - torch.gather(Q_B, dim=-1, index=j_B) - torch.gather(Q_C, dim=-1, index=j_C) + + # Find extrema + is_extrema = (left_derivative < 0) & (right_derivative >= 0) + is_extrema[..., 0] |= ~is_extrema.any(dim=-1) # In case all derivatives are zero, take the first one as extrema. + where_extrema_batch, where_extrema_index = torch.where(is_extrema) + + # Calculate objective value at extrema + extrema_a = y_div_x[where_extrema_batch, where_extrema_index] # (num_extrema,) + MAX_ELEMENTS = 4096 ** 2 # Split into small batches to avoid OOM in case there are too many extrema.(~1G) + SPLIT_SIZE = MAX_ELEMENTS // x.shape[-1] + extrema_value = torch.cat([ + _compute_residual(extrema_a_split[:, None], xyw[extrema_i_split, :, :], trunc) + for extrema_a_split, extrema_i_split in zip(extrema_a.split(SPLIT_SIZE), where_extrema_batch.split(SPLIT_SIZE)) + ]) # (num_extrema,) + + # Find minima among corresponding extrema + minima, indices = scatter_min(size=batch_size, dim=0, index=where_extrema_batch, src=extrema_value) # (batch_size,) + index = where_extrema_index[indices] + + a = torch.gather(y, dim=-1, index=index[..., None]) / torch.gather(x, dim=-1, index=index[..., None]).clamp_min(eps) + a = a.reshape(batch_shape) + loss = minima.reshape(batch_shape) + index = index.reshape(batch_shape) + + return a, loss, index + + +def align_depth_scale(depth_src: torch.Tensor, depth_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None): + """ + Align `depth_src` to `depth_tgt` with given constant weights. + + ### Parameters: + - `depth_src: torch.Tensor` of shape (..., N) + - `depth_tgt: torch.Tensor` of shape (..., N) + + """ + scale, _, _ = align(depth_src, depth_tgt, weight, trunc) + + return scale + + +def align_depth_affine(depth_src: torch.Tensor, depth_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None): + """ + Align `depth_src` to `depth_tgt` with given constant weights. + + ### Parameters: + - `depth_src: torch.Tensor` of shape (..., N) + - `depth_tgt: torch.Tensor` of shape (..., N) + - `weight: torch.Tensor` of shape (..., N) + - `trunc: float` or tensor of shape (..., N) or None + + ### Returns: + - `scale: torch.Tensor` of shape (...). + - `shift: torch.Tensor` of shape (...). + """ + dtype, device = depth_src.dtype, depth_src.device + + # Flatten batch dimensions for simplicity + batch_shape, n = depth_src.shape[:-1], depth_src.shape[-1] + batch_size = math.prod(batch_shape) + depth_src, depth_tgt, weight = depth_src.reshape(batch_size, n), depth_tgt.reshape(batch_size, n), weight.reshape(batch_size, n) + + # Here, we take anchors only for non-zero weights. + # Although the results will be still correct even anchor points have zero weight, + # it is wasting computation and may cause instability in some cases, e.g. too many extrema. + anchors_where_batch, anchors_where_n = torch.where(weight > 0) + + # Stop gradient when solving optimal anchors + with torch.no_grad(): + depth_src_anchor = depth_src[anchors_where_batch, anchors_where_n] # (anchors) + depth_tgt_anchor = depth_tgt[anchors_where_batch, anchors_where_n] # (anchors) + + depth_src_anchored = depth_src[anchors_where_batch, :] - depth_src_anchor[..., None] # (anchors, n) + depth_tgt_anchored = depth_tgt[anchors_where_batch, :] - depth_tgt_anchor[..., None] # (anchors, n) + weight_anchored = weight[anchors_where_batch, :] # (anchors, n) + + scale, loss, index = align(depth_src_anchored, depth_tgt_anchored, weight_anchored, trunc) # (anchors) + + loss, index_anchor = scatter_min(size=batch_size, dim=0, index=anchors_where_batch, src=loss) # (batch_size,) + + # Reproduce by indexing for shorter compute graph + index_1 = anchors_where_n[index_anchor] # (batch_size,) + index_2 = index[index_anchor] # (batch_size,) + + tgt_1, src_1 = torch.gather(depth_tgt, dim=1, index=index_1[..., None]).squeeze(-1), torch.gather(depth_src, dim=1, index=index_1[..., None]).squeeze(-1) + tgt_2, src_2 = torch.gather(depth_tgt, dim=1, index=index_2[..., None]).squeeze(-1), torch.gather(depth_src, dim=1, index=index_2[..., None]).squeeze(-1) + + scale = (tgt_2 - tgt_1) / torch.where(src_2 != src_1, src_2 - src_1, 1e-7) + shift = tgt_1 - scale * src_1 + + scale, shift = scale.reshape(batch_shape), shift.reshape(batch_shape) + + return scale, shift + +def align_depth_affine_irls(depth_src: torch.Tensor, depth_tgt: torch.Tensor, weight: Optional[torch.Tensor], max_iter: int = 100, eps: float = 1e-12): + """ + Align `depth_src` to `depth_tgt` with given constant weights using IRLS. + """ + dtype, device = depth_src.dtype, depth_src.device + + w = weight + x = torch.stack([depth_src, torch.ones_like(depth_src)], dim=-1) + y = depth_tgt + + for i in range(max_iter): + beta = (x.transpose(-1, -2) @ (w * y)) @ (x.transpose(-1, -2) @ (w[..., None] * x)).inverse().transpose(-2, -1) + w = 1 / (y - (x @ beta[..., None])[..., 0]).abs().clamp_min(eps) + + return beta[..., 0], beta[..., 1] + + +def align_points_scale(points_src: torch.Tensor, points_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None): + """ + ### Parameters: + - `points_src: torch.Tensor` of shape (..., N, 3) + - `points_tgt: torch.Tensor` of shape (..., N, 3) + - `weight: torch.Tensor` of shape (..., N) + + ### Returns: + - `a: torch.Tensor` of shape (...). Only positive solutions are garunteed. You should filter out negative scales before using it. + - `b: torch.Tensor` of shape (...) + """ + dtype, device = points_src.dtype, points_src.device + + scale, _, _ = align(points_src.flatten(-2), points_tgt.flatten(-2), weight[..., None].expand_as(points_src).flatten(-2), trunc) + + return scale + + +def align_points_scale_z_shift(points_src: torch.Tensor, points_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None): + """ + Align `points_src` to `points_tgt` with respect to a shared xyz scale and z shift. + It is similar to `align_affine` but scale and shift are applied to different dimensions. + + ### Parameters: + - `points_src: torch.Tensor` of shape (..., N, 3) + - `points_tgt: torch.Tensor` of shape (..., N, 3) + - `weights: torch.Tensor` of shape (..., N) + + ### Returns: + - `scale: torch.Tensor` of shape (...). + - `shift: torch.Tensor` of shape (..., 3). x and y shifts are zeros. + """ + dtype, device = points_src.dtype, points_src.device + + # Flatten batch dimensions for simplicity + batch_shape, n = points_src.shape[:-2], points_src.shape[-2] + batch_size = math.prod(batch_shape) + points_src, points_tgt, weight = points_src.reshape(batch_size, n, 3), points_tgt.reshape(batch_size, n, 3), weight.reshape(batch_size, n) + + # Take anchors + anchor_where_batch, anchor_where_n = torch.where(weight > 0) + with torch.no_grad(): + zeros = torch.zeros(anchor_where_batch.shape[0], device=device, dtype=dtype) + points_src_anchor = torch.stack([zeros, zeros, points_src[anchor_where_batch, anchor_where_n, 2]], dim=-1) # (anchors, 3) + points_tgt_anchor = torch.stack([zeros, zeros, points_tgt[anchor_where_batch, anchor_where_n, 2]], dim=-1) # (anchors, 3) + + points_src_anchored = points_src[anchor_where_batch, :, :] - points_src_anchor[..., None, :] # (anchors, n, 3) + points_tgt_anchored = points_tgt[anchor_where_batch, :, :] - points_tgt_anchor[..., None, :] # (anchors, n, 3) + weight_anchored = weight[anchor_where_batch, :, None].expand(-1, -1, 3) # (anchors, n, 3) + + # Solve optimal scale and shift for each anchor + MAX_ELEMENTS = 2 ** 20 + scale, loss, index = split_batch_fwd(align, MAX_ELEMENTS // n, points_src_anchored.flatten(-2), points_tgt_anchored.flatten(-2), weight_anchored.flatten(-2), trunc) # (anchors,) + + loss, index_anchor = scatter_min(size=batch_size, dim=0, index=anchor_where_batch, src=loss) # (batch_size,) + + # Reproduce by indexing for shorter compute graph + index_2 = index[index_anchor] # (batch_size,) [0, 3n) + index_1 = anchor_where_n[index_anchor] * 3 + index_2 % 3 # (batch_size,) [0, 3n) + + zeros = torch.zeros((batch_size, n), device=device, dtype=dtype) + points_tgt_00z, points_src_00z = torch.stack([zeros, zeros, points_tgt[..., 2]], dim=-1), torch.stack([zeros, zeros, points_src[..., 2]], dim=-1) + tgt_1, src_1 = torch.gather(points_tgt_00z.flatten(-2), dim=1, index=index_1[..., None]).squeeze(-1), torch.gather(points_src_00z.flatten(-2), dim=1, index=index_1[..., None]).squeeze(-1) + tgt_2, src_2 = torch.gather(points_tgt.flatten(-2), dim=1, index=index_2[..., None]).squeeze(-1), torch.gather(points_src.flatten(-2), dim=1, index=index_2[..., None]).squeeze(-1) + + scale = (tgt_2 - tgt_1) / torch.where(src_2 != src_1, src_2 - src_1, 1.0) + shift = torch.gather(points_tgt_00z, dim=1, index=(index_1 // 3)[..., None, None].expand(-1, -1, 3)).squeeze(-2) - scale[..., None] * torch.gather(points_src_00z, dim=1, index=(index_1 // 3)[..., None, None].expand(-1, -1, 3)).squeeze(-2) + scale, shift = scale.reshape(batch_shape), shift.reshape(*batch_shape, 3) + + return scale, shift + + +def align_points_scale_xyz_shift(points_src: torch.Tensor, points_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None, max_iters: int = 30, eps: float = 1e-6): + """ + Align `points_src` to `points_tgt` with respect to a shared xyz scale and z shift. + It is similar to `align_affine` but scale and shift are applied to different dimensions. + + ### Parameters: + - `points_src: torch.Tensor` of shape (..., N, 3) + - `points_tgt: torch.Tensor` of shape (..., N, 3) + - `weights: torch.Tensor` of shape (..., N) + + ### Returns: + - `scale: torch.Tensor` of shape (...). + - `shift: torch.Tensor` of shape (..., 3) + """ + dtype, device = points_src.dtype, points_src.device + + # Flatten batch dimensions for simplicity + batch_shape, n = points_src.shape[:-2], points_src.shape[-2] + batch_size = math.prod(batch_shape) + points_src, points_tgt, weight = points_src.reshape(batch_size, n, 3), points_tgt.reshape(batch_size, n, 3), weight.reshape(batch_size, n) + + # Take anchors + anchor_where_batch, anchor_where_n = torch.where(weight > 0) + + with torch.no_grad(): + points_src_anchor = points_src[anchor_where_batch, anchor_where_n] # (anchors, 3) + points_tgt_anchor = points_tgt[anchor_where_batch, anchor_where_n] # (anchors, 3) + + points_src_anchored = points_src[anchor_where_batch, :, :] - points_src_anchor[..., None, :] # (anchors, n, 3) + points_tgt_anchored = points_tgt[anchor_where_batch, :, :] - points_tgt_anchor[..., None, :] # (anchors, n, 3) + weight_anchored = weight[anchor_where_batch, :, None].expand(-1, -1, 3) # (anchors, n, 3) + + # Solve optimal scale and shift for each anchor + MAX_ELEMENTS = 2 ** 20 + scale, loss, index = split_batch_fwd(align, MAX_ELEMENTS // 2, points_src_anchored.flatten(-2), points_tgt_anchored.flatten(-2), weight_anchored.flatten(-2), trunc) # (anchors,) + + # Get optimal scale and shift for each batch element + loss, index_anchor = scatter_min(size=batch_size, dim=0, index=anchor_where_batch, src=loss) # (batch_size,) + + index_2 = index[index_anchor] # (batch_size,) [0, 3n) + index_1 = anchor_where_n[index_anchor] * 3 + index_2 % 3 # (batch_size,) [0, 3n) + + src_1, tgt_1 = torch.gather(points_src.flatten(-2), dim=1, index=index_1[..., None]).squeeze(-1), torch.gather(points_tgt.flatten(-2), dim=1, index=index_1[..., None]).squeeze(-1) + src_2, tgt_2 = torch.gather(points_src.flatten(-2), dim=1, index=index_2[..., None]).squeeze(-1), torch.gather(points_tgt.flatten(-2), dim=1, index=index_2[..., None]).squeeze(-1) + + scale = (tgt_2 - tgt_1) / torch.where(src_2 != src_1, src_2 - src_1, 1.0) + shift = torch.gather(points_tgt, dim=1, index=(index_1 // 3)[..., None, None].expand(-1, -1, 3)).squeeze(-2) - scale[..., None] * torch.gather(points_src, dim=1, index=(index_1 // 3)[..., None, None].expand(-1, -1, 3)).squeeze(-2) + + scale, shift = scale.reshape(batch_shape), shift.reshape(*batch_shape, 3) + + return scale, shift + + +def align_points_z_shift(points_src: torch.Tensor, points_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None, max_iters: int = 30, eps: float = 1e-6): + """ + Align `points_src` to `points_tgt` with respect to a Z-axis shift. + + ### Parameters: + - `points_src: torch.Tensor` of shape (..., N, 3) + - `points_tgt: torch.Tensor` of shape (..., N, 3) + - `weights: torch.Tensor` of shape (..., N) + + ### Returns: + - `scale: torch.Tensor` of shape (...). + - `shift: torch.Tensor` of shape (..., 3) + """ + dtype, device = points_src.dtype, points_src.device + + shift, _, _ = align(torch.ones_like(points_src[..., 2]), points_tgt[..., 2] - points_src[..., 2], weight, trunc) + shift = torch.stack([torch.zeros_like(shift), torch.zeros_like(shift), shift], dim=-1) + + return shift + + +def align_points_xyz_shift(points_src: torch.Tensor, points_tgt: torch.Tensor, weight: Optional[torch.Tensor], trunc: Optional[Union[float, torch.Tensor]] = None, max_iters: int = 30, eps: float = 1e-6): + """ + Align `points_src` to `points_tgt` with respect to a Z-axis shift. + + ### Parameters: + - `points_src: torch.Tensor` of shape (..., N, 3) + - `points_tgt: torch.Tensor` of shape (..., N, 3) + - `weights: torch.Tensor` of shape (..., N) + + ### Returns: + - `scale: torch.Tensor` of shape (...). + - `shift: torch.Tensor` of shape (..., 3) + """ + dtype, device = points_src.dtype, points_src.device + + shift, _, _ = align(torch.ones_like(points_src).swapaxes(-2, -1), (points_tgt - points_src).swapaxes(-2, -1), weight[..., None, :], trunc) + + return shift + + +def align_affine_lstsq(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Solve `min sum_i w_i * (a * x_i + b - y_i ) ^ 2`, where `a` and `b` are scalars, with respect to `a` and `b` using least squares. + + ### Parameters: + - `x: torch.Tensor` of shape (..., N) + - `y: torch.Tensor` of shape (..., N) + - `w: torch.Tensor` of shape (..., N) + + ### Returns: + - `a: torch.Tensor` of shape (...,) + - `b: torch.Tensor` of shape (...,) + """ + w_sqrt = torch.ones_like(x) if w is None else w.sqrt() + A = torch.stack([w_sqrt * x, torch.ones_like(x)], dim=-1) + B = (w_sqrt * y)[..., None] + a, b = torch.linalg.lstsq(A, B)[0].squeeze(-1).unbind(-1) + return a, b \ No newline at end of file diff --git a/MoGe/moge/utils/download.py b/MoGe/moge/utils/download.py new file mode 100644 index 0000000000000000000000000000000000000000..886edbccc81cc0c3daed4d858f641097bdfceee2 --- /dev/null +++ b/MoGe/moge/utils/download.py @@ -0,0 +1,55 @@ +from pathlib import Path +from typing import * +import requests + +from tqdm import tqdm + + +__all__ = ["download_file", "download_bytes"] + + +def download_file(url: str, filepath: Union[str, Path], headers: dict = None, resume: bool = True) -> None: + # Ensure headers is a dict if not provided + headers = headers or {} + + # Initialize local variables + file_path = Path(filepath) + downloaded_bytes = 0 + + # Check if we should resume the download + if resume and file_path.exists(): + downloaded_bytes = file_path.stat().st_size + headers['Range'] = f"bytes={downloaded_bytes}-" + + # Make a GET request to fetch the file + with requests.get(url, stream=True, headers=headers) as response: + response.raise_for_status() # This will raise an HTTPError if the status is 4xx/5xx + + # Calculate the total size to download + total_size = downloaded_bytes + int(response.headers.get('content-length', 0)) + + # Display a progress bar while downloading + with ( + tqdm(desc=f"Downloading {file_path.name}", total=total_size, unit='B', unit_scale=True, leave=False) as pbar, + open(file_path, 'ab') as file, + ): + # Set the initial position of the progress bar + pbar.update(downloaded_bytes) + + # Write the content to the file in chunks + for chunk in response.iter_content(chunk_size=4096): + file.write(chunk) + pbar.update(len(chunk)) + + +def download_bytes(url: str, headers: dict = None) -> bytes: + # Ensure headers is a dict if not provided + headers = headers or {} + + # Make a GET request to fetch the file + with requests.get(url, stream=True, headers=headers) as response: + response.raise_for_status() # This will raise an HTTPError if the status is 4xx/5xx + + # Read the content of the response + return response.content + \ No newline at end of file diff --git a/MoGe/moge/utils/geometry_numpy.py b/MoGe/moge/utils/geometry_numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..0d4274dc6666eb2188e26431f9010c3d4be3d4e1 --- /dev/null +++ b/MoGe/moge/utils/geometry_numpy.py @@ -0,0 +1,401 @@ +from typing import * +from functools import partial +import math + +import cv2 +import numpy as np +from scipy.signal import fftconvolve +import numpy as np +import utils3d + +from .tools import timeit + + +def weighted_mean_numpy(x: np.ndarray, w: np.ndarray = None, axis: Union[int, Tuple[int,...]] = None, keepdims: bool = False, eps: float = 1e-7) -> np.ndarray: + if w is None: + return np.mean(x, axis=axis) + else: + w = w.astype(x.dtype) + return (x * w).mean(axis=axis) / np.clip(w.mean(axis=axis), eps, None) + + +def harmonic_mean_numpy(x: np.ndarray, w: np.ndarray = None, axis: Union[int, Tuple[int,...]] = None, keepdims: bool = False, eps: float = 1e-7) -> np.ndarray: + if w is None: + return 1 / (1 / np.clip(x, eps, None)).mean(axis=axis) + else: + w = w.astype(x.dtype) + return 1 / (weighted_mean_numpy(1 / (x + eps), w, axis=axis, keepdims=keepdims, eps=eps) + eps) + + +def normalized_view_plane_uv_numpy(width: int, height: int, aspect_ratio: float = None, dtype: np.dtype = np.float32) -> np.ndarray: + "UV with left-top corner as (-width / diagonal, -height / diagonal) and right-bottom corner as (width / diagonal, height / diagonal)" + if aspect_ratio is None: + aspect_ratio = width / height + + span_x = aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5 + span_y = 1 / (1 + aspect_ratio ** 2) ** 0.5 + + u = np.linspace(-span_x * (width - 1) / width, span_x * (width - 1) / width, width, dtype=dtype) + v = np.linspace(-span_y * (height - 1) / height, span_y * (height - 1) / height, height, dtype=dtype) + u, v = np.meshgrid(u, v, indexing='xy') + uv = np.stack([u, v], axis=-1) + return uv + + +def focal_to_fov_numpy(focal: np.ndarray): + return 2 * np.arctan(0.5 / focal) + + +def fov_to_focal_numpy(fov: np.ndarray): + return 0.5 / np.tan(fov / 2) + + +def intrinsics_to_fov_numpy(intrinsics: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + fov_x = focal_to_fov_numpy(intrinsics[..., 0, 0]) + fov_y = focal_to_fov_numpy(intrinsics[..., 1, 1]) + return fov_x, fov_y + + +def point_map_to_depth_legacy_numpy(points: np.ndarray): + height, width = points.shape[-3:-1] + diagonal = (height ** 2 + width ** 2) ** 0.5 + uv = normalized_view_plane_uv_numpy(width, height, dtype=points.dtype) # (H, W, 2) + _, uv = np.broadcast_arrays(points[..., :2], uv) + + # Solve least squares problem + b = (uv * points[..., 2:]).reshape(*points.shape[:-3], -1) # (..., H * W * 2) + A = np.stack([points[..., :2], -uv], axis=-1).reshape(*points.shape[:-3], -1, 2) # (..., H * W * 2, 2) + + M = A.swapaxes(-2, -1) @ A + solution = (np.linalg.inv(M + 1e-6 * np.eye(2)) @ (A.swapaxes(-2, -1) @ b[..., None])).squeeze(-1) + focal, shift = solution + + depth = points[..., 2] + shift[..., None, None] + fov_x = np.arctan(width / diagonal / focal) * 2 + fov_y = np.arctan(height / diagonal / focal) * 2 + return depth, fov_x, fov_y, shift + + +def solve_optimal_focal_shift(uv: np.ndarray, xyz: np.ndarray): + "Solve `min |focal * xy / (z + shift) - uv|` with respect to shift and focal" + from scipy.optimize import least_squares + uv, xy, z = uv.reshape(-1, 2), xyz[..., :2].reshape(-1, 2), xyz[..., 2].reshape(-1) + + def fn(uv: np.ndarray, xy: np.ndarray, z: np.ndarray, shift: np.ndarray): + xy_proj = xy / (z + shift)[: , None] + f = (xy_proj * uv).sum() / np.square(xy_proj).sum() + err = (f * xy_proj - uv).ravel() + return err + + solution = least_squares(partial(fn, uv, xy, z), x0=0, ftol=1e-3, method='lm') + optim_shift = solution['x'].squeeze().astype(np.float32) + + xy_proj = xy / (z + optim_shift)[: , None] + optim_focal = (xy_proj * uv).sum() / np.square(xy_proj).sum() + + return optim_shift, optim_focal + + +def solve_optimal_shift(uv: np.ndarray, xyz: np.ndarray, focal: float): + "Solve `min |focal * xy / (z + shift) - uv|` with respect to shift" + from scipy.optimize import least_squares + uv, xy, z = uv.reshape(-1, 2), xyz[..., :2].reshape(-1, 2), xyz[..., 2].reshape(-1) + + def fn(uv: np.ndarray, xy: np.ndarray, z: np.ndarray, shift: np.ndarray): + xy_proj = xy / (z + shift)[: , None] + err = (focal * xy_proj - uv).ravel() + return err + + solution = least_squares(partial(fn, uv, xy, z), x0=0, ftol=1e-3, method='lm') + optim_shift = solution['x'].squeeze().astype(np.float32) + + return optim_shift + + +def recover_focal_shift_numpy(points: np.ndarray, mask: np.ndarray = None, focal: float = None, downsample_size: Tuple[int, int] = (64, 64)): + import cv2 + assert points.shape[-1] == 3, "Points should (H, W, 3)" + + height, width = points.shape[-3], points.shape[-2] + diagonal = (height ** 2 + width ** 2) ** 0.5 + + uv = normalized_view_plane_uv_numpy(width=width, height=height) + + if mask is None: + points_lr = cv2.resize(points, downsample_size, interpolation=cv2.INTER_LINEAR).reshape(-1, 3) + uv_lr = cv2.resize(uv, downsample_size, interpolation=cv2.INTER_LINEAR).reshape(-1, 2) + else: + (points_lr, uv_lr), mask_lr = mask_aware_nearest_resize_numpy((points, uv), mask, downsample_size) + + if points_lr.size < 2: + return 1., 0. + + if focal is None: + focal, shift = solve_optimal_focal_shift(uv_lr, points_lr) + else: + shift = solve_optimal_shift(uv_lr, points_lr, focal) + + return focal, shift + + +def mask_aware_nearest_resize_numpy( + inputs: Union[np.ndarray, Tuple[np.ndarray, ...], None], + mask: np.ndarray, + size: Tuple[int, int], + return_index: bool = False +) -> Tuple[Union[np.ndarray, Tuple[np.ndarray, ...], None], np.ndarray, Tuple[np.ndarray, ...]]: + """ + Resize 2D map by nearest interpolation. Return the nearest neighbor index and mask of the resized map. + + ### Parameters + - `inputs`: a single or a list of input 2D map(s) of shape (..., H, W, ...). + - `mask`: input 2D mask of shape (..., H, W) + - `size`: target size (width, height) + + ### Returns + - `*resized_maps`: resized map(s) of shape (..., target_height, target_width, ...). + - `resized_mask`: mask of the resized map of shape (..., target_height, target_width) + - `nearest_idx`: if return_index is True, nearest neighbor index of the resized map of shape (..., target_height, target_width) for each dimension. + """ + height, width = mask.shape[-2:] + target_width, target_height = size + filter_h_f, filter_w_f = max(1, height / target_height), max(1, width / target_width) + filter_h_i, filter_w_i = math.ceil(filter_h_f), math.ceil(filter_w_f) + filter_size = filter_h_i * filter_w_i + padding_h, padding_w = filter_h_i // 2 + 1, filter_w_i // 2 + 1 + + # Window the original mask and uv + uv = utils3d.numpy.image_pixel_center(width=width, height=height, dtype=np.float32) + indices = np.arange(height * width, dtype=np.int32).reshape(height, width) + padded_uv = np.full((height + 2 * padding_h, width + 2 * padding_w, 2), 0, dtype=np.float32) + padded_uv[padding_h:padding_h + height, padding_w:padding_w + width] = uv + padded_mask = np.full((*mask.shape[:-2], height + 2 * padding_h, width + 2 * padding_w), False, dtype=bool) + padded_mask[..., padding_h:padding_h + height, padding_w:padding_w + width] = mask + padded_indices = np.full((height + 2 * padding_h, width + 2 * padding_w), 0, dtype=np.int32) + padded_indices[padding_h:padding_h + height, padding_w:padding_w + width] = indices + windowed_uv = utils3d.numpy.sliding_window_2d(padded_uv, (filter_h_i, filter_w_i), 1, axis=(0, 1)) + windowed_mask = utils3d.numpy.sliding_window_2d(padded_mask, (filter_h_i, filter_w_i), 1, axis=(-2, -1)) + windowed_indices = utils3d.numpy.sliding_window_2d(padded_indices, (filter_h_i, filter_w_i), 1, axis=(0, 1)) + + # Gather the target pixels's local window + target_centers = utils3d.numpy.image_uv(width=target_width, height=target_height, dtype=np.float32) * np.array([width, height], dtype=np.float32) + target_lefttop = target_centers - np.array((filter_w_f / 2, filter_h_f / 2), dtype=np.float32) + target_window = np.round(target_lefttop).astype(np.int32) + np.array((padding_w, padding_h), dtype=np.int32) + + target_window_centers = windowed_uv[target_window[..., 1], target_window[..., 0], :, :, :].reshape(target_height, target_width, 2, filter_size) # (target_height, tgt_width, 2, filter_size) + target_window_mask = windowed_mask[..., target_window[..., 1], target_window[..., 0], :, :].reshape(*mask.shape[:-2], target_height, target_width, filter_size) # (..., target_height, tgt_width, filter_size) + target_window_indices = windowed_indices[target_window[..., 1], target_window[..., 0], :, :].reshape(*([-1] * (mask.ndim - 2)), target_height, target_width, filter_size) # (target_height, tgt_width, filter_size) + + # Compute nearest neighbor in the local window for each pixel + dist = np.square(target_window_centers - target_centers[..., None]) + dist = dist[..., 0, :] + dist[..., 1, :] + dist = np.where(target_window_mask, dist, np.inf) # (..., target_height, tgt_width, filter_size) + nearest_in_window = np.argmin(dist, axis=-1, keepdims=True) # (..., target_height, tgt_width, 1) + nearest_idx = np.take_along_axis(target_window_indices, nearest_in_window, axis=-1).squeeze(-1) # (..., target_height, tgt_width) + nearest_i, nearest_j = nearest_idx // width, nearest_idx % width + target_mask = np.any(target_window_mask, axis=-1) + batch_indices = [np.arange(n).reshape([1] * i + [n] + [1] * (mask.ndim - i - 1)) for i, n in enumerate(mask.shape[:-2])] + + index = (*batch_indices, nearest_i, nearest_j) + + if inputs is None: + outputs = None + elif isinstance(inputs, np.ndarray): + outputs = inputs[index] + elif isinstance(inputs, Sequence): + outputs = tuple(x[index] for x in inputs) + else: + raise ValueError(f'Invalid input type: {type(inputs)}') + + if return_index: + return outputs, target_mask, index + else: + return outputs, target_mask + + +def mask_aware_area_resize_numpy(image: np.ndarray, mask: np.ndarray, target_width: int, target_height: int) -> Tuple[Tuple[np.ndarray, ...], np.ndarray]: + """ + Resize 2D map by nearest interpolation. Return the nearest neighbor index and mask of the resized map. + + ### Parameters + - `image`: Input 2D image of shape (..., H, W, C) + - `mask`: Input 2D mask of shape (..., H, W) + - `target_width`: target width of the resized map + - `target_height`: target height of the resized map + + ### Returns + - `nearest_idx`: Nearest neighbor index of the resized map of shape (..., target_height, target_width). + - `target_mask`: Mask of the resized map of shape (..., target_height, target_width) + """ + height, width = mask.shape[-2:] + + if image.shape[-2:] == (height, width): + omit_channel_dim = True + else: + omit_channel_dim = False + if omit_channel_dim: + image = image[..., None] + + image = np.where(mask[..., None], image, 0) + + filter_h_f, filter_w_f = max(1, height / target_height), max(1, width / target_width) + filter_h_i, filter_w_i = math.ceil(filter_h_f) + 1, math.ceil(filter_w_f) + 1 + filter_size = filter_h_i * filter_w_i + padding_h, padding_w = filter_h_i // 2 + 1, filter_w_i // 2 + 1 + + # Window the original mask and uv (non-copy) + uv = utils3d.numpy.image_pixel_center(width=width, height=height, dtype=np.float32) + indices = np.arange(height * width, dtype=np.int32).reshape(height, width) + padded_uv = np.full((height + 2 * padding_h, width + 2 * padding_w, 2), 0, dtype=np.float32) + padded_uv[padding_h:padding_h + height, padding_w:padding_w + width] = uv + padded_mask = np.full((*mask.shape[:-2], height + 2 * padding_h, width + 2 * padding_w), False, dtype=bool) + padded_mask[..., padding_h:padding_h + height, padding_w:padding_w + width] = mask + padded_indices = np.full((height + 2 * padding_h, width + 2 * padding_w), 0, dtype=np.int32) + padded_indices[padding_h:padding_h + height, padding_w:padding_w + width] = indices + windowed_uv = utils3d.numpy.sliding_window_2d(padded_uv, (filter_h_i, filter_w_i), 1, axis=(0, 1)) + windowed_mask = utils3d.numpy.sliding_window_2d(padded_mask, (filter_h_i, filter_w_i), 1, axis=(-2, -1)) + windowed_indices = utils3d.numpy.sliding_window_2d(padded_indices, (filter_h_i, filter_w_i), 1, axis=(0, 1)) + + # Gather the target pixels's local window + target_center = utils3d.numpy.image_uv(width=target_width, height=target_height, dtype=np.float32) * np.array([width, height], dtype=np.float32) + target_lefttop = target_center - np.array((filter_w_f / 2, filter_h_f / 2), dtype=np.float32) + target_bottomright = target_center + np.array((filter_w_f / 2, filter_h_f / 2), dtype=np.float32) + target_window = np.floor(target_lefttop).astype(np.int32) + np.array((padding_w, padding_h), dtype=np.int32) + + target_window_centers = windowed_uv[target_window[..., 1], target_window[..., 0], :, :, :].reshape(target_height, target_width, 2, filter_size) # (target_height, tgt_width, 2, filter_size) + target_window_mask = windowed_mask[..., target_window[..., 1], target_window[..., 0], :, :].reshape(*mask.shape[:-2], target_height, target_width, filter_size) # (..., target_height, tgt_width, filter_size) + target_window_indices = windowed_indices[target_window[..., 1], target_window[..., 0], :, :].reshape(target_height, target_width, filter_size) # (target_height, tgt_width, filter_size) + + # Compute pixel area in the local windows + target_window_lefttop = np.maximum(target_window_centers - 0.5, target_lefttop[..., None]) + target_window_bottomright = np.minimum(target_window_centers + 0.5, target_bottomright[..., None]) + target_window_area = (target_window_bottomright - target_window_lefttop).clip(0, None) + target_window_area = np.where(target_window_mask, target_window_area[..., 0, :] * target_window_area[..., 1, :], 0) + + # Weighted sum by area + target_window_image = image.reshape(*image.shape[:-3], height * width, -1)[..., target_window_indices, :].swapaxes(-2, -1) + target_mask = np.sum(target_window_area, axis=-1) >= 0.25 + target_image = weighted_mean_numpy(target_window_image, target_window_area[..., None, :], axis=-1) + + if omit_channel_dim: + target_image = target_image[..., 0] + + return target_image, target_mask + + +def norm3d(x: np.ndarray) -> np.ndarray: + "Faster `np.linalg.norm(x, axis=-1)` for 3D vectors" + return np.sqrt(np.square(x[..., 0]) + np.square(x[..., 1]) + np.square(x[..., 2])) + + +def depth_occlusion_edge_numpy(depth: np.ndarray, mask: np.ndarray, kernel_size: int = 3, tol: float = 0.1): + disp = np.where(mask, 1 / depth, 0) + disp_pad = np.pad(disp, (kernel_size // 2, kernel_size // 2), constant_values=0) + mask_pad = np.pad(mask, (kernel_size // 2, kernel_size // 2), constant_values=False) + disp_window = utils3d.numpy.sliding_window_2d(disp_pad, (kernel_size, kernel_size), 1, axis=(-2, -1)) # [..., H, W, kernel_size ** 2] + mask_window = utils3d.numpy.sliding_window_2d(mask_pad, (kernel_size, kernel_size), 1, axis=(-2, -1)) # [..., H, W, kernel_size ** 2] + + disp_mean = weighted_mean_numpy(disp_window, mask_window, axis=(-2, -1)) + fg_edge_mask = mask & (disp > (1 + tol) * disp_mean) + bg_edge_mask = mask & (disp_mean > (1 + tol) * disp) + return fg_edge_mask, bg_edge_mask + + +def disk_kernel(radius: int) -> np.ndarray: + """ + Generate disk kernel with given radius. + + Args: + radius (int): Radius of the disk (in pixels). + + Returns: + np.ndarray: (2*radius+1, 2*radius+1) normalized convolution kernel. + """ + # Create coordinate grid centered at (0,0) + L = np.arange(-radius, radius + 1) + X, Y = np.meshgrid(L, L) + # Generate disk: region inside circle with radius R is 1 + kernel = ((X**2 + Y**2) <= radius**2).astype(np.float32) + # Normalize the kernel + kernel /= np.sum(kernel) + return kernel + + +def disk_blur(image: np.ndarray, radius: int) -> np.ndarray: + """ + Apply disk blur to an image using FFT convolution. + + Args: + image (np.ndarray): Input image, can be grayscale or color. + radius (int): Blur radius (in pixels). + + Returns: + np.ndarray: Blurred image. + """ + if radius == 0: + return image + kernel = disk_kernel(radius) + if image.ndim == 2: + blurred = fftconvolve(image, kernel, mode='same') + elif image.ndim == 3: + channels = [] + for i in range(image.shape[2]): + blurred_channel = fftconvolve(image[..., i], kernel, mode='same') + channels.append(blurred_channel) + blurred = np.stack(channels, axis=-1) + else: + raise ValueError("Image must be 2D or 3D.") + return blurred + + +def depth_of_field( + img: np.ndarray, + disp: np.ndarray, + focus_disp : float, + max_blur_radius : int = 10, +) -> np.ndarray: + """ + Apply depth of field effect to an image. + + Args: + img (numpy.ndarray): (H, W, 3) input image. + depth (numpy.ndarray): (H, W) depth map of the scene. + focus_depth (float): Focus depth of the lens. + strength (float): Strength of the depth of field effect. + max_blur_radius (int): Maximum blur radius (in pixels). + + Returns: + numpy.ndarray: (H, W, 3) output image with depth of field effect applied. + """ + # Precalculate dialated depth map for each blur radius + max_disp = np.max(disp) + disp = disp / max_disp + focus_disp = focus_disp / max_disp + dilated_disp = [] + for radius in range(max_blur_radius + 1): + dilated_disp.append(cv2.dilate(disp, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*radius+1, 2*radius+1)), iterations=1)) + + # Determine the blur radius for each pixel based on the depth map + blur_radii = np.clip(abs(disp - focus_disp) * max_blur_radius, 0, max_blur_radius).astype(np.int32) + for radius in range(max_blur_radius + 1): + dialted_blur_radii = np.clip(abs(dilated_disp[radius] - focus_disp) * max_blur_radius, 0, max_blur_radius).astype(np.int32) + mask = (dialted_blur_radii >= radius) & (dialted_blur_radii >= blur_radii) & (dilated_disp[radius] > disp) + blur_radii[mask] = dialted_blur_radii[mask] + blur_radii = np.clip(blur_radii, 0, max_blur_radius) + blur_radii = cv2.blur(blur_radii, (5, 5)) + + # Precalculate the blured image for each blur radius + unique_radii = np.unique(blur_radii) + precomputed = {} + for radius in range(max_blur_radius + 1): + if radius not in unique_radii: + continue + precomputed[radius] = disk_blur(img, radius) + + # Composit the blured image for each pixel + output = np.zeros_like(img) + for r in unique_radii: + mask = blur_radii == r + output[mask] = precomputed[r][mask] + + return output diff --git a/MoGe/moge/utils/geometry_torch.py b/MoGe/moge/utils/geometry_torch.py new file mode 100644 index 0000000000000000000000000000000000000000..3536ab7749ac61c3b50ccbd07d1f8e2f4077c7bc --- /dev/null +++ b/MoGe/moge/utils/geometry_torch.py @@ -0,0 +1,323 @@ +from typing import * +import math +from collections import namedtuple + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.types +import utils3d + +from .tools import timeit +from .geometry_numpy import solve_optimal_focal_shift, solve_optimal_shift + + +def weighted_mean(x: torch.Tensor, w: torch.Tensor = None, dim: Union[int, torch.Size] = None, keepdim: bool = False, eps: float = 1e-7) -> torch.Tensor: + if w is None: + return x.mean(dim=dim, keepdim=keepdim) + else: + w = w.to(x.dtype) + return (x * w).mean(dim=dim, keepdim=keepdim) / w.mean(dim=dim, keepdim=keepdim).add(eps) + + +def harmonic_mean(x: torch.Tensor, w: torch.Tensor = None, dim: Union[int, torch.Size] = None, keepdim: bool = False, eps: float = 1e-7) -> torch.Tensor: + if w is None: + return x.add(eps).reciprocal().mean(dim=dim, keepdim=keepdim).reciprocal() + else: + w = w.to(x.dtype) + return weighted_mean(x.add(eps).reciprocal(), w, dim=dim, keepdim=keepdim, eps=eps).add(eps).reciprocal() + + +def geometric_mean(x: torch.Tensor, w: torch.Tensor = None, dim: Union[int, torch.Size] = None, keepdim: bool = False, eps: float = 1e-7) -> torch.Tensor: + if w is None: + return x.add(eps).log().mean(dim=dim).exp() + else: + w = w.to(x.dtype) + return weighted_mean(x.add(eps).log(), w, dim=dim, keepdim=keepdim, eps=eps).exp() + + +def normalized_view_plane_uv(width: int, height: int, aspect_ratio: float = None, dtype: torch.dtype = None, device: torch.device = None) -> torch.Tensor: + "UV with left-top corner as (-width / diagonal, -height / diagonal) and right-bottom corner as (width / diagonal, height / diagonal)" + if aspect_ratio is None: + aspect_ratio = width / height + + span_x = aspect_ratio / (1 + aspect_ratio ** 2) ** 0.5 + span_y = 1 / (1 + aspect_ratio ** 2) ** 0.5 + + u = torch.linspace(-span_x * (width - 1) / width, span_x * (width - 1) / width, width, dtype=dtype, device=device) + v = torch.linspace(-span_y * (height - 1) / height, span_y * (height - 1) / height, height, dtype=dtype, device=device) + u, v = torch.meshgrid(u, v, indexing='xy') + uv = torch.stack([u, v], dim=-1) + return uv + + +def gaussian_blur_2d(input: torch.Tensor, kernel_size: int, sigma: float) -> torch.Tensor: + kernel = torch.exp(-(torch.arange(-kernel_size // 2 + 1, kernel_size // 2 + 1, dtype=input.dtype, device=input.device) ** 2) / (2 * sigma ** 2)) + kernel = kernel / kernel.sum() + kernel = (kernel[:, None] * kernel[None, :]).reshape(1, 1, kernel_size, kernel_size) + input = F.pad(input, (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2), mode='replicate') + input = F.conv2d(input, kernel, groups=input.shape[1]) + return input + + +def focal_to_fov(focal: torch.Tensor): + return 2 * torch.atan(0.5 / focal) + + +def fov_to_focal(fov: torch.Tensor): + return 0.5 / torch.tan(fov / 2) + + +def angle_diff_vec3(v1: torch.Tensor, v2: torch.Tensor, eps: float = 1e-12): + return torch.atan2(torch.cross(v1, v2, dim=-1).norm(dim=-1) + eps, (v1 * v2).sum(dim=-1)) + +def intrinsics_to_fov(intrinsics: torch.Tensor): + """ + Returns field of view in radians from normalized intrinsics matrix. + ### Parameters: + - intrinsics: torch.Tensor of shape (..., 3, 3) + + ### Returns: + - fov_x: torch.Tensor of shape (...) + - fov_y: torch.Tensor of shape (...) + """ + focal_x = intrinsics[..., 0, 0] + focal_y = intrinsics[..., 1, 1] + return 2 * torch.atan(0.5 / focal_x), 2 * torch.atan(0.5 / focal_y) + + +def point_map_to_depth_legacy(points: torch.Tensor): + height, width = points.shape[-3:-1] + diagonal = (height ** 2 + width ** 2) ** 0.5 + uv = normalized_view_plane_uv(width, height, dtype=points.dtype, device=points.device) # (H, W, 2) + + # Solve least squares problem + b = (uv * points[..., 2:]).flatten(-3, -1) # (..., H * W * 2) + A = torch.stack([points[..., :2], -uv.expand_as(points[..., :2])], dim=-1).flatten(-4, -2) # (..., H * W * 2, 2) + + M = A.transpose(-2, -1) @ A + solution = (torch.inverse(M + 1e-6 * torch.eye(2).to(A)) @ (A.transpose(-2, -1) @ b[..., None])).squeeze(-1) + focal, shift = solution.unbind(-1) + + depth = points[..., 2] + shift[..., None, None] + fov_x = torch.atan(width / diagonal / focal) * 2 + fov_y = torch.atan(height / diagonal / focal) * 2 + return depth, fov_x, fov_y, shift + + +def view_plane_uv_to_focal(uv: torch.Tensor): + normed_uv = normalized_view_plane_uv(width=uv.shape[-2], height=uv.shape[-3], device=uv.device, dtype=uv.dtype) + focal = (uv * normed_uv).sum() / uv.square().sum().add(1e-12) + return focal + + +def recover_focal_shift(points: torch.Tensor, mask: torch.Tensor = None, focal: torch.Tensor = None, downsample_size: Tuple[int, int] = (64, 64)): + """ + Recover the depth map and FoV from a point map with unknown z shift and focal. + + Note that it assumes: + - the optical center is at the center of the map + - the map is undistorted + - the map is isometric in the x and y directions + + ### Parameters: + - `points: torch.Tensor` of shape (..., H, W, 3) + - `downsample_size: Tuple[int, int]` in (height, width), the size of the downsampled map. Downsampling produces approximate solution and is efficient for large maps. + + ### Returns: + - `focal`: torch.Tensor of shape (...) the estimated focal length, relative to the half diagonal of the map + - `shift`: torch.Tensor of shape (...) Z-axis shift to translate the point map to camera space + """ + shape = points.shape + height, width = points.shape[-3], points.shape[-2] + diagonal = (height ** 2 + width ** 2) ** 0.5 + + points = points.reshape(-1, *shape[-3:]) + mask = None if mask is None else mask.reshape(-1, *shape[-3:-1]) + focal = focal.reshape(-1) if focal is not None else None + uv = normalized_view_plane_uv(width, height, dtype=points.dtype, device=points.device) # (H, W, 2) + + points_lr = F.interpolate(points.permute(0, 3, 1, 2), downsample_size, mode='nearest').permute(0, 2, 3, 1) + uv_lr = F.interpolate(uv.unsqueeze(0).permute(0, 3, 1, 2), downsample_size, mode='nearest').squeeze(0).permute(1, 2, 0) + mask_lr = None if mask is None else F.interpolate(mask.to(torch.float32).unsqueeze(1), downsample_size, mode='nearest').squeeze(1) > 0 + + uv_lr_np = uv_lr.cpu().numpy() + points_lr_np = points_lr.detach().cpu().numpy() + focal_np = focal.cpu().numpy() if focal is not None else None + mask_lr_np = None if mask is None else mask_lr.cpu().numpy() + optim_shift, optim_focal = [], [] + for i in range(points.shape[0]): + points_lr_i_np = points_lr_np[i] if mask is None else points_lr_np[i][mask_lr_np[i]] + uv_lr_i_np = uv_lr_np if mask is None else uv_lr_np[mask_lr_np[i]] + if uv_lr_i_np.shape[0] < 2: + optim_focal.append(1) + optim_shift.append(0) + continue + if focal is None: + optim_shift_i, optim_focal_i = solve_optimal_focal_shift(uv_lr_i_np, points_lr_i_np) + optim_focal.append(float(optim_focal_i)) + else: + optim_shift_i = solve_optimal_shift(uv_lr_i_np, points_lr_i_np, focal_np[i]) + optim_shift.append(float(optim_shift_i)) + optim_shift = torch.tensor(optim_shift, device=points.device, dtype=points.dtype).reshape(shape[:-3]) + + if focal is None: + optim_focal = torch.tensor(optim_focal, device=points.device, dtype=points.dtype).reshape(shape[:-3]) + else: + optim_focal = focal.reshape(shape[:-3]) + + return optim_focal, optim_shift + + +def mask_aware_nearest_resize( + inputs: Union[torch.Tensor, Sequence[torch.Tensor], None], + mask: torch.BoolTensor, + size: Tuple[int, int], + return_index: bool = False +) -> Tuple[Union[torch.Tensor, Sequence[torch.Tensor], None], torch.BoolTensor, Tuple[torch.LongTensor, ...]]: + """ + Resize 2D map by nearest interpolation. Return the nearest neighbor index and mask of the resized map. + + ### Parameters + - `inputs`: a single or a list of input 2D map(s) of shape (..., H, W, ...). + - `mask`: input 2D mask of shape (..., H, W) + - `size`: target size (target_width, target_height) + + ### Returns + - `*resized_maps`: resized map(s) of shape (..., target_height, target_width, ...). + - `resized_mask`: mask of the resized map of shape (..., target_height, target_width) + - `nearest_idx`: if return_index is True, nearest neighbor index of the resized map of shape (..., target_height, target_width) for each dimension, . + """ + height, width = mask.shape[-2:] + target_width, target_height = size + device = mask.device + filter_h_f, filter_w_f = max(1, height / target_height), max(1, width / target_width) + filter_h_i, filter_w_i = math.ceil(filter_h_f), math.ceil(filter_w_f) + filter_size = filter_h_i * filter_w_i + padding_h, padding_w = filter_h_i // 2 + 1, filter_w_i // 2 + 1 + + # Window the original mask and uv + uv = utils3d.torch.image_pixel_center(width=width, height=height, dtype=torch.float32, device=device) + indices = torch.arange(height * width, dtype=torch.long, device=device).reshape(height, width) + padded_uv = torch.full((height + 2 * padding_h, width + 2 * padding_w, 2), 0, dtype=torch.float32, device=device) + padded_uv[padding_h:padding_h + height, padding_w:padding_w + width] = uv + padded_mask = torch.full((*mask.shape[:-2], height + 2 * padding_h, width + 2 * padding_w), False, dtype=torch.bool, device=device) + padded_mask[..., padding_h:padding_h + height, padding_w:padding_w + width] = mask + padded_indices = torch.full((height + 2 * padding_h, width + 2 * padding_w), 0, dtype=torch.long, device=device) + padded_indices[padding_h:padding_h + height, padding_w:padding_w + width] = indices + windowed_uv = utils3d.torch.sliding_window_2d(padded_uv, (filter_h_i, filter_w_i), 1, dim=(0, 1)) + windowed_mask = utils3d.torch.sliding_window_2d(padded_mask, (filter_h_i, filter_w_i), 1, dim=(-2, -1)) + windowed_indices = utils3d.torch.sliding_window_2d(padded_indices, (filter_h_i, filter_w_i), 1, dim=(0, 1)) + + # Gather the target pixels's local window + target_uv = utils3d.torch.image_uv(width=target_width, height=target_height, dtype=torch.float32, device=device) * torch.tensor([width, height], dtype=torch.float32, device=device) + target_lefttop = target_uv - torch.tensor((filter_w_f / 2, filter_h_f / 2), dtype=torch.float32, device=device) + target_window = torch.round(target_lefttop).long() + torch.tensor((padding_w, padding_h), dtype=torch.long, device=device) + + target_window_uv = windowed_uv[target_window[..., 1], target_window[..., 0], :, :, :].reshape(target_height, target_width, 2, filter_size) # (target_height, tgt_width, 2, filter_size) + target_window_mask = windowed_mask[..., target_window[..., 1], target_window[..., 0], :, :].reshape(*mask.shape[:-2], target_height, target_width, filter_size) # (..., target_height, tgt_width, filter_size) + target_window_indices = windowed_indices[target_window[..., 1], target_window[..., 0], :, :].reshape(target_height, target_width, filter_size) # (target_height, tgt_width, filter_size) + target_window_indices = target_window_indices.expand_as(target_window_mask) + + # Compute nearest neighbor in the local window for each pixel + dist = torch.where(target_window_mask, torch.norm(target_window_uv - target_uv[..., None], dim=-2), torch.inf) # (..., target_height, tgt_width, filter_size) + nearest = torch.argmin(dist, dim=-1, keepdim=True) # (..., target_height, tgt_width, 1) + nearest_idx = torch.gather(target_window_indices, index=nearest, dim=-1).squeeze(-1) # (..., target_height, tgt_width) + target_mask = torch.any(target_window_mask, dim=-1) + nearest_i, nearest_j = nearest_idx // width, nearest_idx % width + batch_indices = [torch.arange(n, device=device).reshape([1] * i + [n] + [1] * (mask.dim() - i - 1)) for i, n in enumerate(mask.shape[:-2])] + + index = (*batch_indices, nearest_i, nearest_j) + + if inputs is None: + outputs = None + elif isinstance(inputs, torch.Tensor): + outputs = inputs[index] + elif isinstance(inputs, Sequence): + outputs = tuple(x[index] for x in inputs) + else: + raise ValueError(f'Invalid input type: {type(inputs)}') + + if return_index: + return outputs, target_mask, index + else: + return outputs, target_mask + + +def theshold_depth_change(depth: torch.Tensor, mask: torch.Tensor, pooler: Literal['min', 'max'], rtol: float = 0.2, kernel_size: int = 3): + *batch_shape, height, width = depth.shape + depth = depth.reshape(-1, 1, height, width) + mask = mask.reshape(-1, 1, height, width) + if pooler =='max': + pooled_depth = F.max_pool2d(torch.where(mask, depth, -torch.inf), kernel_size, stride=1, padding=kernel_size // 2) + output_mask = pooled_depth > depth * (1 + rtol) + elif pooler =='min': + pooled_depth = -F.max_pool2d(-torch.where(mask, depth, torch.inf), kernel_size, stride=1, padding=kernel_size // 2) + output_mask = pooled_depth < depth * (1 - rtol) + else: + raise ValueError(f'Unsupported pooler: {pooler}') + output_mask = output_mask.reshape(*batch_shape, height, width) + return output_mask + + +def depth_occlusion_edge(depth: torch.FloatTensor, mask: torch.BoolTensor, kernel_size: int = 3, tol: float = 0.1): + device, dtype = depth.device, depth.dtype + + disp = torch.where(mask, 1 / depth, 0) + disp_pad = F.pad(disp, (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2), value=0) + mask_pad = F.pad(mask, (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2), value=False) + disp_window = utils3d.torch.sliding_window_2d(disp_pad, (kernel_size, kernel_size), 1, dim=(-2, -1)).flatten(-2) # [..., H, W, kernel_size ** 2] + mask_window = utils3d.torch.sliding_window_2d(mask_pad, (kernel_size, kernel_size), 1, dim=(-2, -1)).flatten(-2) # [..., H, W, kernel_size ** 2] + + x = torch.linspace(-kernel_size // 2, kernel_size // 2, kernel_size, device=device, dtype=dtype) + A = torch.stack([*torch.meshgrid(x, x, indexing='xy'), torch.ones((kernel_size, kernel_size), device=device, dtype=dtype)], dim=-1).reshape(kernel_size ** 2, 3) # [kernel_size ** 2, 3] + A = mask_window[..., None] * A + I = torch.eye(3, device=device, dtype=dtype) + + affine_disp_window = (disp_window[..., None, :] @ A @ torch.inverse(A.mT @ A + 1e-5 * I) @ A.mT).clamp_min(1e-12)[..., 0, :] # [..., H, W, kernel_size ** 2] + diff = torch.where(mask_window, torch.maximum(affine_disp_window, disp_window) / torch.minimum(affine_disp_window, disp_window) - 1, 0) + + edge_mask = mask & (diff > tol).any(dim=-1) + + disp_mean = weighted_mean(disp_window, mask_window, dim=-1) + fg_edge_mask = edge_mask & (disp > disp_mean) + # fg_edge_mask = edge_mask & theshold_depth_change(depth, mask, pooler='max', rtol=tol, kernel_size=kernel_size) + bg_edge_mask = edge_mask & ~fg_edge_mask + return fg_edge_mask, bg_edge_mask + + +def depth_occlusion_edge(depth: torch.FloatTensor, mask: torch.BoolTensor, kernel_size: int = 3, tol: float = 0.1): + device, dtype = depth.device, depth.dtype + + disp = torch.where(mask, 1 / depth, 0) + disp_pad = F.pad(disp, (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2), value=0) + mask_pad = F.pad(mask, (kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2), value=False) + disp_window = utils3d.torch.sliding_window_2d(disp_pad, (kernel_size, kernel_size), 1, dim=(-2, -1)) # [..., H, W, kernel_size ** 2] + mask_window = utils3d.torch.sliding_window_2d(mask_pad, (kernel_size, kernel_size), 1, dim=(-2, -1)) # [..., H, W, kernel_size ** 2] + + disp_mean = weighted_mean(disp_window, mask_window, dim=(-2, -1)) + fg_edge_mask = mask & (disp / disp_mean > 1 + tol) + bg_edge_mask = mask & (disp_mean / disp > 1 + tol) + + fg_edge_mask = fg_edge_mask & F.max_pool2d(bg_edge_mask.float(), kernel_size + 2, stride=1, padding=kernel_size // 2 + 1).bool() + bg_edge_mask = bg_edge_mask & F.max_pool2d(fg_edge_mask.float(), kernel_size + 2, stride=1, padding=kernel_size // 2 + 1).bool() + + return fg_edge_mask, bg_edge_mask + + +def dilate_with_mask(input: torch.Tensor, mask: torch.BoolTensor, filter: Literal['min', 'max', 'mean', 'median'] = 'mean', iterations: int = 1) -> torch.Tensor: + kernel = torch.tensor([[False, True, False], [True, True, True], [False, True, False]], device=input.device, dtype=torch.bool) + for _ in range(iterations): + input_window = utils3d.torch.sliding_window_2d(F.pad(input, (1, 1, 1, 1), mode='constant', value=0), window_size=3, stride=1, dim=(-2, -1)) + mask_window = kernel & utils3d.torch.sliding_window_2d(F.pad(mask, (1, 1, 1, 1), mode='constant', value=False), window_size=3, stride=1, dim=(-2, -1)) + if filter =='min': + input = torch.where(mask, input, torch.where(mask_window, input_window, torch.inf).min(dim=(-2, -1)).values) + elif filter =='max': + input = torch.where(mask, input, torch.where(mask_window, input_window, -torch.inf).max(dim=(-2, -1)).values) + elif filter == 'mean': + input = torch.where(mask, input, torch.where(mask_window, input_window, torch.nan).nanmean(dim=(-2, -1))) + elif filter =='median': + input = torch.where(mask, input, torch.where(mask_window, input_window, torch.nan).flatten(-2).nanmedian(dim=-1).values) + mask = mask_window.any(dim=(-2, -1)) + return input, mask \ No newline at end of file diff --git a/MoGe/moge/utils/io.py b/MoGe/moge/utils/io.py new file mode 100644 index 0000000000000000000000000000000000000000..4c7b9fd2e9b5327833c6d75d1862248ed1650fd9 --- /dev/null +++ b/MoGe/moge/utils/io.py @@ -0,0 +1,233 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +from typing import IO +import zipfile +import json +import io +from typing import * +from pathlib import Path +import re +from PIL import Image, PngImagePlugin + +import numpy as np +import cv2 + +from .tools import timeit + + +def save_glb( + save_path: Union[str, os.PathLike], + vertices: np.ndarray, + faces: np.ndarray, + vertex_uvs: np.ndarray, + texture: np.ndarray, +): + import trimesh + import trimesh.visual + from PIL import Image + + trimesh.Trimesh( + vertices=vertices, + faces=faces, + visual = trimesh.visual.texture.TextureVisuals( + uv=vertex_uvs, + material=trimesh.visual.material.PBRMaterial( + baseColorTexture=Image.fromarray(texture), + metallicFactor=0.5, + roughnessFactor=1.0 + ) + ), + process=False + ).export(save_path) + + +def save_ply( + save_path: Union[str, os.PathLike], + vertices: np.ndarray, + faces: np.ndarray, + vertex_colors: np.ndarray, +): + import trimesh + import trimesh.visual + from PIL import Image + + trimesh.Trimesh( + vertices=vertices, + faces=faces, + vertex_colors=vertex_colors, + process=False + ).export(save_path) + + + +def read_image(path: Union[str, os.PathLike, IO]) -> np.ndarray: + """ + Read a image, return uint8 RGB array of shape (H, W, 3). + """ + if isinstance(path, (str, os.PathLike)): + data = Path(path).read_bytes() + else: + data = path.read() + image = cv2.cvtColor(cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) + return image + + +def write_image(path: Union[str, os.PathLike, IO], image: np.ndarray, quality: int = 95): + """ + Write a image, input uint8 RGB array of shape (H, W, 3). + """ + data = cv2.imencode('.jpg', cv2.cvtColor(image, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, quality])[1].tobytes() + if isinstance(path, (str, os.PathLike)): + Path(path).write_bytes(data) + else: + path.write(data) + + +def read_depth(path: Union[str, os.PathLike, IO]) -> Tuple[np.ndarray, float]: + """ + Read a depth image, return float32 depth array of shape (H, W). + """ + if isinstance(path, (str, os.PathLike)): + data = Path(path).read_bytes() + else: + data = path.read() + pil_image = Image.open(io.BytesIO(data)) + near = float(pil_image.info.get('near')) + far = float(pil_image.info.get('far')) + unit = float(pil_image.info.get('unit')) if 'unit' in pil_image.info else None + depth = np.array(pil_image) + mask_nan, mask_inf = depth == 0, depth == 65535 + depth = (depth.astype(np.float32) - 1) / 65533 + depth = near ** (1 - depth) * far ** depth + depth[mask_nan] = np.nan + depth[mask_inf] = np.inf + return depth, unit + + +def write_depth( + path: Union[str, os.PathLike, IO], + depth: np.ndarray, + unit: float = None, + max_range: float = 1e5, + compression_level: int = 7, +): + """ + Encode and write a depth image as 16-bit PNG format. + ### Parameters: + - `path: Union[str, os.PathLike, IO]` + The file path or file object to write to. + - `depth: np.ndarray` + The depth array, float32 array of shape (H, W). + May contain `NaN` for invalid values and `Inf` for infinite values. + - `unit: float = None` + The unit of the depth values. + + Depth values are encoded as follows: + - 0: unknown + - 1 ~ 65534: depth values in logarithmic + - 65535: infinity + + metadata is stored in the PNG file as text fields: + - `near`: the minimum depth value + - `far`: the maximum depth value + - `unit`: the unit of the depth values (optional) + """ + mask_values, mask_nan, mask_inf = np.isfinite(depth), np.isnan(depth),np.isinf(depth) + + depth = depth.astype(np.float32) + mask_finite = depth + near = max(depth[mask_values].min(), 1e-5) + far = max(near * 1.1, min(depth[mask_values].max(), near * max_range)) + depth = 1 + np.round((np.log(np.nan_to_num(depth, nan=0).clip(near, far) / near) / np.log(far / near)).clip(0, 1) * 65533).astype(np.uint16) # 1~65534 + depth[mask_nan] = 0 + depth[mask_inf] = 65535 + + pil_image = Image.fromarray(depth) + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text('near', str(near)) + pnginfo.add_text('far', str(far)) + if unit is not None: + pnginfo.add_text('unit', str(unit)) + pil_image.save(path, pnginfo=pnginfo, compress_level=compression_level) + + +def read_segmentation(path: Union[str, os.PathLike, IO]) -> Tuple[np.ndarray, Dict[str, int]]: + """ + Read a segmentation mask + ### Parameters: + - `path: Union[str, os.PathLike, IO]` + The file path or file object to read from. + ### Returns: + - `Tuple[np.ndarray, Dict[str, int]]` + A tuple containing: + - `mask`: uint8 or uint16 numpy.ndarray of shape (H, W). + - `labels`: Dict[str, int]. The label mapping, a dictionary of {label_name: label_id}. + """ + if isinstance(path, (str, os.PathLike)): + data = Path(path).read_bytes() + else: + data = path.read() + pil_image = Image.open(io.BytesIO(data)) + labels = json.loads(pil_image.info['labels']) if 'labels' in pil_image.info else None + mask = np.array(pil_image) + return mask, labels + + +def write_segmentation(path: Union[str, os.PathLike, IO], mask: np.ndarray, labels: Dict[str, int] = None, compression_level: int = 7): + """ + Write a segmentation mask and label mapping, as PNG format. + ### Parameters: + - `path: Union[str, os.PathLike, IO]` + The file path or file object to write to. + - `mask: np.ndarray` + The segmentation mask, uint8 or uint16 array of shape (H, W). + - `labels: Dict[str, int] = None` + The label mapping, a dictionary of {label_name: label_id}. + - `compression_level: int = 7` + The compression level for PNG compression. + """ + assert mask.dtype == np.uint8 or mask.dtype == np.uint16, f"Unsupported dtype {mask.dtype}" + pil_image = Image.fromarray(mask) + pnginfo = PngImagePlugin.PngInfo() + if labels is not None: + labels_json = json.dumps(labels, ensure_ascii=True, separators=(',', ':')) + pnginfo.add_text('labels', labels_json) + pil_image.save(path, pnginfo=pnginfo, compress_level=compression_level) + + + +def read_normal(path: Union[str, os.PathLike, IO]) -> np.ndarray: + """ + Read a normal image, return float32 normal array of shape (H, W, 3). + """ + if isinstance(path, (str, os.PathLike)): + data = Path(path).read_bytes() + else: + data = path.read() + normal = cv2.cvtColor(cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB) + mask_nan = np.all(normal == 0, axis=-1) + normal = (normal.astype(np.float32) / 65535 - 0.5) * [2.0, -2.0, -2.0] + normal = normal / (np.sqrt(np.square(normal[..., 0]) + np.square(normal[..., 1]) + np.square(normal[..., 2])) + 1e-12) + normal[mask_nan] = np.nan + return normal + + +def write_normal(path: Union[str, os.PathLike, IO], normal: np.ndarray, compression_level: int = 7) -> np.ndarray: + """ + Write a normal image, input float32 normal array of shape (H, W, 3). + """ + mask_nan = np.isnan(normal).any(axis=-1) + normal = ((normal * [0.5, -0.5, -0.5] + 0.5).clip(0, 1) * 65535).astype(np.uint16) + normal[mask_nan] = 0 + data = cv2.imencode('.png', cv2.cvtColor(normal, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_PNG_COMPRESSION, compression_level])[1].tobytes() + if isinstance(path, (str, os.PathLike)): + Path(path).write_bytes(data) + else: + path.write(data) + + +def read_meta(path: Union[str, os.PathLike, IO]) -> Dict[str, Any]: + return json.loads(Path(path).read_text()) + +def write_meta(path: Union[str, os.PathLike, IO], meta: Dict[str, Any]): + Path(path).write_text(json.dumps(meta)) \ No newline at end of file diff --git a/MoGe/moge/utils/panorama.py b/MoGe/moge/utils/panorama.py new file mode 100644 index 0000000000000000000000000000000000000000..3f9d121c3c189770a7fd9f88be66f74f1ba5cfd3 --- /dev/null +++ b/MoGe/moge/utils/panorama.py @@ -0,0 +1,191 @@ +import os +os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1' +from pathlib import Path +from typing import * +import itertools +import json +import warnings + +import cv2 +import numpy as np +from numpy import ndarray +from tqdm import tqdm, trange +from scipy.sparse import csr_array, hstack, vstack +from scipy.ndimage import convolve +from scipy.sparse.linalg import lsmr + +import utils3d + + +def get_panorama_cameras(): + vertices, _ = utils3d.numpy.icosahedron() + intrinsics = utils3d.numpy.intrinsics_from_fov(fov_x=np.deg2rad(90), fov_y=np.deg2rad(90)) + extrinsics = utils3d.numpy.extrinsics_look_at([0, 0, 0], vertices, [0, 0, 1]).astype(np.float32) + return extrinsics, [intrinsics] * len(vertices) + + +def spherical_uv_to_directions(uv: np.ndarray): + theta, phi = (1 - uv[..., 0]) * (2 * np.pi), uv[..., 1] * np.pi + directions = np.stack([np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)], axis=-1) + return directions + + +def directions_to_spherical_uv(directions: np.ndarray): + directions = directions / np.linalg.norm(directions, axis=-1, keepdims=True) + u = 1 - np.arctan2(directions[..., 1], directions[..., 0]) / (2 * np.pi) % 1.0 + v = np.arccos(directions[..., 2]) / np.pi + return np.stack([u, v], axis=-1) + + +def split_panorama_image(image: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray, resolution: int): + height, width = image.shape[:2] + uv = utils3d.numpy.image_uv(width=resolution, height=resolution) + splitted_images = [] + for i in range(len(extrinsics)): + spherical_uv = directions_to_spherical_uv(utils3d.numpy.unproject_cv(uv, extrinsics=extrinsics[i], intrinsics=intrinsics[i])) + pixels = utils3d.numpy.uv_to_pixel(spherical_uv, width=width, height=height).astype(np.float32) + + splitted_image = cv2.remap(image, pixels[..., 0], pixels[..., 1], interpolation=cv2.INTER_LINEAR) + splitted_images.append(splitted_image) + return splitted_images + + +def poisson_equation(width: int, height: int, wrap_x: bool = False, wrap_y: bool = False) -> Tuple[csr_array, ndarray]: + grid_index = np.arange(height * width).reshape(height, width) + grid_index = np.pad(grid_index, ((0, 0), (1, 1)), mode='wrap' if wrap_x else 'edge') + grid_index = np.pad(grid_index, ((1, 1), (0, 0)), mode='wrap' if wrap_y else 'edge') + + data = np.array([[-4, 1, 1, 1, 1]], dtype=np.float32).repeat(height * width, axis=0).reshape(-1) + indices = np.stack([ + grid_index[1:-1, 1:-1], + grid_index[:-2, 1:-1], # up + grid_index[2:, 1:-1], # down + grid_index[1:-1, :-2], # left + grid_index[1:-1, 2:] # right + ], axis=-1).reshape(-1) + indptr = np.arange(0, height * width * 5 + 1, 5) + A = csr_array((data, indices, indptr), shape=(height * width, height * width)) + + return A + + +def grad_equation(width: int, height: int, wrap_x: bool = False, wrap_y: bool = False) -> Tuple[csr_array, np.ndarray]: + grid_index = np.arange(width * height).reshape(height, width) + if wrap_x: + grid_index = np.pad(grid_index, ((0, 0), (0, 1)), mode='wrap') + if wrap_y: + grid_index = np.pad(grid_index, ((0, 1), (0, 0)), mode='wrap') + + data = np.concatenate([ + np.concatenate([ + np.ones((grid_index.shape[0], grid_index.shape[1] - 1), dtype=np.float32).reshape(-1, 1), # x[i,j] + -np.ones((grid_index.shape[0], grid_index.shape[1] - 1), dtype=np.float32).reshape(-1, 1), # x[i,j-1] + ], axis=1).reshape(-1), + np.concatenate([ + np.ones((grid_index.shape[0] - 1, grid_index.shape[1]), dtype=np.float32).reshape(-1, 1), # x[i,j] + -np.ones((grid_index.shape[0] - 1, grid_index.shape[1]), dtype=np.float32).reshape(-1, 1), # x[i-1,j] + ], axis=1).reshape(-1), + ]) + indices = np.concatenate([ + np.concatenate([ + grid_index[:, :-1].reshape(-1, 1), + grid_index[:, 1:].reshape(-1, 1), + ], axis=1).reshape(-1), + np.concatenate([ + grid_index[:-1, :].reshape(-1, 1), + grid_index[1:, :].reshape(-1, 1), + ], axis=1).reshape(-1), + ]) + indptr = np.arange(0, grid_index.shape[0] * (grid_index.shape[1] - 1) * 2 + (grid_index.shape[0] - 1) * grid_index.shape[1] * 2 + 1, 2) + A = csr_array((data, indices, indptr), shape=(grid_index.shape[0] * (grid_index.shape[1] - 1) + (grid_index.shape[0] - 1) * grid_index.shape[1], height * width)) + + return A + + +def merge_panorama_depth(width: int, height: int, distance_maps: List[np.ndarray], pred_masks: List[np.ndarray], extrinsics: List[np.ndarray], intrinsics: List[np.ndarray]): + if max(width, height) > 256: + panorama_depth_init, _ = merge_panorama_depth(width // 2, height // 2, distance_maps, pred_masks, extrinsics, intrinsics) + panorama_depth_init = cv2.resize(panorama_depth_init, (width, height), cv2.INTER_LINEAR) + else: + panorama_depth_init = None + + uv = utils3d.numpy.image_uv(width=width, height=height) + spherical_directions = spherical_uv_to_directions(uv) + + # Warp each view to the panorama + panorama_log_distance_grad_maps, panorama_grad_masks = [], [] + panorama_log_distance_laplacian_maps, panorama_laplacian_masks = [], [] + panorama_pred_masks = [] + for i in range(len(distance_maps)): + projected_uv, projected_depth = utils3d.numpy.project_cv(spherical_directions, extrinsics=extrinsics[i], intrinsics=intrinsics[i]) + projection_valid_mask = (projected_depth > 0) & (projected_uv > 0).all(axis=-1) & (projected_uv < 1).all(axis=-1) + + projected_pixels = utils3d.numpy.uv_to_pixel(np.clip(projected_uv, 0, 1), width=distance_maps[i].shape[1], height=distance_maps[i].shape[0]).astype(np.float32) + + log_splitted_distance = np.log(distance_maps[i]) + panorama_log_distance_map = np.where(projection_valid_mask, cv2.remap(log_splitted_distance, projected_pixels[..., 0], projected_pixels[..., 1], cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE), 0) + panorama_pred_mask = projection_valid_mask & (cv2.remap(pred_masks[i].astype(np.uint8), projected_pixels[..., 0], projected_pixels[..., 1], cv2.INTER_NEAREST, borderMode=cv2.BORDER_REPLICATE) > 0) + + # calculate gradient map + padded = np.pad(panorama_log_distance_map, ((0, 0), (0, 1)), mode='wrap') + grad_x, grad_y = padded[:, :-1] - padded[:, 1:], padded[:-1, :] - padded[1:, :] + + padded = np.pad(panorama_pred_mask, ((0, 0), (0, 1)), mode='wrap') + mask_x, mask_y = padded[:, :-1] & padded[:, 1:], padded[:-1, :] & padded[1:, :] + + panorama_log_distance_grad_maps.append((grad_x, grad_y)) + panorama_grad_masks.append((mask_x, mask_y)) + + # calculate laplacian map + padded = np.pad(panorama_log_distance_map, ((1, 1), (0, 0)), mode='edge') + padded = np.pad(padded, ((0, 0), (1, 1)), mode='wrap') + laplacian = convolve(padded, np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=np.float32))[1:-1, 1:-1] + + padded = np.pad(panorama_pred_mask, ((1, 1), (0, 0)), mode='edge') + padded = np.pad(padded, ((0, 0), (1, 1)), mode='wrap') + mask = convolve(padded.astype(np.uint8), np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.uint8))[1:-1, 1:-1] == 5 + + panorama_log_distance_laplacian_maps.append(laplacian) + panorama_laplacian_masks.append(mask) + + panorama_pred_masks.append(panorama_pred_mask) + + panorama_log_distance_grad_x = np.stack([grad_map[0] for grad_map in panorama_log_distance_grad_maps], axis=0) + panorama_log_distance_grad_y = np.stack([grad_map[1] for grad_map in panorama_log_distance_grad_maps], axis=0) + panorama_grad_mask_x = np.stack([mask_map[0] for mask_map in panorama_grad_masks], axis=0) + panorama_grad_mask_y = np.stack([mask_map[1] for mask_map in panorama_grad_masks], axis=0) + + panorama_log_distance_grad_x = np.sum(panorama_log_distance_grad_x * panorama_grad_mask_x, axis=0) / np.sum(panorama_grad_mask_x, axis=0).clip(1e-3) + panorama_log_distance_grad_y = np.sum(panorama_log_distance_grad_y * panorama_grad_mask_y, axis=0) / np.sum(panorama_grad_mask_y, axis=0).clip(1e-3) + + panorama_laplacian_maps = np.stack(panorama_log_distance_laplacian_maps, axis=0) + panorama_laplacian_masks = np.stack(panorama_laplacian_masks, axis=0) + panorama_laplacian_map = np.sum(panorama_laplacian_maps * panorama_laplacian_masks, axis=0) / np.sum(panorama_laplacian_masks, axis=0).clip(1e-3) + + grad_x_mask = np.any(panorama_grad_mask_x, axis=0).reshape(-1) + grad_y_mask = np.any(panorama_grad_mask_y, axis=0).reshape(-1) + grad_mask = np.concatenate([grad_x_mask, grad_y_mask]) + laplacian_mask = np.any(panorama_laplacian_masks, axis=0).reshape(-1) + + # Solve overdetermined system + A = vstack([ + grad_equation(width, height, wrap_x=True, wrap_y=False)[grad_mask], + poisson_equation(width, height, wrap_x=True, wrap_y=False)[laplacian_mask], + ]) + b = np.concatenate([ + panorama_log_distance_grad_x.reshape(-1)[grad_x_mask], + panorama_log_distance_grad_y.reshape(-1)[grad_y_mask], + panorama_laplacian_map.reshape(-1)[laplacian_mask] + ]) + x, *_ = lsmr( + A, b, + atol=1e-5, btol=1e-5, + x0=np.log(panorama_depth_init).reshape(-1) if panorama_depth_init is not None else None, + show=False, + ) + + panorama_depth = np.exp(x).reshape(height, width).astype(np.float32) + panorama_mask = np.any(panorama_pred_masks, axis=0) + + return panorama_depth, panorama_mask + diff --git a/MoGe/moge/utils/pipeline.py b/MoGe/moge/utils/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..daa522e987317e949899d4159e61d7a7066e1fba --- /dev/null +++ b/MoGe/moge/utils/pipeline.py @@ -0,0 +1,503 @@ +from typing import * +from abc import abstractmethod +from queue import Empty, Full +from threading import Thread +from queue import Queue +from multiprocessing import Process +from threading import Thread, Event +import multiprocessing +import threading +import inspect +import time +import uuid +from copy import deepcopy +import itertools +import functools + +__all__ = [ + 'Node', + 'Link', + 'ConcurrentNode', + 'Worker', + 'WorkerFunction', + 'Provider', + 'ProviderFunction', + 'Sequential', + 'Batch', + 'Unbatch', + 'Parallel', + 'Graph', + 'Buffer', +] + +TERMINATE_CHECK_INTERVAL = 0.5 + + +class _ItemWrapper: + def __init__(self, data: Any, id: Union[int, List[int]] = None): + self.data = data + self.id = id + + +class Terminate(Exception): + pass + + +def _get_queue_item(queue: Queue, terminate_flag: Event, timeout: float = None) -> _ItemWrapper: + while True: + try: + item: _ItemWrapper = queue.get(block=True, timeout=TERMINATE_CHECK_INTERVAL if timeout is None else min(timeout, TERMINATE_CHECK_INTERVAL)) + if terminate_flag.is_set(): + raise Terminate() + return item + except Empty: + if terminate_flag.is_set(): + raise Terminate() + + if timeout is not None: + timeout -= TERMINATE_CHECK_INTERVAL + if timeout <= 0: + raise Empty() + + +def _put_queue_item(queue: Queue, item: _ItemWrapper, terminate_flag: Event): + while True: + try: + queue.put(item, block=True, timeout=TERMINATE_CHECK_INTERVAL) + if terminate_flag.is_set(): + raise Terminate() + return + except Full: + if terminate_flag.is_set(): + raise Terminate() + +class Node: + def __init__(self, in_buffer_size: int = 1, out_buffer_size: int = 1) -> None: + self.input: Queue = Queue(maxsize=in_buffer_size) + self.output: Queue = Queue(maxsize=out_buffer_size) + self.in_buffer_size = in_buffer_size + self.out_buffer_size = out_buffer_size + + @abstractmethod + def start(self): + pass + + @abstractmethod + def terminate(self): + pass + + def stop(self): + self.terminate() + self.join() + + @abstractmethod + def join(self): + pass + + def put(self, data: Any, key: str = None, block: bool = True) -> None: + item = _ItemWrapper(data) + self.input.put(item, block=block) + + def get(self, key: str = None, block: bool = True) -> Any: + item: _ItemWrapper = self.output.get(block=block) + return item.data + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.terminate() + self.join() + + +class ConcurrentNode(Node): + job: Union[Thread, Process] + + def __init__(self, running_as: Literal['thread', 'process'] = 'thread', in_buffer_size: int = 1, out_buffer_size: int = 1) -> None: + super().__init__(in_buffer_size, out_buffer_size) + self.running_as = running_as + + @abstractmethod + def _loop_fn(self, input: Queue, output: Queue, terminate_flag: Event): + pass + + def start(self): + if self.running_as == 'thread': + terminate_flag = threading.Event() + job = Thread(target=self._loop_fn, args=(self.input, self.output, terminate_flag)) + elif self.running_as == 'process': + terminate_flag = multiprocessing.Event() + job = Process(target=self._loop_fn, args=(self.input, self.output, terminate_flag)) + job.start() + self.job = job + self.terminate_flag = terminate_flag + + def terminate(self): + self.terminate_flag.set() + + def join(self): + self.job.join() + + +class Worker(ConcurrentNode): + def __init__(self, running_as: Literal['thread', 'process'] = 'thread', in_buffer_size: int = 0, out_buffer_size: int = 0) -> None: + super().__init__(running_as, in_buffer_size, out_buffer_size) + + def init(self) -> None: + """ + This method is called the the thread is started, to initialize any resources that is only held in the thread. + """ + pass + + @abstractmethod + def work(self, *args, **kwargs) -> Union[Any, Dict[str, Any]]: + """ + This method defines the job that the node should do for each input item. + A item obtained from the input queue is passed as arguments to this method, and the result is placed in the output queue. + The method is executed concurrently with other nodes. + """ + pass + + def _loop_fn(self, input: Queue, output: Queue, terminate_flag: Event): + self.init() + try: + while True: + item = _get_queue_item(input, terminate_flag) + result = self.work(item.data) + _put_queue_item(output, _ItemWrapper(result, item.id), terminate_flag) + + except Terminate: + return + + +class Provider(ConcurrentNode): + """ + A node that provides data to successive nodes. It takes no input and provides data to the output queue. + """ + def __init__(self, running_as: Literal['thread', 'process'], out_buffer_size: int = 1) -> None: + super().__init__(running_as, 0, out_buffer_size) + + def init(self) -> None: + """ + This method is called the the thread or process is started, to initialize any resources that is only held in the thread or process. + """ + pass + + @abstractmethod + def provide(self) -> Generator[Any, None, None]: + pass + + def _loop_fn(self, input: Queue, output: Queue, terminate_flag: Event): + self.init() + try: + for data in self.provide(): + _put_queue_item(output, _ItemWrapper(data), terminate_flag) + except Terminate: + return + + +class WorkerFunction(Worker): + def __init__(self, fn: Callable, running_as: 'thread', in_buffer_size: int = 1, out_buffer_size: int = 1) -> None: + super().__init__(running_as, in_buffer_size, out_buffer_size) + self.fn = fn + + def work(self, *args, **kwargs): + return self.fn(*args, **kwargs) + + +class ProviderFunction(Provider): + def __init__(self, fn: Callable, running_as: 'thread', out_buffer_size: int = 1) -> None: + super().__init__(running_as, out_buffer_size) + self.fn = fn + + def provide(self): + for item in self.fn(): + yield item + + +class Link: + def __init__(self, src: Queue, dst: Queue): + self.src = src + self.dst = dst + + def _thread_fn(self): + try: + while True: + item = _get_queue_item(self.src, self.terminate_flag) + _put_queue_item(self.dst, item, self.terminate_flag) + except Terminate: + return + + def start(self): + self.terminate_flag = threading.Event() + self.thread = Thread(target=self._thread_fn) + self.thread.start() + + def terminate(self): + self.terminate_flag.set() + + def join(self): + self.thread.join() + + +class Graph(Node): + """ + Graph pipeline of nodes and links + """ + nodes: List[Node] + links: List[Link] + + def __init__(self, in_buffer_size: int = 1, out_buffer_size: int = 1): + super().__init__(in_buffer_size, out_buffer_size) + self.nodes = [] + self.links = [] + + def add(self, node: Node): + self.nodes.append(node) + + def link(self, src: Union[Node, Tuple[Node, str]], dst: Union[Node, Tuple[Node, str]]): + """ + Links the output of the source node to the input of the destination node. + If the source or destination node is None, the pipeline's input or output is used. + """ + src_queue = self.input if src is None else src.output + dst_queue = self.output if dst is None else dst.input + self.links.append(Link(src_queue, dst_queue)) + + def chain(self, nodes: Iterable[Node]): + """ + Link the output of each node to the input of the next node. + """ + nodes = list(nodes) + for i in range(len(nodes) - 1): + self.link(nodes[i], nodes[i + 1]) + + def start(self): + for node in self.nodes: + node.start() + for link in self.links: + link.start() + + def terminate(self): + for node in self.nodes: + node.terminate() + for link in self.links: + link.terminate() + + def join(self): + for node in self.nodes: + node.join() + for link in self.links: + link.join() + + def __iter__(self): + providers = [node for node in self.nodes if isinstance(node, Provider)] + if len(providers) == 0: + raise ValueError("No provider node found in the pipeline. If you want to iterate over the pipeline, the pipeline must be driven by a provider node.") + with self: + # while all(provider.job.is_alive() for provider in providers): + while True: + yield self.get() + + def __call__(self, data: Any) -> Any: + """ + Submit data to the pipeline's input queue, and return the output data asynchronously. + NOTE: The pipeline must be streamed (i.e., every output item is uniquely associated with an input item) for this to work. + """ + # TODO + + +class Sequential(Graph): + """ + Pipeline of nodes in sequential order, where each node takes the output of the previous node as input. + The order of input and output items is preserved (FIFO) + """ + def __init__(self, nodes: List[Union[Node, Callable]], function_running_as: Literal['thread', 'process'] = 'thread', in_buffer_size: int = 1, out_buffer_size: int = 1): + """ + Initialize the pipeline with a list of nodes to execute sequentially. + ### Parameters: + - nodes: List of nodes or functions to execute sequentially. Generator functions are wrapped in provider nodes, and other functions are wrapped in worker nodes. + - function_running_as: Whether to wrap the function as a thread or process worker. Defaults to 'thread'. + - in_buffer_size: Maximum size of the input queue of the pipeline. Defaults to 0 (unlimited). + - out_buffer_size: Maximum size of the output queue of the pipeline. Defaults to 0 (unlimited). + """ + super().__init__(in_buffer_size, out_buffer_size) + for node in nodes: + if isinstance(node, Node): + pass + elif isinstance(node, Callable): + if inspect.isgeneratorfunction(node): + node = ProviderFunction(node, function_running_as) + else: + node = WorkerFunction(node, function_running_as) + else: + raise ValueError(f"Invalid node type: {type(node)}") + self.add(node) + self.chain([None, *self.nodes, None]) + + +class Parallel(Node): + """ + A FIFO node that runs multiple nodes in parallel to process the input items. Each input item is handed to one of the nodes whoever is available. + NOTE: It is FIFO if and only if all the nested nodes are FIFO. + """ + nodes: List[Node] + + def __init__(self, nodes: Iterable[Node], in_buffer_size: int = 1, out_buffer_size: int = 1, function_running_as: Literal['thread', 'process'] = 'thread'): + super().__init__(in_buffer_size, out_buffer_size) + self.nodes = [] + for node in nodes: + if isinstance(node, Node): + pass + elif isinstance(node, Callable): + if inspect.isgeneratorfunction(node): + node = ProviderFunction(node, function_running_as) + else: + node = WorkerFunction(node, function_running_as) + else: + raise ValueError(f"Invalid node type: {type(node)}") + self.nodes.append(node) + self.output_order = Queue() + self.lock = threading.Lock() + + def _in_thread_fn(self, node: Node): + try: + while True: + with self.lock: + # A better idea: first make sure its node is vacant, then get it a new item. + # Currently we will not be able to know which node is busy util there is at least one item already waiting in the queue of the node. + # This could lead to suboptimal scheduling. + item = _get_queue_item(self.input, self.terminate_flag) + self.output_order.put(node.output) + _put_queue_item(node.input, item, self.terminate_flag) + except Terminate: + return + + def _out_thread_fn(self): + try: + while True: + queue = _get_queue_item(self.output_order, self.terminate_flag) + item = _get_queue_item(queue, self.terminate_flag) + _put_queue_item(self.output, item, self.terminate_flag) + except Terminate: + return + + def start(self): + self.terminate_flag = threading.Event() + self.in_threads = [] + for node in self.nodes: + thread = Thread(target=self._in_thread_fn, args=(node,)) + thread.start() + self.in_threads.append(thread) + thread = Thread(target=self._out_thread_fn) + thread.start() + self.out_thread = thread + for node in self.nodes: + node.start() + + def terminate(self): + self.terminate_flag.set() + for node in self.nodes: + node.terminate() + + def join(self): + for thread in self.in_threads: + thread.join() + self.out_thread.join() + + +class UnorderedParallel(Graph): + """ + Pipeline of nodes in parallel, where each input item is handed to one of the nodes whoever is available. + NOTE: The order of the output items is NOT guaranteed to be the same as the input items, depending on how fast the nodes handle their input. + """ + def __init__(self, nodes: List[Union[Node, Callable]], function_running_as: Literal['thread', 'process'] = 'thread', in_buffer_size: int = 1, out_buffer_size: int = 1): + """ + Initialize the pipeline with a list of nodes to execute in parallel. If a function is given, it is wrapped in a worker node. + ### Parameters: + - nodes: List of nodes or functions to execute in parallel. Generator functions are wrapped in provider nodes, and other functions are wrapped in worker nodes. + - function_running_as: Whether to wrap the function as a thread or process worker. Defaults to 'thread'. + - in_buffer_size: Maximum size of the input queue of the pipeline. Defaults to 0 (unlimited). + - out_buffer_size: Maximum size of the output queue of the pipeline. Defaults to 0 (unlimited). + """ + super().__init__(in_buffer_size, out_buffer_size) + for node in nodes: + if isinstance(node, Node): + pass + elif isinstance(node, Callable): + if inspect.isgeneratorfunction(node): + node = ProviderFunction(node, function_running_as) + else: + node = WorkerFunction(node, function_running_as) + else: + raise ValueError(f"Invalid node type: {type(node)}") + self.add(node) + for i in range(len(nodes)): + self.chain([None, self.nodes[i], None]) + + +class Batch(ConcurrentNode): + """ + Groups every `batch_size` items into a batch (a list of items) and passes the batch to successive nodes. + The `patience` parameter specifies the maximum time to wait for a batch to be filled before sending it to the next node, + i.e., when the earliest item in the batch is out of `patience` seconds, the batch is sent regardless of its size. + """ + def __init__(self, batch_size: int, patience: float = None, in_buffer_size: int = 1, out_buffer_size: int = 1): + assert batch_size > 0, "Batch size must be greater than 0." + super().__init__('thread', in_buffer_size, out_buffer_size) + self.batch_size = batch_size + self.patience = patience + + def _loop_fn(self, input: Queue, output: Queue, terminate_flag: Event): + try: + while True: + batch_id, batch_data = [], [] + # Try to fill the batch + for i in range(self.batch_size): + if i == 0 or self.patience is None: + timeout = None + else: + timeout = self.patience - (time.time() - earliest_time) + if timeout < 0: + break + try: + item = _get_queue_item(input, terminate_flag, timeout) + except Empty: + break + + if i == 0: + earliest_time = time.time() + batch_data.append(item.data) + batch_id.append(item.id) + + batch = _ItemWrapper(batch_data, batch_id) + _put_queue_item(output, batch, terminate_flag) + except Terminate: + return + + +class Unbatch(ConcurrentNode): + """ + Ungroups every batch (a list of items) into individual items and passes them to successive nodes. + """ + def __init__(self, in_buffer_size: int = 1, out_buffer_size: int = 1): + super().__init__('thread', in_buffer_size, out_buffer_size) + + def _loop_fn(self, input: Queue, output: Queue, terminate_flag: Event): + try: + while True: + batch = _get_queue_item(input, terminate_flag) + for id, data in zip(batch.id or itertools.repeat(None), batch.data): + item = _ItemWrapper(data, id) + _put_queue_item(output, item, terminate_flag) + except Terminate: + return + + +class Buffer(Node): + "A FIFO node that buffers items in a queue. Usefull achieve better temporal balance when its successor node has a variable processing time." + def __init__(self, size: int): + super().__init__(size, size) + self.size = size + self.input = self.output = Queue(maxsize=size) \ No newline at end of file diff --git a/MoGe/moge/utils/tools.py b/MoGe/moge/utils/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..3687f6938fe34433d149a1a8405be7eed5f23c37 --- /dev/null +++ b/MoGe/moge/utils/tools.py @@ -0,0 +1,289 @@ +from typing import * +import time +from pathlib import Path +from numbers import Number +from functools import wraps +import warnings +import math +import json +import os +import importlib +import importlib.util + + +def catch_exception(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except Exception as e: + import traceback + print(f"Exception in {fn.__name__}", end='r') + # print({', '.join(repr(arg) for arg in args)}, {', '.join(f'{k}={v!r}' for k, v in kwargs.items())}) + traceback.print_exc(chain=False) + time.sleep(0.1) + return None + return wrapper + + +class CallbackOnException: + def __init__(self, callback: Callable, exception: type): + self.exception = exception + self.callback = callback + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if isinstance(exc_val, self.exception): + self.callback() + return True + return False + +def traverse_nested_dict_keys(d: Dict[str, Dict]) -> Generator[Tuple[str, ...], None, None]: + for k, v in d.items(): + if isinstance(v, dict): + for sub_key in traverse_nested_dict_keys(v): + yield (k, ) + sub_key + else: + yield (k, ) + + +def get_nested_dict(d: Dict[str, Dict], keys: Tuple[str, ...], default: Any = None): + for k in keys: + d = d.get(k, default) + if d is None: + break + return d + +def set_nested_dict(d: Dict[str, Dict], keys: Tuple[str, ...], value: Any): + for k in keys[:-1]: + d = d.setdefault(k, {}) + d[keys[-1]] = value + + +def key_average(list_of_dicts: list) -> Dict[str, Any]: + """ + Returns a dictionary with the average value of each key in the input list of dictionaries. + """ + _nested_dict_keys = set() + for d in list_of_dicts: + _nested_dict_keys.update(traverse_nested_dict_keys(d)) + _nested_dict_keys = sorted(_nested_dict_keys) + result = {} + for k in _nested_dict_keys: + values = [] + for d in list_of_dicts: + v = get_nested_dict(d, k) + if v is not None and not math.isnan(v): + values.append(v) + avg = sum(values) / len(values) if values else float('nan') + set_nested_dict(result, k, avg) + return result + + +def flatten_nested_dict(d: Dict[str, Any], parent_key: Tuple[str, ...] = None) -> Dict[Tuple[str, ...], Any]: + """ + Flattens a nested dictionary into a single-level dictionary, with keys as tuples. + """ + items = [] + if parent_key is None: + parent_key = () + for k, v in d.items(): + new_key = parent_key + (k, ) + if isinstance(v, MutableMapping): + items.extend(flatten_nested_dict(v, new_key).items()) + else: + items.append((new_key, v)) + return dict(items) + + +def unflatten_nested_dict(d: Dict[str, Any]) -> Dict[str, Any]: + """ + Unflattens a single-level dictionary into a nested dictionary, with keys as tuples. + """ + result = {} + for k, v in d.items(): + sub_dict = result + for k_ in k[:-1]: + if k_ not in sub_dict: + sub_dict[k_] = {} + sub_dict = sub_dict[k_] + sub_dict[k[-1]] = v + return result + + +def read_jsonl(file): + import json + with open(file, 'r') as f: + data = f.readlines() + return [json.loads(line) for line in data] + + +def write_jsonl(data: List[dict], file): + import json + with open(file, 'w') as f: + for item in data: + f.write(json.dumps(item) + '\n') + + +def to_hierachical_dataframe(data: List[Dict[Tuple[str, ...], Any]]): + import pandas as pd + data = [flatten_nested_dict(d) for d in data] + df = pd.DataFrame(data) + df = df.sort_index(axis=1) + df.columns = pd.MultiIndex.from_tuples(df.columns) + return df + + +def recursive_replace(d: Union[List, Dict, str], mapping: Dict[str, str]): + if isinstance(d, str): + for old, new in mapping.items(): + d = d.replace(old, new) + elif isinstance(d, list): + for i, item in enumerate(d): + d[i] = recursive_replace(item, mapping) + elif isinstance(d, dict): + for k, v in d.items(): + d[k] = recursive_replace(v, mapping) + return d + + +class timeit: + _history: Dict[str, List['timeit']] = {} + + def __init__(self, name: str = None, verbose: bool = True, average: bool = False): + self.name = name + self.verbose = verbose + self.start = None + self.end = None + self.average = average + if average and name not in timeit._history: + timeit._history[name] = [] + + def __call__(self, func: Callable): + import inspect + if inspect.iscoroutinefunction(func): + async def wrapper(*args, **kwargs): + with timeit(self.name or func.__qualname__): + ret = await func(*args, **kwargs) + return ret + return wrapper + else: + def wrapper(*args, **kwargs): + with timeit(self.name or func.__qualname__): + ret = func(*args, **kwargs) + return ret + return wrapper + + def __enter__(self): + self.start = time.time() + return self + + @property + def time(self) -> float: + assert self.start is not None, "Time not yet started." + assert self.end is not None, "Time not yet ended." + return self.end - self.start + + @property + def average_time(self) -> float: + assert self.average, "Average time not available." + return sum(t.time for t in timeit._history[self.name]) / len(timeit._history[self.name]) + + @property + def history(self) -> List['timeit']: + return timeit._history.get(self.name, []) + + def __exit__(self, exc_type, exc_val, exc_tb): + self.end = time.time() + if self.average: + timeit._history[self.name].append(self) + if self.verbose: + if self.average: + avg = self.average_time + print(f"{self.name or 'It'} took {avg:.6f} seconds in average.") + else: + print(f"{self.name or 'It'} took {self.time:.6f} seconds.") + + +def strip_common_prefix_suffix(strings: List[str]) -> List[str]: + first = strings[0] + + for start in range(len(first)): + if any(s[start] != strings[0][start] for s in strings): + break + + for end in range(1, min(len(s) for s in strings)): + if any(s[-end] != first[-end] for s in strings): + break + + return [s[start:len(s) - end + 1] for s in strings] + + +def multithead_execute(inputs: List[Any], num_workers: int, pbar = None): + from concurrent.futures import ThreadPoolExecutor + from contextlib import nullcontext + from tqdm import tqdm + + if pbar is not None: + pbar.total = len(inputs) if hasattr(inputs, '__len__') else None + else: + pbar = tqdm(total=len(inputs) if hasattr(inputs, '__len__') else None) + + def decorator(fn: Callable): + with ( + ThreadPoolExecutor(max_workers=num_workers) as executor, + pbar + ): + pbar.refresh() + @catch_exception + @suppress_traceback + def _fn(input): + ret = fn(input) + pbar.update() + return ret + executor.map(_fn, inputs) + executor.shutdown(wait=True) + + return decorator + + +def suppress_traceback(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except Exception as e: + e.__traceback__ = e.__traceback__.tb_next.tb_next + raise + return wrapper + + +class no_warnings: + def __init__(self, action: str = 'ignore', **kwargs): + self.action = action + self.filter_kwargs = kwargs + + def __call__(self, fn): + @wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + warnings.simplefilter(self.action, **self.filter_kwargs) + return fn(*args, **kwargs) + return wrapper + + def __enter__(self): + self.warnings_manager = warnings.catch_warnings() + self.warnings_manager.__enter__() + warnings.simplefilter(self.action, **self.filter_kwargs) + + def __exit__(self, exc_type, exc_val, exc_tb): + self.warnings_manager.__exit__(exc_type, exc_val, exc_tb) + + +def import_file_as_module(file_path: Union[str, os.PathLike], module_name: str): + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module \ No newline at end of file diff --git a/MoGe/moge/utils/vis.py b/MoGe/moge/utils/vis.py new file mode 100644 index 0000000000000000000000000000000000000000..cb9c2378b58ec26ac5067b7ffcbd749a8ad968ce --- /dev/null +++ b/MoGe/moge/utils/vis.py @@ -0,0 +1,65 @@ +from typing import * + +import numpy as np +import matplotlib + + +def colorize_depth(depth: np.ndarray, mask: np.ndarray = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray: + if mask is None: + depth = np.where(depth > 0, depth, np.nan) + else: + depth = np.where((depth > 0) & mask, depth, np.nan) + disp = 1 / depth + if normalize: + min_disp, max_disp = np.nanquantile(disp, 0.001), np.nanquantile(disp, 0.99) + disp = (disp - min_disp) / (max_disp - min_disp) + colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disp)[..., :3], 0) + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_depth_affine(depth: np.ndarray, mask: np.ndarray = None, cmap: str = 'Spectral') -> np.ndarray: + if mask is not None: + depth = np.where(mask, depth, np.nan) + + min_depth, max_depth = np.nanquantile(depth, 0.001), np.nanquantile(depth, 0.999) + depth = (depth - min_depth) / (max_depth - min_depth) + colored = np.nan_to_num(matplotlib.colormaps[cmap](depth)[..., :3], 0) + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_disparity(disparity: np.ndarray, mask: np.ndarray = None, normalize: bool = True, cmap: str = 'Spectral') -> np.ndarray: + if mask is not None: + disparity = np.where(mask, disparity, np.nan) + + if normalize: + min_disp, max_disp = np.nanquantile(disparity, 0.001), np.nanquantile(disparity, 0.999) + disparity = (disparity - min_disp) / (max_disp - min_disp) + colored = np.nan_to_num(matplotlib.colormaps[cmap](1.0 - disparity)[..., :3], 0) + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_segmentation(segmentation: np.ndarray, cmap: str = 'Set1') -> np.ndarray: + colored = matplotlib.colormaps[cmap]((segmentation % 20) / 20)[..., :3] + colored = np.ascontiguousarray((colored.clip(0, 1) * 255).astype(np.uint8)) + return colored + + +def colorize_normal(normal: np.ndarray, mask: np.ndarray = None) -> np.ndarray: + if mask is not None: + normal = np.where(mask[..., None], normal, 0) + normal = normal * [0.5, -0.5, -0.5] + 0.5 + normal = (normal.clip(0, 1) * 255).astype(np.uint8) + return normal + + +def colorize_error_map(error_map: np.ndarray, mask: np.ndarray = None, cmap: str = 'plasma', value_range: Tuple[float, float] = None): + vmin, vmax = value_range if value_range is not None else (np.nanmin(error_map), np.nanmax(error_map)) + cmap = matplotlib.colormaps[cmap] + colorized_error_map = cmap(((error_map - vmin) / (vmax - vmin)).clip(0, 1))[..., :3] + if mask is not None: + colorized_error_map = np.where(mask[..., None], colorized_error_map, 0) + colorized_error_map = np.ascontiguousarray((colorized_error_map.clip(0, 1) * 255).astype(np.uint8)) + return colorized_error_map diff --git a/MoGe/moge/utils/webfile.py b/MoGe/moge/utils/webfile.py new file mode 100644 index 0000000000000000000000000000000000000000..1e98abf8413e1c9f408849b74f4d2025d25511b6 --- /dev/null +++ b/MoGe/moge/utils/webfile.py @@ -0,0 +1,73 @@ +import requests +from typing import * + +__all__ = ["WebFile"] + + +class WebFile: + def __init__(self, url: str, session: Optional[requests.Session] = None, headers: Optional[Dict[str, str]] = None, size: Optional[int] = None): + self.url = url + self.session = session or requests.Session() + self.session.headers.update(headers or {}) + self._offset = 0 + self.size = size if size is not None else self._fetch_size() + + def _fetch_size(self): + with self.session.get(self.url, stream=True) as response: + response.raise_for_status() + content_length = response.headers.get("Content-Length") + if content_length is None: + raise ValueError("Missing Content-Length in header") + return int(content_length) + + def _fetch_data(self, offset: int, n: int) -> bytes: + headers = {"Range": f"bytes={offset}-{min(offset + n - 1, self.size)}"} + response = self.session.get(self.url, headers=headers) + response.raise_for_status() + return response.content + + def seekable(self) -> bool: + return True + + def tell(self) -> int: + return self._offset + + def available(self) -> int: + return self.size - self._offset + + def seek(self, offset: int, whence: int = 0) -> None: + if whence == 0: + new_offset = offset + elif whence == 1: + new_offset = self._offset + offset + elif whence == 2: + new_offset = self.size + offset + else: + raise ValueError("Invalid value for whence") + + self._offset = max(0, min(new_offset, self.size)) + + def read(self, n: Optional[int] = None) -> bytes: + if n is None or n < 0: + n = self.available() + else: + n = min(n, self.available()) + + if n == 0: + return b'' + + data = self._fetch_data(self._offset, n) + self._offset += len(data) + + return data + + def close(self) -> None: + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + pass + + \ No newline at end of file diff --git a/MoGe/moge/utils/webzipfile.py b/MoGe/moge/utils/webzipfile.py new file mode 100644 index 0000000000000000000000000000000000000000..25ed1d3cd34720335eb001d77a278539ffef569b --- /dev/null +++ b/MoGe/moge/utils/webzipfile.py @@ -0,0 +1,128 @@ +from typing import * +import io +import os +from zipfile import ( + ZipInfo, BadZipFile, ZipFile, ZipExtFile, + sizeFileHeader, structFileHeader, stringFileHeader, + _FH_SIGNATURE, _FH_FILENAME_LENGTH, _FH_EXTRA_FIELD_LENGTH, _FH_GENERAL_PURPOSE_FLAG_BITS, + _MASK_COMPRESSED_PATCH, _MASK_STRONG_ENCRYPTION, _MASK_UTF_FILENAME, _MASK_ENCRYPTED +) +import struct +from requests import Session + +from .webfile import WebFile + + +class _SharedWebFile(WebFile): + def __init__(self, webfile: WebFile, pos: int): + super().__init__(webfile.url, webfile.session, size=webfile.size) + self.seek(pos) + + +class WebZipFile(ZipFile): + "Lock-free version of ZipFile that reads from a WebFile, allowing for concurrent reads." + def __init__(self, url: str, session: Optional[Session] = None, headers: Optional[Dict[str, str]] = None): + """Open the ZIP file with mode read 'r', write 'w', exclusive create 'x', + or append 'a'.""" + webf = WebFile(url, session=session, headers=headers) + super().__init__(webf, mode='r') + + def open(self, name, mode="r", pwd=None, *, force_zip64=False): + """Return file-like object for 'name'. + + name is a string for the file name within the ZIP file, or a ZipInfo + object. + + mode should be 'r' to read a file already in the ZIP file, or 'w' to + write to a file newly added to the archive. + + pwd is the password to decrypt files (only used for reading). + + When writing, if the file size is not known in advance but may exceed + 2 GiB, pass force_zip64 to use the ZIP64 format, which can handle large + files. If the size is known in advance, it is best to pass a ZipInfo + instance for name, with zinfo.file_size set. + """ + if mode not in {"r", "w"}: + raise ValueError('open() requires mode "r" or "w"') + if pwd and (mode == "w"): + raise ValueError("pwd is only supported for reading files") + if not self.fp: + raise ValueError( + "Attempt to use ZIP archive that was already closed") + + assert mode == "r", "Only read mode is supported for now" + + # Make sure we have an info object + if isinstance(name, ZipInfo): + # 'name' is already an info object + zinfo = name + elif mode == 'w': + zinfo = ZipInfo(name) + zinfo.compress_type = self.compression + zinfo._compresslevel = self.compresslevel + else: + # Get info object for name + zinfo = self.getinfo(name) + + if mode == 'w': + return self._open_to_write(zinfo, force_zip64=force_zip64) + + if self._writing: + raise ValueError("Can't read from the ZIP file while there " + "is an open writing handle on it. " + "Close the writing handle before trying to read.") + + # Open for reading: + self._fileRefCnt += 1 + zef_file = _SharedWebFile(self.fp, zinfo.header_offset) + + try: + # Skip the file header: + fheader = zef_file.read(sizeFileHeader) + if len(fheader) != sizeFileHeader: + raise BadZipFile("Truncated file header") + fheader = struct.unpack(structFileHeader, fheader) + if fheader[_FH_SIGNATURE] != stringFileHeader: + raise BadZipFile("Bad magic number for file header") + + fname = zef_file.read(fheader[_FH_FILENAME_LENGTH]) + if fheader[_FH_EXTRA_FIELD_LENGTH]: + zef_file.seek(fheader[_FH_EXTRA_FIELD_LENGTH], whence=1) + + if zinfo.flag_bits & _MASK_COMPRESSED_PATCH: + # Zip 2.7: compressed patched data + raise NotImplementedError("compressed patched data (flag bit 5)") + + if zinfo.flag_bits & _MASK_STRONG_ENCRYPTION: + # strong encryption + raise NotImplementedError("strong encryption (flag bit 6)") + + if fheader[_FH_GENERAL_PURPOSE_FLAG_BITS] & _MASK_UTF_FILENAME: + # UTF-8 filename + fname_str = fname.decode("utf-8") + else: + fname_str = fname.decode(self.metadata_encoding or "cp437") + + if fname_str != zinfo.orig_filename: + raise BadZipFile( + 'File name in directory %r and header %r differ.' + % (zinfo.orig_filename, fname)) + + # check for encrypted flag & handle password + is_encrypted = zinfo.flag_bits & _MASK_ENCRYPTED + if is_encrypted: + if not pwd: + pwd = self.pwd + if pwd and not isinstance(pwd, bytes): + raise TypeError("pwd: expected bytes, got %s" % type(pwd).__name__) + if not pwd: + raise RuntimeError("File %r is encrypted, password " + "required for extraction" % name) + else: + pwd = None + + return ZipExtFile(zef_file, mode, zinfo, pwd, True) + except: + zef_file.close() + raise \ No newline at end of file diff --git a/MoGe/pyproject.toml b/MoGe/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..b16a6fe479ce91c4f70c7d4e62564821c710897e --- /dev/null +++ b/MoGe/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "moge" +version = "1.0.0" +description = "MoGe: Unlocking Accurate Monocular Geometry Estimation for Open-Domain Images with Optimal Training Supervision" +readme = "README.md" +license = {text = "MIT"} +dependencies = [ + "click", + "opencv-python", + "scipy", + "matplotlib", + "trimesh", + "pillow", + "huggingface_hub", + "numpy", + "torch>=2.0.0", + "torchvision", + "gradio", + "utils3d @ git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900" +] +requires-python = ">=3.9" + +[project.urls] +Homepage = "https://github.com/microsoft/MoGe" + +[tool.setuptools.packages.find] +where = ["."] +include = ["moge*"] + +[project.scripts] +moge = "moge.scripts.cli:main" \ No newline at end of file diff --git a/MoGe/pyrightconfig.json b/MoGe/pyrightconfig.json new file mode 100644 index 0000000000000000000000000000000000000000..deb3aa62afbda00a7c7413b9eefa6f0ec18fb72b --- /dev/null +++ b/MoGe/pyrightconfig.json @@ -0,0 +1,10 @@ +{ + "include": [ + "moge", + "scripts", + "baselines" + ], + "ignore": [ + "**" + ] +} \ No newline at end of file diff --git a/MoGe/requirements.txt b/MoGe/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..58a1999ea84abc75f138a5152f1e1f8c3ad9a88f --- /dev/null +++ b/MoGe/requirements.txt @@ -0,0 +1,13 @@ +# The versions are not specified since MoGe should be compatible with most versions of the packages. +# If incompatibilities are found, consider upgrading to latest versions or installing the following recommended version of the package. +# torch # >= 2.0.0 +# torchvision +gradio # ==2.8.13 +click # ==8.1.7 +opencv-python # ==4.10.0.84 +scipy # ==1.14.1 +matplotlib # ==3.9.2 +trimesh # ==4.5.1 +pillow # ==10.4.0 +huggingface_hub # ==0.25.2 +git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900 diff --git a/README.md b/README.md index 33fade4c3d2b85ed8391a5b400a5b38c67a125d4..92747eb973fe065acb7023fffe7ebb13cac4caed 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ emoji: 🌖 colorFrom: gray colorTo: yellow sdk: gradio -sdk_version: 6.0.2 +sdk_version: 5.45.0 # ← ADD THIS LINE app_file: app.py pinned: false short_description: Any4D Demo diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..97da1befe8e6c19edbeb55a6bad3b3bce773f809 --- /dev/null +++ b/app.py @@ -0,0 +1,1295 @@ +import os +import tempfile +import uuid +import traceback +import spaces + +import gradio as gr +import torch +import cv2 +import numpy as np +from pathlib import Path +from glob import glob +from natsort import natsorted + +import rerun as rr +import rerun.blueprint as rrb +from gradio_rerun import Rerun + +import hydra +from tqdm import tqdm +import matplotlib.pyplot as plt +import matplotlib as mpl +import matplotlib.cm as cm +from matplotlib.colors import hsv_to_rgb +from PIL import Image +from huggingface_hub import hf_hub_download + +from mapanything.utils.image import load_images, rgb +from mapanything.utils.misc import seed_everything +from mapanything.utils.moge_inference import load_moge_model +from mapanything.utils.inference import loss_of_one_batch_multi_view +from mapanything.models import init_model +from mapanything.utils.geometry import ( + quaternion_to_rotation_matrix, + recover_pinhole_intrinsics_from_ray_directions, + normals_edge, + depth_edge, + points_to_normals, +) + +def log_data_to_rerun(image, depthmap, pose, intrinsics, pts3d, mask, base_name, pts_name, viz_mask=None): + # Log camera info and loaded data + height, width = image.shape[0], image.shape[1] + rr.log( + base_name, + rr.Transform3D( + translation=pose[:3, 3], + mat3x3=pose[:3, :3], + from_parent=False, + ), + ) + rr.log( + f"{base_name}/pinhole", + rr.Pinhole( + image_from_camera=intrinsics, + height=height, + width=width, + camera_xyz=rr.ViewCoordinates.RDF, + ), + ) + rr.log( + f"{base_name}/pinhole/rgb", + rr.Image(image), + ) + + +def log_points_to_rerun(image, pts3d, pts_name, mask=None): + # Log points in 3D + if mask is None: + filtered_pts = pts3d + filtered_pts_col = image + else: + filtered_pts = pts3d[mask] + filtered_pts_col = image[mask] + rr.log( + pts_name, + rr.Points3D( + positions=filtered_pts.reshape(-1, 3), + colors=filtered_pts_col.reshape(-1, 3), + ), + ) + + +def init_hydra_config(config_path, overrides=None): + "Initialize Hydra config" + config_dir = os.path.dirname(config_path) + config_name = os.path.basename(config_path).split(".")[0] + relative_path = os.path.relpath(config_dir, os.path.dirname(__file__)) + hydra.core.global_hydra.GlobalHydra.instance().clear() + hydra.initialize(version_base=None, config_path=relative_path) + if overrides is not None: + cfg = hydra.compose(config_name=config_name, overrides=overrides) + else: + cfg = hydra.compose(config_name=config_name) + + return cfg + + +def log_point_tracks(point_tracks): + """ + Log the point tracks for visualization. + """ + # Create colors based on x-position of initial points + initial_positions = point_tracks[0] + x_coords = initial_positions[:, 0] + x_min = x_coords.min() + x_max = x_coords.max() + + norm = mpl.colors.Normalize(vmin=x_min, vmax=x_max) + normalized_x = norm(x_coords) + colormap = cm.get_cmap('rainbow') + track_colors = colormap(normalized_x)[:, :3] # RGB only + + for track_idx in range(point_tracks.shape[1]): + track = point_tracks[:, track_idx, :] + rr.log( + f"pred/point_tracks/track_{track_idx}", + rr.LineStrips3D( + strips=[track], + colors=[track_colors[track_idx]], + ) + ) + + +def init_inference_model(config, ckpt_path, device): + "Initialize the model for inference" + # Load the model + if isinstance(config, dict): + config_path = config["path"] + overrrides = config["config_overrides"] + model_args = init_hydra_config(config_path, overrides=overrrides) + model = init_model(model_args.model.model_str, model_args.model.model_config) + else: + config_path = config + model_args = init_hydra_config(config_path) + model = init_model(model_args.model_str, model_args.model_config) + model.to(device) + if ckpt_path is not None: + print("Loading model from: ", ckpt_path) + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) + print(model.load_state_dict(ckpt["model"], strict=False)) + model.to(device) + # Set the model to eval mode + model.eval() + + return model + + +@torch.no_grad() +def sample_inference(model, views, device, use_amp): + # Run inference + result = loss_of_one_batch_multi_view( + views, + model, + None, + device, + use_amp=use_amp, + ) + + return result + + +class Any4DProcessor: + """Handles Any4D processing and Rerun visualization""" + + def __init__(self, config_dir, checkpoint_path, machine="local"): + self.config_dir = config_dir + self.machine = machine + self.checkpoint_path = checkpoint_path + + # Download checkpoint from Model Hub if not provided locally + if checkpoint_path is None or not os.path.exists(checkpoint_path): + print("📥 Downloading checkpoint from Hugging Face Model Hub...") + self.checkpoint_path = hf_hub_download( + repo_id="airlabshare/any4d-checkpoint", + filename="any4d_4v_combined.pth", + cache_dir="./checkpoints" + ) + print(f"✓ Checkpoint downloaded to: {self.checkpoint_path}") + else: + self.checkpoint_path = checkpoint_path + print(f"✓ Using local checkpoint: {self.checkpoint_path}") + + seed_everything(0) + + def get_config(self): + """Get model configuration""" + return { + "path": f"{self.config_dir}/train.yaml", + "config_overrides": [ + f"machine={self.machine}", + "model=any4d", + "model.encoder.uses_torch_hub=false", + "model/task=images_only", + ], + "checkpoint_path": self.checkpoint_path, + "trained_with_amp": True, + "data_norm_type": "dinov2", + } + + def extract_frames_from_video(self, video_path, max_frames=100): + """Extract frames from video""" + cap = cv2.VideoCapture(video_path) + frames = [] + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + interval = max(1, total_frames // max_frames) + + frame_count = 0 + while cap.isOpened() and len(frames) < max_frames: + ret, frame = cap.read() + if not ret: + break + + if frame_count % interval == 0: + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append(frame_rgb) + + frame_count += 1 + + cap.release() + return frames + + @spaces.GPU + def process_video_streaming( + self, + recording_id: str, + video_path: str, + max_frames: int = 40, + img_idx: int = 0, + use_scene_flow_mask_refinement: bool = True, + progress=gr.Progress() + ): + """ + Process video and stream results to embedded Rerun viewer. + + This function yields data incrementally to the Rerun viewer + embedded in the Gradio interface. + """ + + # Initialize models inside GPU context to avoid pickling issues + device = "cuda" if torch.cuda.is_available() else "cpu" + high_level_config = self.get_config() + + print("Loading models inside GPU context...") + model = init_inference_model( + high_level_config, + self.checkpoint_path, + device + ) + moge_model = load_moge_model(model_code_path="MoGe", device=device) + print("✓ Models loaded") + + # Create recording stream + rec = rr.RecordingStream( + application_id="any4d_visualization", + recording_id=recording_id + ) + stream = rec.binary_stream() + + try: + progress(0, desc="Extracting frames...") + + # Extract frames + frames = self.extract_frames_from_video(video_path, max_frames=max_frames) + + # Validate img_idx + if img_idx >= len(frames): + raise gr.Error(f"Reference frame index ({img_idx}) must be less than number of frames ({len(frames)})") + + # Save to temp directory + temp_dir = tempfile.mkdtemp() + image_paths = [] + for idx, frame in enumerate(frames): + frame_path = os.path.join(temp_dir, f"frame_{idx:05d}.jpg") + cv2.imwrite(frame_path, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) + image_paths.append(frame_path) + + progress(0.2, desc="Loading images...") + + # Select frame range - img_idx is now a user parameter + start_idx = 0 + end_idx = len(image_paths) + + image_list = [image_paths[img_idx]] + image_list += [image_paths[idx] for idx in range(start_idx, end_idx, 1)] + + # Load images + views = load_images( + image_list, + # size=(width, height), + verbose=True, + norm_type="dinov2", + patch_size=14, + compute_moge_mask=True, + moge_model=moge_model, + binary_mask_path=None + ) + + progress(0.4, desc="Running inference...") + + # Run inference + pred_result = sample_inference( + model, + views, + device, + use_amp=True + ) + + progress(0.6, desc="Creating visualization...") + + # Set up coordinate system + rec.log("pred", rr.ViewCoordinates.RDF, static=True) + + # Create blueprint for better initial view + blueprint = rrb.Blueprint( + rrb.Spatial3DView( + origin="pred", + name="3D Scene", + background=[255, 255, 255], # White color (RGB) + line_grid=rrb.archetypes.LineGrid3D( + visible=False, + ), + ), + collapse_panels=True, + ) + rec.send_blueprint(blueprint) + + # Yield initial setup + yield stream.read() + + # Visualize each frame pair + num_views = len(views) + for idx in range(1, num_views): + progress_val = 0.6 + (0.4 * idx / (num_views - 1)) + progress(progress_val, desc=f"Visualizing frame {idx}/{num_views-1}") + + cur_views = [views[0], views[idx]] + cur_pred_result = { + "view1": pred_result["view1"], + "view2": pred_result[f"view{idx+1}"], + "pred1": pred_result["pred1"], + "pred2": pred_result[f"pred{idx+1}"] + } + + # Set time for animation + rec.set_time_seconds("stable_time", 0.2*idx) + + self.log_visualization_data( + rec, + cur_views, + cur_pred_result, + img_norm_type=high_level_config["data_norm_type"], + use_scene_flow_type="allo_scene_flow", + use_scene_flow_mask_refinement=use_scene_flow_mask_refinement, + ) + + # Yield data to viewer incrementally + yield stream.read() + + progress(1.0, desc="Complete!") + + # Final yield to ensure all data is sent + yield stream.read() + + except Exception as e: + error_msg = f"Error: {str(e)}\n{traceback.format_exc()}" + print(error_msg) + raise gr.Error(error_msg) + + def log_visualization_data(self, rec, views, pred_output, img_norm_type, use_scene_flow_type, use_scene_flow_mask_refinement=True, start_view_idx=0): + # Visualize the results + viz_mask_0 = None + depth_z_0 = None + + for view_idx, view in enumerate(views): + image = rgb(view["img"], norm_type=img_norm_type) + + # Visualize the predicted pointmaps + pts_name = f"pred/pointcloud_view_{view_idx+start_view_idx}" + pts_key = "pts3d" + pred_pts3d = pred_output[f"pred{view_idx+1}"][pts_key][0].cpu() + pred_pts3d_0 = pred_output[f"pred1"]["pts3d"][0].cpu() + + # Get the non ambiguous class mask if available + non_ambiguous_mask = view["non_ambiguous_mask"].cpu() + + # Calculate normal mask + normals, normals_mask = points_to_normals(pred_pts3d.numpy(), mask=non_ambiguous_mask.numpy()) + normal_edges = normals_edge(normals, tol=5, mask=normals_mask) + + # Calculate depth mask + depth_z = pred_output[f"pred{view_idx+1}"]["pts3d_cam"][...,2:3][0].squeeze(-1).cpu().numpy() + depth_edges = depth_edge(depth_z, rtol=0.03, mask=non_ambiguous_mask.numpy()) + + # Combine both edge types + mask = ~(depth_edges & normal_edges) + + # Combine with non ambiguous mask + mask = non_ambiguous_mask.numpy() & mask + + # Close Depth mask + close_depth_mask = depth_z < 40.0 + mask = mask & close_depth_mask + + if view_idx == 0: + viz_mask_0 = mask + + kernel_size = 3 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + viz_mask_0 = viz_mask_0.astype(np.uint8) * 255 # Convert bool → uint8 + viz_mask_0 = cv2.erode(viz_mask_0, kernel, iterations=3) + + depth_z_0 = depth_z.copy() + + if "cam_quats" not in pred_output[f"pred{view_idx+1}"].keys(): + # Visualize the image + base_name = f"pred/image_view_{view_idx+start_view_idx}" + rec.log( + f"{base_name}/pinhole/rgb", + rr.Image(image), + ) + # Visualize the pointmaps + log_points_to_rerun(image[0], pred_pts3d, pts_name, mask=mask) + # Visualize the mask if available + if mask is not None: + rec.log( + f"{base_name}/pinhole/mask", + rr.SegmentationImage(mask.numpy().astype(int)), + ) + else: + base_name = f"pred/image_view_{view_idx+start_view_idx}" + cam_quats = pred_output[f"pred{view_idx+1}"]["cam_quats"][0].cpu() + cam_trans = pred_output[f"pred{view_idx+1}"]["cam_trans"][0].cpu() + ray_directions = pred_output[f"pred{view_idx+1}"]["ray_directions"][0].cpu() + ray_depth = pred_output[f"pred{view_idx+1}"]["depth_along_ray"][0].cpu() + local_pts3d = ray_directions * ray_depth + depth_z = local_pts3d[..., 2:] + if mask is not None: + ambiguous_mask = ~mask + depth_z[ambiguous_mask] = 0 + viz_mask = mask + else: + viz_mask = None + + if "motion_mask" in pred_output[f"pred{view_idx+1}"].keys() and view_idx == 0: + pred_motion_mask = pred_output[f"pred{view_idx+1}"]["motion_mask"][0].cpu().numpy() + viz_mask = viz_mask * pred_motion_mask + + cam_rot = quaternion_to_rotation_matrix(cam_quats) + cam_pose = torch.eye(4) + cam_pose[:3, :3] = cam_rot + cam_pose[:3, 3] = cam_trans + cam_intrinsics = recover_pinhole_intrinsics_from_ray_directions(ray_directions) + self.log_data_to_rerun( + rec, + image[0], + depth_z, + cam_pose, + cam_intrinsics, + pred_pts3d, + mask, + base_name, + pts_name, + viz_mask=viz_mask, + ) + + # Get second camera pose + cam_quats_0 = pred_output[f"pred{1}"]["cam_quats"][0].cpu() + cam_trans_0 = pred_output[f"pred{1}"]["cam_trans"][0].cpu() + cam_rot_0 = quaternion_to_rotation_matrix(cam_quats_0) + cam_pose_0 = torch.eye(4) + cam_pose_0[:3, :3] = cam_rot_0 + cam_pose_0[:3, 3] = cam_trans_0 + + if "scene_flow" in pred_output[f"pred{view_idx+1}"].keys(): + if use_scene_flow_type == "allo_scene_flow": + # Log the predicted scene flow vectors + pred_scene_flow_vectors = pred_output[f"pred{view_idx+1}"]["scene_flow"][0].cpu() + + # Get dynamic scene flow mask from scene flow + motion_sf_mask = pred_scene_flow_vectors.norm(dim=-1) > 1e-1 + scene_flow_mask = viz_mask_0 & motion_sf_mask.numpy() & views[0]["binary_mask"].cpu().numpy() + + mask = (scene_flow_mask > 0) + depth_values = depth_z_0[mask] + if len(depth_values) > 0: + depth_mean = np.median(depth_values) + depth_std = np.std(depth_values) + else: + depth_mean, depth_std = 0, 0 + depth_tolerance = 0.5 * depth_std + depth_min = depth_mean - depth_tolerance + depth_max = depth_mean + depth_tolerance + depth_consistent_mask = np.logical_and( + mask, + np.logical_and(depth_z_0 >= depth_min, depth_z_0 <= depth_max) + ) + scene_flow_mask_refined = depth_consistent_mask.astype(np.uint8) * 255 + + if view_idx == 1: + # Use refined mask or basic mask based on user preference + mask_to_use = scene_flow_mask_refined if use_scene_flow_mask_refinement else viz_mask_0 + self.log_scene_flow_to_rerun(rec, image[0], pred_pts3d_0, pred_scene_flow_vectors, f"pred/scene_flow_{view_idx+start_view_idx}", mask=mask_to_use) + + return scene_flow_mask_refined + + def log_data_to_rerun(self, rec, image, depthmap, pose, intrinsics, pts3d, mask, base_name, pts_name, viz_mask=None): + # Log camera info and loaded data + height, width = image.shape[0], image.shape[1] + rec.log( + base_name, + rr.Transform3D( + translation=pose[:3, 3], + mat3x3=pose[:3, :3], + from_parent=False, + ), + ) + rec.log( + f"{base_name}/pinhole", + rr.Pinhole( + image_from_camera=intrinsics, + height=height, + width=width, + camera_xyz=rr.ViewCoordinates.RDF, + ), + ) + rec.log( + f"{base_name}/pinhole/rgb", + rr.Image(image), + ) + # Log points in 3D + filtered_pts = pts3d[mask] + filtered_pts_col = image[mask] + rec.log( + pts_name, + rr.Points3D( + positions=filtered_pts.reshape(-1, 3), + colors=filtered_pts_col.reshape(-1, 3), + ), + ) + + def log_scene_flow_to_rerun(self, rec, image, pts3d, scene_flow_vecs, base_name, mask=None, scene_bounds=None, colormap_name='rainbow'): + """ + Log scene flow in 3D with color visualization based on flow magnitude and direction + """ + # Filter points based on mask if provided + if mask is None: + filtered_pts = pts3d + filtered_scene_flow_vecs = scene_flow_vecs + filtered_pts_col = image + else: + filtered_pts = pts3d[mask] + filtered_scene_flow_vecs = scene_flow_vecs[mask] + filtered_pts_col = image[mask] + + # Check if we have any valid points + if filtered_pts.numel() == 0: + print(f"Warning: No valid points found for {base_name}. Skipping visualization.") + return scene_bounds + + # Reshape tensors + filtered_pts = filtered_pts.reshape(-1, 3) + filtered_scene_flow_vecs = filtered_scene_flow_vecs.reshape(-1, 3) + + # Sample a subset of points to avoid overcrowding visualization + max_arrows = 10000 + if filtered_pts.shape[0] > max_arrows: + flow_magnitudes = torch.norm(filtered_scene_flow_vecs, dim=1) + + if flow_magnitudes.max() > 1e-6: + probabilities = 0.2 + 0.8 * (flow_magnitudes / (flow_magnitudes.max() + 1e-6)) + probabilities_np = probabilities.cpu().numpy() + probabilities_np = probabilities_np / probabilities_np.sum() + + indices_np = np.random.choice( + filtered_pts.shape[0], + size=max_arrows, + replace=False, + p=probabilities_np + ) + indices = torch.tensor(indices_np, device=filtered_pts.device) + else: + indices = torch.randperm(filtered_pts.shape[0], device=filtered_pts.device)[:max_arrows] + + sampled_pts = filtered_pts[indices] + sampled_vectors = filtered_scene_flow_vecs[indices] + else: + sampled_pts = filtered_pts + sampled_vectors = filtered_scene_flow_vecs + + # Convert to numpy + sampled_pts_np = sampled_pts.cpu().numpy() + sampled_vectors_np = sampled_vectors.cpu().numpy() + + # Calculate flow magnitudes + flow_magnitudes = np.linalg.norm(sampled_vectors_np, axis=1) + + # Calculate bounds for magnitude if not provided + if scene_bounds is None: + if len(flow_magnitudes) == 0: + return (0, 1) + mag_min = flow_magnitudes.min() + mag_max = flow_magnitudes.max() + if mag_min == mag_max: + mag_max = mag_min + 1e-6 + else: + mag_min, mag_max = scene_bounds + + # Create colors based on flow magnitude and direction + if len(flow_magnitudes) > 0 and mag_max > mag_min: + # Normalize flow vectors for direction + normalized_flow = sampled_vectors_np / (flow_magnitudes[:, np.newaxis] + 1e-8) + + # Hue from XZ plane angle + hue = np.arctan2(normalized_flow[:, 2], normalized_flow[:, 0]) + hue = (hue + np.pi) / (2 * np.pi) + + # Saturation and value from magnitude + normalized_magnitude = np.clip((flow_magnitudes - mag_min) / (mag_max - mag_min + 1e-8), 0, 1) + saturation = 0.3 + 0.7 * normalized_magnitude + value = 0.5 + 0.5 * normalized_magnitude + + # Convert HSV to RGB + hsv = np.stack([hue, saturation, value], axis=1) + colors = hsv_to_rgb(hsv) + else: + colors = np.ones((len(sampled_vectors_np), 3)) * 0.5 + + # Log flow vectors as arrows + rec.log( + f"{base_name}/scene_flow", + rr.Arrows3D( + origins=sampled_pts_np, + vectors=sampled_vectors_np, + colors=colors, + ), + ) + + return (mag_min, mag_max) + + +# def create_app(config_dir, checkpoint_path): +# """Create Gradio app with embedded Rerun viewer""" + +# # Initialize processor (without loading models yet) +# processor = Any4DProcessor(config_dir, checkpoint_path) + +# # Create Gradio interface with custom CSS for scrolling and gallery styling +# with gr.Blocks( +# title="Any4D Scene Flow Visualization", +# css=""" +# .gradio-container { +# max-height: 100vh; +# overflow-y: auto !important; +# } +# footer { +# display: none !important; +# } +# /* Style clickable thumbnails */ +# .clickable-thumbnail { +# cursor: pointer; +# transition: transform 0.2s ease; +# } +# .clickable-thumbnail:hover { +# transform: scale(1.05); +# } +# .scene-info { +# text-align: center; +# margin-top: 8px; +# } +# """ +# ) as demo: + +# gr.Markdown(""" +# # 🎬 [Any4D](https://any-4d.github.io/) Rerun Scene Flow Visualization + +# Upload a video or choose an example video from the thumbnails below to visualize 3D reconstruction and scene flow in real-time. + +# Getting Started: +# 1. Upload a video using the input panel on the left. +# 2. Adjust the parameters as needed: +# - **Max Frames to Extract**: Number of frames to extract from the video for processing. +# - **Reference Frame Index**: Select which frame to use as the reference for scene flow computation. +# - **Use Scene Flow Mask Refinement**: Enable to reduce background smearing in scene flow visualization. +# 3. Click "Process Video" to start the visualization. +# 4. Explore the live 3D visualization on the right panel. + +# Note: Inference time for Any4D for 100 frames is approximately less than 1.5 seconds. However, downloading model weights and visualizing 3D points may take tens of seconds. +# Please be patient or, for faster visualization, use a local machine to run our demo from our [GitHub repository](https://github.com/Any-4D/Any4D). +# """) + +# with gr.Row(): +# with gr.Column(scale=1): +# gr.Markdown("### 📥 Input") + +# video_input = gr.Video( +# label="Upload Video", +# sources=["upload"] +# ) + +# with gr.Row(): +# max_frames_slider = gr.Slider( +# minimum=10, +# maximum=200, +# value=40, +# step=1, +# label="Max Frames to Extract", +# info="Number of frames to extract from video" +# ) + +# # Frame preview section +# with gr.Row(): +# frame_preview = gr.Image( +# label="Reference Frame Preview", +# type="numpy", +# interactive=False +# ) + +# with gr.Row(): +# reference_frame_slider = gr.Slider( +# minimum=0, +# maximum=100, +# value=0, +# step=1, +# label="Reference Frame Index", +# info="Select which frame to use as reference for scene flow" +# ) + +# with gr.Row(): +# use_mask_refinement = gr.Checkbox( +# value=False, +# label="Use Scene Flow Mask Refinement", +# info=" Removes background scene-flow smearing near boundaries by assuming dominant motion in scene. Do not use for scenes with multiple dynamic objects." +# ) + +# process_btn = gr.Button("🚀 Process Video", variant="primary", size="lg") + +# status_text = gr.Textbox( +# label="Status", +# lines=2, +# interactive=False +# ) + +# with gr.Column(scale=3): +# gr.Markdown("### 🎥 Live 3D Visualization") + +# # Embedded Rerun viewer with streaming enabled - reduced height for better scrolling +# viewer = Rerun( +# streaming=True, +# height=700, +# panel_states={ +# "time": "collapsed", +# "blueprint": "collapsed", +# "selection": "collapsed", +# }, +# ) + +# # Store extracted frames in session state only +# extracted_frames = gr.State([]) +# loading_example = gr.State(False) # Flag to prevent auto-extract on example load +# preset_ref_state = gr.State(None) # Store preset reference frame from examples + +# # Add Example Videos section below the main interface +# gr.Markdown("---") +# gr.Markdown("## 📂 Example Videos") +# gr.Markdown("Click any thumbnail to load the example video with preset configurations") + +# # Define example configurations +# example_videos = [ +# {"path": "examples/layup_fpv.mp4", "max_frames": 100, "ref_frame": 50, "mask_refinement": True, "name": "Basketball Layup", "thumbnail": "examples/layup_fpv.jpg"}, +# {"path": "examples/cat_jump.mp4", "max_frames": 40, "ref_frame": 0, "mask_refinement": True, "name": "Cat Jump", "thumbnail": "examples/cat_jump.jpg"}, +# {"path": "examples/climbing.mp4", "max_frames": 50, "ref_frame": 0, "mask_refinement": True, "name": "Climbing", "thumbnail": "examples/climbing.jpg"}, +# {"path": "examples/lady-running.mp4", "max_frames": 30, "ref_frame": 15, "mask_refinement": True, "name": "Lady Running", "thumbnail": "examples/lady-running.jpg"}, +# {"path": "examples/mobile_robot.mp4", "max_frames": 50, "ref_frame": 0, "mask_refinement": True, "name": "Mobile Robot", "thumbnail": "examples/mobile_robot.jpg"}, +# {"path": "examples/quadruped_climbing.mp4", "max_frames": 90, "ref_frame": 45, "mask_refinement": True, "name": "Quadruped Climbing", "thumbnail": "examples/quadruped_climbing.jpg"}, +# {"path": "examples/seagull.mp4", "max_frames": 50, "ref_frame": 25, "mask_refinement": True, "name": "Seagull", "thumbnail": "examples/seagull.jpg"}, +# {"path": "examples/Humanoid_Walking.mp4", "max_frames": 50, "ref_frame": 25, "mask_refinement": True, "name": "Humanoid Walking", "thumbnail": "examples/Humanoid_Walking.jpg"}, +# {"path": "examples/car.mp4", "max_frames": 90, "ref_frame": 0, "mask_refinement": True, "name": "Car", "thumbnail": "examples/car.jpg"}, +# {"path": "examples/Bigfoot.mp4", "max_frames": 60, "ref_frame": 30, "mask_refinement": True, "name": "Bigfoot", "thumbnail": "examples/Bigfoot.jpg"}, +# {"path": "examples/humanoid_football.mp4", "max_frames": 22, "ref_frame": 10, "mask_refinement": True, "name": "humanoid_football", "thumbnail": "examples/humanoid_football.png"}, +# {"path": "examples/motorbike_desert.mp4", "max_frames": 50, "ref_frame": 25, "mask_refinement": True, "name": "Motorbike Desert", "thumbnail": "examples/motorbike_desert.jpg"}, +# ] + +# # Function to load example video and extract frames with preset config +# def load_example_video(video_path, max_frames, ref_frame, mask_refinement): +# """Load an example video with its configuration and extract frames""" +# print(f"[Example] Loading {video_path} with ref_frame={ref_frame}, max_frames={max_frames}") +# # Extract frames with the preset reference frame +# frame_preview_img, ref_slider_update, frames = extract_and_preview(video_path, max_frames, preset_ref_frame=ref_frame) +# print(f"[Example] Extracted {len(frames)} frames, slider update value: {ref_slider_update}") +# # Return outputs in order: max_frames_slider, use_mask_refinement, loading_example, video_input, reference_frame_slider, frame_preview, extracted_frames, preset_ref_state +# return max_frames, mask_refinement, True, video_path, ref_slider_update, frame_preview_img, frames, ref_frame + +# # Create thumbnail grid (4 columns per row) +# for i in range(0, len(example_videos), 4): +# with gr.Row(): +# for j in range(4): +# video_idx = i + j +# if video_idx < len(example_videos): +# video = example_videos[video_idx] +# with gr.Column(scale=1, elem_classes=["clickable-thumbnail"]): +# # Clickable thumbnail image +# video_thumb = gr.Image( +# value=video.get("thumbnail", video["path"]), # Use thumbnail if available, else video path +# label="", +# height=150, +# interactive=False, +# show_label=False, +# elem_id=f"video_thumb_{video['name'].replace(' ', '_')}", +# sources=[], +# ) + +# # Video name below thumbnail +# gr.Markdown( +# f"**{video['name']}**", +# elem_classes=["scene-info"], +# ) + +# # Connect thumbnail click to load video +# video_thumb.select( +# fn=lambda p=video["path"], m=video["max_frames"], r=video["ref_frame"], mr=video["mask_refinement"]: +# load_example_video(p, m, r, mr), +# outputs=[max_frames_slider, use_mask_refinement, loading_example, video_input, reference_frame_slider, frame_preview, extracted_frames, preset_ref_state], +# ) +# else: +# # Empty column to maintain grid structure +# with gr.Column(scale=1): +# pass + +# # Function to extract and preview frames +# def extract_and_preview(video_file, max_frames, preset_ref_frame=None, progress=gr.Progress()): +# if video_file is None: +# return None, gr.update(maximum=0, value=0), [] + +# try: +# progress(0, desc="Extracting frames...") + +# # Extract frames +# cap = cv2.VideoCapture(video_file) +# frames = [] +# total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) +# interval = max(1, total_frames // max_frames) + +# frame_count = 0 +# while cap.isOpened() and len(frames) < max_frames: +# ret, frame = cap.read() +# if not ret: +# break + +# if frame_count % interval == 0: +# frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) +# frames.append(frame_rgb) + +# frame_count += 1 + +# cap.release() + +# num_frames = len(frames) + +# # Use preset reference frame if provided, otherwise use middle frame +# if preset_ref_frame is not None and preset_ref_frame < num_frames: +# default_idx = preset_ref_frame +# else: +# default_idx = num_frames // 2 if num_frames > 0 else 0 + +# # Return frame preview and update slider +# return ( +# frames[default_idx] if num_frames > 0 else None, +# gr.update(maximum=max(0, num_frames - 1), value=default_idx), +# frames +# ) + +# except Exception as e: +# print(f"Error extracting frames: {e}") +# return None, gr.update(maximum=0, value=0), [] + +# # Function to update frame preview when slider changes +# def update_frame_preview(frames, frame_idx): +# if frames is None or len(frames) == 0 or frame_idx >= len(frames): +# return None +# return frames[frame_idx] + +# # Processing function +# def process_video_wrapper(video_file, max_frames, ref_frame_idx, use_mask_ref, progress=gr.Progress()): +# if video_file is None: +# raise gr.Error("Please upload a video first") + +# # Generate a new recording ID for this processing run +# new_rec_id = str(uuid.uuid4()) + +# try: +# # Stream data to embedded viewer +# for stream_data in processor.process_video_streaming( +# new_rec_id, +# video_file, +# max_frames, +# ref_frame_idx, +# use_mask_ref, +# progress +# ): +# # Yield both the stream data AND a status message +# yield stream_data, "🔄 Processing... Please wait." + +# # Final yield with completion message +# mask_status = "with refinement" if use_mask_ref else "without refinement" +# yield stream_data, f"✅ Processing complete! Reference frame: {ref_frame_idx}, Scene flow mask {mask_status}. Explore the 3D visualization above." + +# except Exception as e: +# # On error, yield None for viewer and error message for status +# yield None, f"❌ Error: {str(e)}" + +# # Connect video upload to frame extraction (but not when loading examples) +# def handle_video_input_change(video_file, max_frames, is_loading_example): +# """Handle video input changes, skip if loading from example""" +# print(f"[VideoChange] Triggered with is_loading_example={is_loading_example}, video={video_file}") +# if is_loading_example: +# # Don't re-extract, already done in load_example_video +# # Return no update to avoid overwriting what was set +# print("[VideoChange] Skipping extraction (loading from example)") +# return gr.update(), gr.update(), gr.update(), False +# # Normal extraction for user uploads +# if video_file is None: +# return None, gr.update(maximum=0, value=0), [], False +# print("[VideoChange] Extracting frames for user upload") +# frame_preview_img, ref_slider_update, frames = extract_and_preview(video_file, max_frames) +# return frame_preview_img, ref_slider_update, frames, False + +# video_input.change( +# handle_video_input_change, +# inputs=[video_input, max_frames_slider, loading_example], +# outputs=[frame_preview, reference_frame_slider, extracted_frames, loading_example] +# ) + +# # Handle max_frames changes, respecting preset reference frame +# def handle_max_frames_change(video_file, max_frames, preset_ref): +# """Handle max frames change, respecting preset reference frame""" +# frame_preview_img, ref_slider_update, frames = extract_and_preview(video_file, max_frames, preset_ref_frame=preset_ref) +# # Clear preset_ref after using it once +# return frame_preview_img, ref_slider_update, frames, None + +# max_frames_slider.change( +# handle_max_frames_change, +# inputs=[video_input, max_frames_slider, preset_ref_state], +# outputs=[frame_preview, reference_frame_slider, extracted_frames, preset_ref_state] +# ) + +# # Update preview when reference frame slider changes +# reference_frame_slider.change( +# update_frame_preview, +# inputs=[extracted_frames, reference_frame_slider], +# outputs=[frame_preview] +# ) + +# # Connect button to processing +# process_btn.click( +# process_video_wrapper, +# inputs=[video_input, max_frames_slider, reference_frame_slider, use_mask_refinement], +# outputs=[viewer, status_text] +# ) + +# return demo + +def create_app(config_dir, checkpoint_path): + """Create Gradio app with embedded Rerun viewer""" + + # Initialize processor (without loading models yet) + processor = Any4DProcessor(config_dir, checkpoint_path) + + # Create Gradio interface with custom CSS for scrolling and gallery styling + with gr.Blocks( + title="Any4D Scene Flow Visualization", + css=""" + .gradio-container { + max-height: 100vh; + overflow-y: auto !important; + } + footer { + display: none !important; + } + /* Style clickable thumbnails */ + .clickable-thumbnail { + cursor: pointer; + transition: transform 0.2s ease; + } + .clickable-thumbnail:hover { + transform: scale(1.05); + } + .scene-info { + text-align: center; + margin-top: 8px; + } + """ + ) as demo: + + gr.Markdown(""" + # 🎬 [Any4D](https://any-4d.github.io/) Rerun Scene Flow Visualization + + Upload a video or choose an example video from the thumbnails below to visualize 3D reconstruction and scene flow in real-time. + + Getting Started: + 1. Upload a video using the input panel on the left. + 2. Adjust the parameters as needed: + - **Max Frames to Extract**: Number of frames to extract from the video for processing. + - **Reference Frame Index**: Select which frame to use as the reference for scene flow computation. + - **Use Scene Flow Mask Refinement**: Enable to reduce background smearing in scene flow visualization. + 3. Click "Process Video" to start the visualization. + 4. Explore the live 3D visualization on the right panel. + + Note: Inference time for Any4D for 100 frames is approximately less than 1.5 seconds. However, downloading model weights and visualizing 3D points may take tens of seconds. + Please be patient or, for faster visualization, use a local machine to run our demo from our [GitHub repository](https://github.com/Any-4D/Any4D). + """) + + with gr.Row(): + with gr.Column(scale=1): + gr.Markdown("### 📥 Input") + + video_input = gr.Video( + label="Upload Video", + sources=["upload"] + ) + + with gr.Row(): + max_frames_slider = gr.Slider( + minimum=10, + maximum=200, + value=40, + step=1, + label="Max Frames to Extract", + info="Number of frames to extract from video" + ) + + # Frame preview section + with gr.Row(): + frame_preview = gr.Image( + label="Reference Frame Preview", + type="numpy", + interactive=False + ) + + with gr.Row(): + reference_frame_slider = gr.Slider( + minimum=0, + maximum=100, + value=0, + step=1, + label="Reference Frame Index", + info="Select which frame to use as reference for scene flow" + ) + + with gr.Row(): + use_mask_refinement = gr.Checkbox( + value=False, + label="Use Scene Flow Mask Refinement", + info=" Removes background scene-flow smearing near boundaries by assuming dominant motion in scene. Do not use for scenes with multiple dynamic objects." + ) + + process_btn = gr.Button("🚀 Process Video", variant="primary", size="lg") + + with gr.Column(scale=3): + gr.Markdown("### 🎥 Live 3D Visualization") + + # Embedded Rerun viewer with streaming enabled - reduced height for better scrolling + viewer = Rerun( + streaming=True, + height=700, + panel_states={ + "time": "collapsed", + "blueprint": "collapsed", + "selection": "collapsed", + }, + ) + + status_text = gr.Textbox( + label="Status", + lines=2, + interactive=False + ) + + # Store extracted frames in session state only + extracted_frames = gr.State([]) + loading_example = gr.State(False) # Flag to prevent auto-extract on example load + preset_ref_state = gr.State(None) # Store preset reference frame from examples + + # Add Example Videos section below the main interface + gr.Markdown("---") + gr.Markdown("## 📂 Example Videos") + gr.Markdown("Click any thumbnail to load the example video with preset configurations") + + # Define example configurations + example_videos = [ + {"path": "examples/layup_fpv.mp4", "max_frames": 100, "ref_frame": 50, "mask_refinement": True, "name": "Basketball Layup", "thumbnail": "examples/layup_fpv.jpg"}, + {"path": "examples/cat_jump.mp4", "max_frames": 40, "ref_frame": 0, "mask_refinement": True, "name": "Cat Jump", "thumbnail": "examples/cat_jump.jpg"}, + {"path": "examples/climbing.mp4", "max_frames": 50, "ref_frame": 0, "mask_refinement": True, "name": "Climbing", "thumbnail": "examples/climbing.jpg"}, + {"path": "examples/lady-running.mp4", "max_frames": 30, "ref_frame": 15, "mask_refinement": True, "name": "Lady Running", "thumbnail": "examples/lady-running.jpg"}, + {"path": "examples/mobile_robot.mp4", "max_frames": 50, "ref_frame": 0, "mask_refinement": True, "name": "Mobile Robot", "thumbnail": "examples/mobile_robot.jpg"}, + {"path": "examples/quadruped_climbing.mp4", "max_frames": 90, "ref_frame": 45, "mask_refinement": True, "name": "Quadruped Climbing", "thumbnail": "examples/quadruped_climbing.jpg"}, + {"path": "examples/seagull.mp4", "max_frames": 50, "ref_frame": 25, "mask_refinement": True, "name": "Seagull", "thumbnail": "examples/seagull.jpg"}, + {"path": "examples/Humanoid_Walking.mp4", "max_frames": 50, "ref_frame": 25, "mask_refinement": True, "name": "Humanoid Walking", "thumbnail": "examples/Humanoid_Walking.jpg"}, + {"path": "examples/car.mp4", "max_frames": 90, "ref_frame": 0, "mask_refinement": True, "name": "Car", "thumbnail": "examples/car.jpg"}, + {"path": "examples/Bigfoot.mp4", "max_frames": 60, "ref_frame": 30, "mask_refinement": True, "name": "Bigfoot", "thumbnail": "examples/Bigfoot.jpg"}, + {"path": "examples/humanoid_football.mp4", "max_frames": 22, "ref_frame": 10, "mask_refinement": True, "name": "humanoid_football", "thumbnail": "examples/humanoid_football.png"}, + {"path": "examples/motorbike_desert.mp4", "max_frames": 50, "ref_frame": 25, "mask_refinement": True, "name": "Motorbike Desert", "thumbnail": "examples/motorbike_desert.jpg"}, + ] + + # Function to load example video and extract frames with preset config + def load_example_video(video_path, max_frames, ref_frame, mask_refinement): + """Load an example video with its configuration and extract frames""" + print(f"[Example] Loading {video_path} with ref_frame={ref_frame}, max_frames={max_frames}") + # Extract frames with the preset reference frame + frame_preview_img, ref_slider_update, frames = extract_and_preview(video_path, max_frames, preset_ref_frame=ref_frame) + print(f"[Example] Extracted {len(frames)} frames, slider update value: {ref_slider_update}") + # Return outputs in order: max_frames_slider, use_mask_refinement, loading_example, video_input, reference_frame_slider, frame_preview, extracted_frames, preset_ref_state + return max_frames, mask_refinement, True, video_path, ref_slider_update, frame_preview_img, frames, ref_frame + + # Create thumbnail grid (4 columns per row) + for i in range(0, len(example_videos), 4): + with gr.Row(): + for j in range(4): + video_idx = i + j + if video_idx < len(example_videos): + video = example_videos[video_idx] + with gr.Column(scale=1, elem_classes=["clickable-thumbnail"]): + # Clickable thumbnail image + video_thumb = gr.Image( + value=video.get("thumbnail", video["path"]), # Use thumbnail if available, else video path + label="", + height=150, + interactive=False, + show_label=False, + elem_id=f"video_thumb_{video['name'].replace(' ', '_')}", + sources=[], + ) + + # Video name below thumbnail + gr.Markdown( + f"**{video['name']}**", + elem_classes=["scene-info"], + ) + + # Connect thumbnail click to load video + video_thumb.select( + fn=lambda p=video["path"], m=video["max_frames"], r=video["ref_frame"], mr=video["mask_refinement"]: + load_example_video(p, m, r, mr), + outputs=[max_frames_slider, use_mask_refinement, loading_example, video_input, reference_frame_slider, frame_preview, extracted_frames, preset_ref_state], + ) + else: + # Empty column to maintain grid structure + with gr.Column(scale=1): + pass + + # Function to extract and preview frames + def extract_and_preview(video_file, max_frames, preset_ref_frame=None, progress=gr.Progress()): + if video_file is None: + return None, gr.update(maximum=0, value=0), [] + + try: + progress(0, desc="Extracting frames...") + + # Extract frames + cap = cv2.VideoCapture(video_file) + frames = [] + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + interval = max(1, total_frames // max_frames) + + frame_count = 0 + while cap.isOpened() and len(frames) < max_frames: + ret, frame = cap.read() + if not ret: + break + + if frame_count % interval == 0: + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append(frame_rgb) + + frame_count += 1 + + cap.release() + + num_frames = len(frames) + + # Use preset reference frame if provided, otherwise use middle frame + if preset_ref_frame is not None and preset_ref_frame < num_frames: + default_idx = preset_ref_frame + else: + default_idx = num_frames // 2 if num_frames > 0 else 0 + + # Return frame preview and update slider + return ( + frames[default_idx] if num_frames > 0 else None, + gr.update(maximum=max(0, num_frames - 1), value=default_idx), + frames + ) + + except Exception as e: + print(f"Error extracting frames: {e}") + return None, gr.update(maximum=0, value=0), [] + + # Function to update frame preview when slider changes + def update_frame_preview(frames, frame_idx): + if frames is None or len(frames) == 0 or frame_idx >= len(frames): + return None + return frames[frame_idx] + + # Processing function + def process_video_wrapper(video_file, max_frames, ref_frame_idx, use_mask_ref, progress=gr.Progress()): + if video_file is None: + raise gr.Error("Please upload a video first") + + # Generate a new recording ID for this processing run + new_rec_id = str(uuid.uuid4()) + + try: + # Stream data to embedded viewer + for stream_data in processor.process_video_streaming( + new_rec_id, + video_file, + max_frames, + ref_frame_idx, + use_mask_ref, + progress + ): + # Yield both the stream data AND a status message + yield stream_data, "🔄 Processing... Please wait." + + # Final yield with completion message + mask_status = "with refinement" if use_mask_ref else "without refinement" + yield stream_data, f"✅ Processing complete! Reference frame: {ref_frame_idx}, Scene flow mask {mask_status}. Explore the 3D visualization above." + + except Exception as e: + # On error, yield None for viewer and error message for status + yield None, f"❌ Error: {str(e)}" + + # Connect video upload to frame extraction (but not when loading examples) + def handle_video_input_change(video_file, max_frames, is_loading_example): + """Handle video input changes, skip if loading from example""" + print(f"[VideoChange] Triggered with is_loading_example={is_loading_example}, video={video_file}") + if is_loading_example: + # Don't re-extract, already done in load_example_video + # Return no update to avoid overwriting what was set + print("[VideoChange] Skipping extraction (loading from example)") + return gr.update(), gr.update(), gr.update(), False + # Normal extraction for user uploads + if video_file is None: + return None, gr.update(maximum=0, value=0), [], False + print("[VideoChange] Extracting frames for user upload") + frame_preview_img, ref_slider_update, frames = extract_and_preview(video_file, max_frames) + return frame_preview_img, ref_slider_update, frames, False + + video_input.change( + handle_video_input_change, + inputs=[video_input, max_frames_slider, loading_example], + outputs=[frame_preview, reference_frame_slider, extracted_frames, loading_example] + ) + + # Handle max_frames changes, respecting preset reference frame + def handle_max_frames_change(video_file, max_frames, preset_ref): + """Handle max frames change, respecting preset reference frame""" + frame_preview_img, ref_slider_update, frames = extract_and_preview(video_file, max_frames, preset_ref_frame=preset_ref) + # Clear preset_ref after using it once + return frame_preview_img, ref_slider_update, frames, None + + max_frames_slider.change( + handle_max_frames_change, + inputs=[video_input, max_frames_slider, preset_ref_state], + outputs=[frame_preview, reference_frame_slider, extracted_frames, preset_ref_state] + ) + + # Update preview when reference frame slider changes + reference_frame_slider.change( + update_frame_preview, + inputs=[extracted_frames, reference_frame_slider], + outputs=[frame_preview] + ) + + # Connect button to processing + process_btn.click( + process_video_wrapper, + inputs=[video_input, max_frames_slider, reference_frame_slider, use_mask_refinement], + outputs=[viewer, status_text] + ) + + return demo + +if __name__ == "__main__": + # Your actual paths + CONFIG_DIR = "configs" + CHECKPOINT_PATH = None # Will download from Model Hub + + app = create_app(CONFIG_DIR, CHECKPOINT_PATH) + app.launch( + share=True, + server_name="0.0.0.0", + server_port=7860, + ssr_mode=False + ) \ No newline at end of file diff --git a/configs/machine/default.yaml b/configs/machine/default.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1155bbd9ebee093c7e998b980ba56daa5fea0209 --- /dev/null +++ b/configs/machine/default.yaml @@ -0,0 +1,12 @@ +# Root directory containing all datasets +root_data_dir: ??? +# Dataset metadata directory +mapanything_dataset_metadata_dir: ??? +# Root directory containing pretrained checkpoints for custom models +root_pretrained_checkpoints_dir: ??? +# Root directory to log experiments +root_experiments_dir: ??? +# Root directory containing UniCeption pretrained checkpoints +root_uniception_pretrained_checkpoints_dir: ??? +# Root directory containing external benchmark data +external_benchmark_data_root_data_dir: ??? diff --git a/configs/machine/local.yaml b/configs/machine/local.yaml new file mode 100644 index 0000000000000000000000000000000000000000..faf5f3cd6fe981131971ebae4b2b9ed0194d7bb6 --- /dev/null +++ b/configs/machine/local.yaml @@ -0,0 +1,14 @@ +defaults: + - default + +# Root directory containing all datasets +root_data_dir: "/ocean/projects/cis220039p/mdt2/datasets/dydust3r" +# Root directory containing pretrained checkpoints for custom models +root_pretrained_checkpoints_dir: "../../checkpoints" +# Root directory to log experiments +root_experiments_dir: "../../any4d_experiments" +# Root directory containing UniCeption pretrained checkpoints +root_uniception_pretrained_checkpoints_dir: "../../UniCeption/checkpoints" + +# MapAnything dataset metadata directory +mapanything_dataset_metadata_dir: "/ocean/projects/cis220039p/mdt2/datasets/dydust3r/map-anything/mapanything_dataset_metadata" diff --git a/configs/model/any4d.yaml b/configs/model/any4d.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5c1ea2b23f2885e592e9d18fc879a5ac7fd240c0 --- /dev/null +++ b/configs/model/any4d.yaml @@ -0,0 +1,20 @@ +defaults: + - default + - encoder: dinov2_large + - info_sharing: aat_ifr_24_layers + - pred_head: dpt_pose_scale + - pred_head_scene_flow: dpt_scene_flow + - task: images_only + +# String for model factory +model_str: "any4d" +# Model config +model_config: + name: "any4d" + encoder_config: ${model.encoder} + info_sharing_config: ${model.info_sharing} + pred_head_config: ${model.pred_head} + scene_flow_pred_head_config: ${model.pred_head_scene_flow} + geometric_input_config: ${model.task} +# Image Normalization Type +data_norm_type: ${model.encoder.data_norm_type} diff --git a/configs/model/default.yaml b/configs/model/default.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0926e7bf93b61f81bc4397ec2a16331d10f756be --- /dev/null +++ b/configs/model/default.yaml @@ -0,0 +1,16 @@ +# String for model factory (Options: "mapanything", "mapanything_ablations", "modular_dust3r", "vggt", "pi3") +model_str: ??? +# Model config +model_config: + # Path to pretrained model checkpoint + pretrained_checkpoint_path: null + # Load specific submodules from the checkpoint + load_specific_pretrained_submodules: False + # List of submodules to load from the checkpoint (if load_specific_pretrained_submodules is True) + specific_pretrained_submodules: [] +# Path of a starting checkpoint (to enable backward compatibility with original DUSt3R class) +pretrained: null +# Image normalization type +data_norm_type: ??? +# Torch hub force reload +torch_hub_force_reload: False diff --git a/configs/model/encoder/croco_512.yaml b/configs/model/encoder/croco_512.yaml new file mode 100644 index 0000000000000000000000000000000000000000..37d15cb7ac75f83836b7d1fed956a398b61af9df --- /dev/null +++ b/configs/model/encoder/croco_512.yaml @@ -0,0 +1,16 @@ +# UniCeption encoder string used for selecting encoder class (python3 -m uniception.models.encoders.list) +encoder_str: "croco" +# Name of the encoder +name: "croco_512" +# Data normalization type +data_norm_type: "croco" +# Patch embedding class +patch_embed_cls: "PatchEmbedDust3R" +# Image size +img_size: [512, 512] # This parameter has no influence for PatchEmbedDust3R +# Path to the pretrained encoder checkpoint +pretrained_checkpoint_path: '${machine.root_uniception_pretrained_checkpoints_dir}/encoders/CroCo_Encoder_224.pth' +# Override attributes in the pretrained checkpoint +override_checkpoint_attributes: True +# Flag to indicate whether model class uses torch hub +uses_torch_hub: False diff --git a/configs/model/encoder/croco_512_data_norm_dust3r.yaml b/configs/model/encoder/croco_512_data_norm_dust3r.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e98d9d3b91af5fe8ab35256ee8f0fe36a49b9890 --- /dev/null +++ b/configs/model/encoder/croco_512_data_norm_dust3r.yaml @@ -0,0 +1,16 @@ +# UniCeption encoder string used for selecting encoder class (python3 -m uniception.models.encoders.list) +encoder_str: "croco" +# Name of the encoder +name: "croco_512_img_norm_dust3r" +# Data normalization type +data_norm_type: "dust3r" +# Patch embedding class +patch_embed_cls: "PatchEmbedDust3R" +# Image size +img_size: [512, 512] # This parameter has no influence for PatchEmbedDust3R +# Path to the pretrained encoder checkpoint +pretrained_checkpoint_path: '${machine.root_uniception_pretrained_checkpoints_dir}/encoders/CroCo_Encoder_224.pth' +# Override attributes in the pretrained checkpoint +override_checkpoint_attributes: True +# Flag to indicate whether model class uses torch hub +uses_torch_hub: False diff --git a/configs/model/encoder/dinov2_large.yaml b/configs/model/encoder/dinov2_large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a15ad5aa58cfd7da159b4713d206d54ac0a8f9a7 --- /dev/null +++ b/configs/model/encoder/dinov2_large.yaml @@ -0,0 +1,14 @@ +# UniCeption encoder string used for selecting encoder class (python3 -m uniception.models.encoders.list) +encoder_str: "dinov2" +# Name of the encoder +name: "dinov2_large" +# Data normalization type +data_norm_type: "dinov2" +# ViT size +size: "large" +# Registers +with_registers: False +# Flag to indicate whether model class uses torch hub +uses_torch_hub: True +# Flag to indicate whether to use gradient checkpointing for encoder +gradient_checkpointing: False diff --git a/configs/model/encoder/radio_v2_5_large.yaml b/configs/model/encoder/radio_v2_5_large.yaml new file mode 100644 index 0000000000000000000000000000000000000000..95ce4103aa426bcb3a7be5b48184778d935afac2 --- /dev/null +++ b/configs/model/encoder/radio_v2_5_large.yaml @@ -0,0 +1,10 @@ +# UniCeption encoder string used for selecting encoder class (python3 -m uniception.models.encoders.list) +encoder_str: "radio" +# Name of the encoder +name: "radio_v2.5-large" +# Data normalization type +data_norm_type: "radio" +# Model version +model_version: "radio_v2.5-l" +# Flag to indicate whether model class uses torch hub +uses_torch_hub: True diff --git a/configs/model/info_sharing/aat_ifr_24_layers.yaml b/configs/model/info_sharing/aat_ifr_24_layers.yaml new file mode 100644 index 0000000000000000000000000000000000000000..34f787a24dc3d3facbdc6e387046f1eafbe2d099 --- /dev/null +++ b/configs/model/info_sharing/aat_ifr_24_layers.yaml @@ -0,0 +1,22 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "alternating_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "aat_24_layers_ifr" + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 17] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "24_layers" + # Depth + depth: 24 + # Distinguish Reference and Non-Reference Views + distinguish_ref_and_non_ref_views: True + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/info_sharing/aat_ifr_24_layers_escaling.yaml b/configs/model/info_sharing/aat_ifr_24_layers_escaling.yaml new file mode 100644 index 0000000000000000000000000000000000000000..63705b14ff02c25d6a58f59791a10daa479f907f --- /dev/null +++ b/configs/model/info_sharing/aat_ifr_24_layers_escaling.yaml @@ -0,0 +1,24 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "alternating_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "aat_24_layers_ifr" + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 17] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "24_layers" + # Depth + depth: 24 + # Distinguish Reference and Non-Reference Views + distinguish_ref_and_non_ref_views: True + # Scale Entropy in Attention + use_entropy_scaling: True + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/info_sharing/aat_ifr_24_layers_no_ref_view.yaml b/configs/model/info_sharing/aat_ifr_24_layers_no_ref_view.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4abf1d168eca3ccd6c7a8ca2717c1aee1c0e927b --- /dev/null +++ b/configs/model/info_sharing/aat_ifr_24_layers_no_ref_view.yaml @@ -0,0 +1,22 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "alternating_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "aat_24_layers_ifr_no_ref_view" + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 17] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "24_layers" + # Depth + depth: 24 + # Distinguish Reference and Non-Reference Views + distinguish_ref_and_non_ref_views: False + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/info_sharing/aat_ifr_24_layers_w_view_pe.yaml b/configs/model/info_sharing/aat_ifr_24_layers_w_view_pe.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ef2fb6ccae11b740d2ea081181893147d4ee415c --- /dev/null +++ b/configs/model/info_sharing/aat_ifr_24_layers_w_view_pe.yaml @@ -0,0 +1,26 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "alternating_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "aat_24_layers_ifr_w_view_pe" + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 17] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "24_layers" + # Depth + depth: 24 + # Distinguish Reference and Non-Reference Views + distinguish_ref_and_non_ref_views: True + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False + # Maximum number of views for positional encoding + max_num_views_for_pe: 1000 + # Use random indices within range (1, max_num_views_for_pe) for positional encoding of non reference views + use_rand_idx_pe_for_non_reference_views: True diff --git a/configs/model/info_sharing/aat_ifr_48_layers.yaml b/configs/model/info_sharing/aat_ifr_48_layers.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7249706a5b5031e71218604df44cd377a90dc37b --- /dev/null +++ b/configs/model/info_sharing/aat_ifr_48_layers.yaml @@ -0,0 +1,26 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "alternating_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "aat_48_layers_ifr" + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 23, 35] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "48_layers" + # Depth (this includes both frame-wise and gloabl attention layers) + depth: 48 + # Feature dim (similar to ViT-Large) + dim: 1024 + # Number of heads (similar to ViT-Large) + num_heads: 16 + # Distinguish Reference and Non-Reference Views + distinguish_ref_and_non_ref_views: True + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/info_sharing/aat_ifr_48_layers_escaling.yaml b/configs/model/info_sharing/aat_ifr_48_layers_escaling.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fc8ca1da63b0a523ea537d47515974473820ef9e --- /dev/null +++ b/configs/model/info_sharing/aat_ifr_48_layers_escaling.yaml @@ -0,0 +1,28 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "alternating_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "aat_48_layers_ifr" + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 23, 35] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "48_layers" + # Depth (this includes both frame-wise and gloabl attention layers) + depth: 48 + # Feature dim (similar to ViT-Large) + dim: 1024 + # Number of heads (similar to ViT-Large) + num_heads: 16 + # Distinguish Reference and Non-Reference Views + distinguish_ref_and_non_ref_views: True + # Scale Entropy in Attention + use_entropy_scaling: True + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/info_sharing/aat_ifr_48_layers_no_ref_view.yaml b/configs/model/info_sharing/aat_ifr_48_layers_no_ref_view.yaml new file mode 100644 index 0000000000000000000000000000000000000000..68456c32359c48f5e626a2ef9ecb0eb08faa2d54 --- /dev/null +++ b/configs/model/info_sharing/aat_ifr_48_layers_no_ref_view.yaml @@ -0,0 +1,26 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "alternating_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "aat_48_layers_ifr_no_ref_view" + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 23, 35] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "48_layers" + # Depth (this includes both frame-wise and gloabl attention layers) + depth: 48 + # Feature dim (similar to ViT-Large) + dim: 1024 + # Number of heads (similar to ViT-Large) + num_heads: 16 + # Distinguish Reference and Non-Reference Views + distinguish_ref_and_non_ref_views: False + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/info_sharing/cat_ifr_dust3r.yaml b/configs/model/info_sharing/cat_ifr_dust3r.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e52e1d37cefc92d7121bb6a4076198fef8a72ee6 --- /dev/null +++ b/configs/model/info_sharing/cat_ifr_dust3r.yaml @@ -0,0 +1,18 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "cross_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: "RoPE100" +# Module arguments +module_args: + # Name of the info sharing module + name: "base_cat_ifr_dust3r" + # Number of views + num_views: 2 + # Indices of the intermediate features to be shared (indices start from 0) + indices: [5, 8] + # Normalize intermediate features + norm_intermediate: False + # Load CroCo cross-attention transformer for DUSt3R Init + pretrained_checkpoint_path: '${machine.root_uniception_pretrained_checkpoints_dir}/info_sharing/cross_attn_transformer/Two_View_Cross_Attention_Transformer_CroCo.pth' diff --git a/configs/model/info_sharing/gat_ifr_24_layers.yaml b/configs/model/info_sharing/gat_ifr_24_layers.yaml new file mode 100644 index 0000000000000000000000000000000000000000..60423714da681f71ca2088f101f4b9fa734b254b --- /dev/null +++ b/configs/model/info_sharing/gat_ifr_24_layers.yaml @@ -0,0 +1,24 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "global_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "gat_24_layers_ifr" + # Maximum number of views for positional encoding + max_num_views: 1000 + # Use random indices within range (1, max_num_views) for positional encoding of non reference views + use_rand_idx_pe_for_non_reference_views: True + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 17] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "24_layers" + # Depth + depth: 24 + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/info_sharing/gat_ifr_24_layers_escaling.yaml b/configs/model/info_sharing/gat_ifr_24_layers_escaling.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6143b5b37adde56a6a742b01a9b0d4fc0824fa35 --- /dev/null +++ b/configs/model/info_sharing/gat_ifr_24_layers_escaling.yaml @@ -0,0 +1,26 @@ +# Model type (Options: ["cross_attention", "global_attention", "alternating_attention"]) +model_type: "global_attention" +# Model class type (Options: ["no_intermediate_features", "intermediate_features"]) +model_return_type: "intermediate_features" +# Custom positional encoding (Options: ["RoPEfreq"], Callable Function, null) +custom_positional_encoding: null +# Module arguments +module_args: + # Name of the info sharing module + name: "gat_24_layers_ifr" + # Maximum number of views for positional encoding + max_num_views: 1000 + # Use random indices within range (1, max_num_views) for positional encoding of non reference views + use_rand_idx_pe_for_non_reference_views: True + # Indices of the intermediate features to be shared (indices start from 0) + indices: [11, 17] + # Normalize intermediate features + norm_intermediate: True + # Size string + size: "24_layers" + # Depth + depth: 24 + # Scale Entropy in Attention + use_entropy_scaling: True + # Flag to indicate whether to use gradient checkpointing + gradient_checkpointing: False diff --git a/configs/model/pred_head/adaptor_config/campointmap_pose_confidence_mask.yaml b/configs/model/pred_head/adaptor_config/campointmap_pose_confidence_mask.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c98de936666ba0c0f469df99d75935f34f492cc4 --- /dev/null +++ b/configs/model/pred_head/adaptor_config/campointmap_pose_confidence_mask.yaml @@ -0,0 +1,22 @@ +# Camera Frame Pointmap + Global Camera Pose (Trans + Quats) + Confidence + Mask +input_dim: 5 +scene_rep_dim: 3 +type: "campointmap+pose+confidence+mask" +scene_rep_type: "campointmap+pose" +dense_pred_init_dict: + name: "campointmap+pose+confidence+mask+scale" + pointmap_mode: "z_exp" + pointmap_vmin: ${special_float:"-inf"} + pointmap_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} +pose_pred_init_dict: + name: "campointmap+pose+confidence+mask+scale" + cam_trans_mode: "linear" + cam_trans_vmin: ${special_float:"-inf"} + cam_trans_vmax: ${special_float:"inf"} + quaternions_mode: "linear" + quaternions_normalize: true + quaternions_vmin: ${special_float:"-inf"} + quaternions_vmax: ${special_float:"inf"} diff --git a/configs/model/pred_head/adaptor_config/campointmap_pose_confidence_mask_scale.yaml b/configs/model/pred_head/adaptor_config/campointmap_pose_confidence_mask_scale.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5ef92a540e7e2979c73541ed532ee225b0df3c11 --- /dev/null +++ b/configs/model/pred_head/adaptor_config/campointmap_pose_confidence_mask_scale.yaml @@ -0,0 +1,27 @@ +# Camera Frame Pointmap + Global Camera Pose (Trans + Quats) + Confidence + Mask + Scene-wide Metric Scaling Factor +input_dim: 5 +scene_rep_dim: 3 +type: "campointmap+pose+confidence+mask" +scene_rep_type: "campointmap+pose" +dense_pred_init_dict: + name: "campointmap+pose+confidence+mask+scale" + pointmap_mode: "z_exp" + pointmap_vmin: ${special_float:"-inf"} + pointmap_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} +pose_pred_init_dict: + name: "campointmap+pose+confidence+mask+scale" + cam_trans_mode: "linear" + cam_trans_vmin: ${special_float:"-inf"} + cam_trans_vmax: ${special_float:"inf"} + quaternions_mode: "linear" + quaternions_normalize: true + quaternions_vmin: ${special_float:"-inf"} + quaternions_vmax: ${special_float:"inf"} +scale_pred_init_dict: + name: "campointmap+pose+confidence+mask+scale" + mode: "exp" + vmin: 1e-08 + vmax: ${special_float:"inf"} diff --git a/configs/model/pred_head/adaptor_config/pointmap_confidence.yaml b/configs/model/pred_head/adaptor_config/pointmap_confidence.yaml new file mode 100644 index 0000000000000000000000000000000000000000..96b60f9a47407a048c1f4af698f5de693eb1d109 --- /dev/null +++ b/configs/model/pred_head/adaptor_config/pointmap_confidence.yaml @@ -0,0 +1,13 @@ +# Pointmap + Confidence +input_dim: 4 +scene_rep_dim: 3 +type: "pointmap+confidence" +scene_rep_type: "pointmap" +init_dict: + name: "pointmap+confidence" + pointmap_mode: "exp" + pointmap_vmin: ${special_float:"-inf"} + pointmap_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} diff --git a/configs/model/pred_head/adaptor_config/pointmap_confidence_mask.yaml b/configs/model/pred_head/adaptor_config/pointmap_confidence_mask.yaml new file mode 100644 index 0000000000000000000000000000000000000000..833df59a2a8d941788cbf9ef6e988cd8e0cf5e6b --- /dev/null +++ b/configs/model/pred_head/adaptor_config/pointmap_confidence_mask.yaml @@ -0,0 +1,13 @@ +# Pointmap + Confidence + Mask +input_dim: 5 +scene_rep_dim: 3 +type: "pointmap+confidence+mask" +scene_rep_type: "pointmap" +init_dict: + name: "pointmap+confidence+mask" + pointmap_mode: "exp" + pointmap_vmin: ${special_float:"-inf"} + pointmap_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} diff --git a/configs/model/pred_head/adaptor_config/pointmap_confidence_mask_scale.yaml b/configs/model/pred_head/adaptor_config/pointmap_confidence_mask_scale.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4230d94989652d7b5ee779bebd7d8ac52d8e3cd7 --- /dev/null +++ b/configs/model/pred_head/adaptor_config/pointmap_confidence_mask_scale.yaml @@ -0,0 +1,18 @@ +# Pointmap + Confidence + Mask + Scene-wide Metric Scaling Factor +input_dim: 5 +scene_rep_dim: 3 +type: "pointmap+confidence+mask" +scene_rep_type: "pointmap" +init_dict: + name: "pointmap+confidence+mask" + pointmap_mode: "exp" + pointmap_vmin: ${special_float:"-inf"} + pointmap_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} +scale_pred_init_dict: + name: "campointmap+pose+confidence+mask+scale" + mode: "exp" + vmin: 1e-08 + vmax: ${special_float:"inf"} diff --git a/configs/model/pred_head/adaptor_config/pointmap_factored_raydirs_depth_pose_confidence_mask_scale.yaml b/configs/model/pred_head/adaptor_config/pointmap_factored_raydirs_depth_pose_confidence_mask_scale.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2f93a518ba2322cc8e6915abfa4b1cdf00fb2dbf --- /dev/null +++ b/configs/model/pred_head/adaptor_config/pointmap_factored_raydirs_depth_pose_confidence_mask_scale.yaml @@ -0,0 +1,39 @@ +# Global Pointmaps + Ray Directions on Unit Sphere + Depth along Ray + Global Camera Pose (Trans + Quats) + Confidence + Mask + Global Metric Scaling Factor +input_dim: 9 +scene_rep_dim: 7 +type: "pointmap+raydirs+depth+pose+confidence+mask" +scene_rep_type: "pointmap+raydirs+depth+pose" +dense_pred_init_dict: + name: "pointmap+raydirs+depth+pose+confidence+mask+scale" + pointmap_mode: "exp" + pointmap_vmin: ${special_float:"-inf"} + pointmap_vmax: ${special_float:"inf"} + ray_directions_mode: "linear" + ray_directions_normalize_to_unit_sphere: true + ray_directions_normalize_to_unit_image_plane: false + ray_directions_vmin: ${special_float:"-inf"} + ray_directions_vmax: ${special_float:"inf"} + ray_directions_clamp_min_of_z_dir: false + ray_directions_z_dir_min: ${special_float:"-inf"} + depth_mode: "exp" + depth_vmin: 0 + depth_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} +pose_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask+scale" + cam_trans_mode: "linear" + cam_trans_vmin: ${special_float:"-inf"} + cam_trans_vmax: ${special_float:"inf"} + quaternions_mode: "linear" + quaternions_normalize: true + quaternions_vmin: ${special_float:"-inf"} + quaternions_vmax: ${special_float:"inf"} +scale_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask+scale" + mode: "exp" + vmin: 1e-08 + vmax: ${special_float:"inf"} +# Flag to decide what representaion to use for global pointmaps +use_factored_predictions_for_global_pointmaps: true diff --git a/configs/model/pred_head/adaptor_config/pointmap_raydirs_depth_pose_confidence_mask_scale.yaml b/configs/model/pred_head/adaptor_config/pointmap_raydirs_depth_pose_confidence_mask_scale.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bfb2408fd39ab13bd619b3b503e29f9900433a61 --- /dev/null +++ b/configs/model/pred_head/adaptor_config/pointmap_raydirs_depth_pose_confidence_mask_scale.yaml @@ -0,0 +1,39 @@ +# Global Pointmaps + Ray Directions on Unit Sphere + Depth along Ray + Global Camera Pose (Trans + Quats) + Confidence + Mask + Global Metric Scaling Factor +input_dim: 9 +scene_rep_dim: 7 +type: "pointmap+raydirs+depth+pose+confidence+mask" +scene_rep_type: "pointmap+raydirs+depth+pose" +dense_pred_init_dict: + name: "pointmap+raydirs+depth+pose+confidence+mask+scale" + pointmap_mode: "exp" + pointmap_vmin: ${special_float:"-inf"} + pointmap_vmax: ${special_float:"inf"} + ray_directions_mode: "linear" + ray_directions_normalize_to_unit_sphere: true + ray_directions_normalize_to_unit_image_plane: false + ray_directions_vmin: ${special_float:"-inf"} + ray_directions_vmax: ${special_float:"inf"} + ray_directions_clamp_min_of_z_dir: false + ray_directions_z_dir_min: ${special_float:"-inf"} + depth_mode: "exp" + depth_vmin: 0 + depth_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} +pose_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask+scale" + cam_trans_mode: "linear" + cam_trans_vmin: ${special_float:"-inf"} + cam_trans_vmax: ${special_float:"inf"} + quaternions_mode: "linear" + quaternions_normalize: true + quaternions_vmin: ${special_float:"-inf"} + quaternions_vmax: ${special_float:"inf"} +scale_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask+scale" + mode: "exp" + vmin: 1e-08 + vmax: ${special_float:"inf"} +# Flag to decide what representaion to use for global pointmaps +use_factored_predictions_for_global_pointmaps: false diff --git a/configs/model/pred_head/adaptor_config/raydirs_depth_pose_confidence_mask.yaml b/configs/model/pred_head/adaptor_config/raydirs_depth_pose_confidence_mask.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ee83ea0b976e3d173302d167b51404446b950b11 --- /dev/null +++ b/configs/model/pred_head/adaptor_config/raydirs_depth_pose_confidence_mask.yaml @@ -0,0 +1,29 @@ +# Ray Directions on Unit Sphere + Depth along Ray + Global Camera Pose (Trans + Quats) + Confidence + Mask +input_dim: 6 +scene_rep_dim: 4 +type: "raydirs+depth+pose+confidence+mask" +scene_rep_type: "raydirs+depth+pose" +dense_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask" + ray_directions_mode: "linear" + ray_directions_normalize_to_unit_sphere: true + ray_directions_normalize_to_unit_image_plane: false + ray_directions_vmin: ${special_float:"-inf"} + ray_directions_vmax: ${special_float:"inf"} + ray_directions_clamp_min_of_z_dir: false + ray_directions_z_dir_min: ${special_float:"-inf"} + depth_mode: "exp" + depth_vmin: 0 + depth_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} +pose_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask" + cam_trans_mode: "linear" + cam_trans_vmin: ${special_float:"-inf"} + cam_trans_vmax: ${special_float:"inf"} + quaternions_mode: "linear" + quaternions_normalize: true + quaternions_vmin: ${special_float:"-inf"} + quaternions_vmax: ${special_float:"inf"} diff --git a/configs/model/pred_head/adaptor_config/raydirs_depth_pose_confidence_mask_scale.yaml b/configs/model/pred_head/adaptor_config/raydirs_depth_pose_confidence_mask_scale.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4a939ad9eeed4131990941976bde8124bfbeba15 --- /dev/null +++ b/configs/model/pred_head/adaptor_config/raydirs_depth_pose_confidence_mask_scale.yaml @@ -0,0 +1,34 @@ +# Ray Directions on Unit Sphere + Depth along Ray + Global Camera Pose (Trans + Quats) + Confidence + Mask + Global Metric Scaling Factor +input_dim: 6 +scene_rep_dim: 4 +type: "raydirs+depth+pose+confidence+mask" +scene_rep_type: "raydirs+depth+pose" +dense_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask+scale" + ray_directions_mode: "linear" + ray_directions_normalize_to_unit_sphere: true + ray_directions_normalize_to_unit_image_plane: false + ray_directions_vmin: ${special_float:"-inf"} + ray_directions_vmax: ${special_float:"inf"} + ray_directions_clamp_min_of_z_dir: false + ray_directions_z_dir_min: ${special_float:"-inf"} + depth_mode: "exp" + depth_vmin: 0 + depth_vmax: ${special_float:"inf"} + confidence_type: "exp" + confidence_vmin: 1 + confidence_vmax: ${special_float:"inf"} +pose_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask+scale" + cam_trans_mode: "linear" + cam_trans_vmin: ${special_float:"-inf"} + cam_trans_vmax: ${special_float:"inf"} + quaternions_mode: "linear" + quaternions_normalize: true + quaternions_vmin: ${special_float:"-inf"} + quaternions_vmax: ${special_float:"inf"} +scale_pred_init_dict: + name: "raydirs+depth+pose+confidence+mask+scale" + mode: "exp" + vmin: 1e-08 + vmax: ${special_float:"inf"} diff --git a/configs/model/pred_head/dpt.yaml b/configs/model/pred_head/dpt.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2d05fe0b2695236470388b522d0ed70fcf603b1d --- /dev/null +++ b/configs/model/pred_head/dpt.yaml @@ -0,0 +1,15 @@ +defaults: + - adaptor_config: pointmap_confidence + +type: "dpt" +feature_head: + feature_dim: 256 + hooks: [0, 1, 2, 3] + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +regressor_head: + output_dim: ${model.pred_head.adaptor_config.input_dim} + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +adaptor_type: ${model.pred_head.adaptor_config.type} +adaptor: ${model.pred_head.adaptor_config.init_dict} +# Flag to indicate whether to use gradient checkpointing +gradient_checkpointing: False diff --git a/configs/model/pred_head/dpt_pose.yaml b/configs/model/pred_head/dpt_pose.yaml new file mode 100644 index 0000000000000000000000000000000000000000..24d50ae5dbfa1ada961b0afde5422ce9dac97a8b --- /dev/null +++ b/configs/model/pred_head/dpt_pose.yaml @@ -0,0 +1,19 @@ +defaults: + - adaptor_config: raydirs_depth_pose_confidence_mask + +type: "dpt+pose" +feature_head: + feature_dim: 256 + hooks: [0, 1, 2, 3] + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +regressor_head: + output_dim: ${model.pred_head.adaptor_config.input_dim} + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +pose_head: + num_resconv_block: 2 + rot_representation_dim: 4 +adaptor_type: ${model.pred_head.adaptor_config.type} +dpt_adaptor: ${model.pred_head.adaptor_config.dense_pred_init_dict} +pose_adaptor: ${model.pred_head.adaptor_config.pose_pred_init_dict} +# Flag to indicate whether to use gradient checkpointing +gradient_checkpointing: False diff --git a/configs/model/pred_head/dpt_pose_scale.yaml b/configs/model/pred_head/dpt_pose_scale.yaml new file mode 100644 index 0000000000000000000000000000000000000000..068dc181712e4d92f2e1e1e088f363ed698c2582 --- /dev/null +++ b/configs/model/pred_head/dpt_pose_scale.yaml @@ -0,0 +1,22 @@ +defaults: + - adaptor_config: raydirs_depth_pose_confidence_mask_scale + +type: "dpt+pose" +feature_head: + feature_dim: 256 + hooks: [0, 1, 2, 3] + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +regressor_head: + output_dim: ${model.pred_head.adaptor_config.input_dim} + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +pose_head: + num_resconv_block: 2 + rot_representation_dim: 4 +scale_head: + output_dim: 1 +adaptor_type: ${model.pred_head.adaptor_config.type} +dpt_adaptor: ${model.pred_head.adaptor_config.dense_pred_init_dict} +pose_adaptor: ${model.pred_head.adaptor_config.pose_pred_init_dict} +scale_adaptor: ${model.pred_head.adaptor_config.scale_pred_init_dict} +# Flag to indicate whether to use gradient checkpointing +gradient_checkpointing: False diff --git a/configs/model/pred_head/dpt_scale.yaml b/configs/model/pred_head/dpt_scale.yaml new file mode 100644 index 0000000000000000000000000000000000000000..286eb589259cd07c6ff0d2c24e5a58a6d1c748c5 --- /dev/null +++ b/configs/model/pred_head/dpt_scale.yaml @@ -0,0 +1,18 @@ +defaults: + - adaptor_config: pointmap_confidence_mask_scale + +type: "dpt" +feature_head: + feature_dim: 256 + hooks: [0, 1, 2, 3] + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +regressor_head: + output_dim: ${model.pred_head.adaptor_config.input_dim} + checkpoint_gradient: ${model.pred_head.gradient_checkpointing} +scale_head: + output_dim: 1 +adaptor_type: ${model.pred_head.adaptor_config.type} +adaptor: ${model.pred_head.adaptor_config.init_dict} +scale_adaptor: ${model.pred_head.adaptor_config.scale_pred_init_dict} +# Flag to indicate whether to use gradient checkpointing +gradient_checkpointing: False diff --git a/configs/model/pred_head_scene_flow/adaptor_config/scene_flow.yaml b/configs/model/pred_head_scene_flow/adaptor_config/scene_flow.yaml new file mode 100644 index 0000000000000000000000000000000000000000..42913943b95db3b42d5bb2cd418672875f5c7cdd --- /dev/null +++ b/configs/model/pred_head_scene_flow/adaptor_config/scene_flow.yaml @@ -0,0 +1,10 @@ +# Ray Directions on Unit Sphere + Depth along Ray + Global Camera Pose (Trans + Quats) + Confidence + Mask +input_dim: 3 +scene_rep_dim: 3 +type: "scene_flow" +scene_rep_type: "scene_flow" +dense_pred_init_dict: + name: "scene_flow" + mode: "linear" + vmin: ${special_float:"-inf"} + vmax: ${special_float:"inf"} \ No newline at end of file diff --git a/configs/model/pred_head_scene_flow/dpt_scene_flow.yaml b/configs/model/pred_head_scene_flow/dpt_scene_flow.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dbba110ad44f35832071e55ab41098a80fbacfb2 --- /dev/null +++ b/configs/model/pred_head_scene_flow/dpt_scene_flow.yaml @@ -0,0 +1,11 @@ +defaults: + - adaptor_config: scene_flow + +type: "dpt" +feature_head: + feature_dim: 256 + hooks: [0, 1, 2, 3] +regressor_head: + output_dim: ${model.pred_head_scene_flow.adaptor_config.input_dim} +adaptor_type: ${model.pred_head_scene_flow.adaptor_config.type} +dpt_adaptor: ${model.pred_head_scene_flow.adaptor_config.dense_pred_init_dict} \ No newline at end of file diff --git a/configs/model/task/aug_training.yaml b/configs/model/task/aug_training.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5d3f618484b0cfddee81a989774c4e6785b3e513 --- /dev/null +++ b/configs/model/task/aug_training.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 0.9 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0.05 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0.5 +# Probability of Geometric Inputs with Depths +depth_prob: 0.5 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0.5 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0.5 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0.9 +# Probability for skipping input of the metric scale quantities for the input metric high quality gt depth +# If 0, the metric scale quantities will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0.05 +# Probability for skipping input of the metric scale quantities for the input metric pose +# If 0, the metric scale quantities will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0.05 +# Doppler Input Probability +doppler_prob: 0.5 \ No newline at end of file diff --git a/configs/model/task/calibrated_sfm.yaml b/configs/model/task/calibrated_sfm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8d24c0033e82e7db23c7c6c40428d5c295380c8b --- /dev/null +++ b/configs/model/task/calibrated_sfm.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 diff --git a/configs/model/task/default.yaml b/configs/model/task/default.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b7c4d3c42e16f6aeb0a26f159f2d1c8d2796bd6f --- /dev/null +++ b/configs/model/task/default.yaml @@ -0,0 +1,33 @@ +# Ray Directions Encoder Config +ray_dirs_encoder_config: + name: "ray_dirs_encoder" + in_chans: 3 + encoder_str: "dense_rep_encoder" + apply_pe: false +# Depth Encoder Config +depth_encoder_config: + name: "depth_encoder" + in_chans: 1 + encoder_str: "dense_rep_encoder" + apply_pe: false +# Cam Rotation (Quats) Encoder Config +cam_rot_encoder_config: + name: "cam_rot_quats_encoder" + in_chans: 4 + encoder_str: "global_rep_encoder" +# Cam Translation Encoder Config +cam_trans_encoder_config: + name: "cam_trans_encoder" + in_chans: 3 + encoder_str: "global_rep_encoder" +# Scale Encoder Config +scale_encoder_config: + name: "scale_encoder" + in_chans: 1 + encoder_str: "global_rep_encoder" +# Scene Flow Encoder Config +scene_flow_encoder_config: + name: "scene_flow_encoder" + in_chans: 3 + encoder_str: "dense_rep_encoder" + apply_pe: false \ No newline at end of file diff --git a/configs/model/task/depth_aug_training.yaml b/configs/model/task/depth_aug_training.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cda5e4531062eb61ab48ee77b739fd28a7c84060 --- /dev/null +++ b/configs/model/task/depth_aug_training.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0.5 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 0 \ No newline at end of file diff --git a/configs/model/task/depth_aug_training_with_image_dropout.yaml b/configs/model/task/depth_aug_training_with_image_dropout.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c525db48c8d2ba2e43a3b6abdc038ecbfe9137e1 --- /dev/null +++ b/configs/model/task/depth_aug_training_with_image_dropout.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 0.25 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 0 \ No newline at end of file diff --git a/configs/model/task/depth_completion.yaml b/configs/model/task/depth_completion.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0461a2c7fce649a4b1fe4e55b80955375197a88e --- /dev/null +++ b/configs/model/task/depth_completion.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 1 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0.9 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 diff --git a/configs/model/task/doppler.yaml b/configs/model/task/doppler.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8932851c3d8c9055bd6cbf9136072f18cf141370 --- /dev/null +++ b/configs/model/task/doppler.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 1 \ No newline at end of file diff --git a/configs/model/task/doppler_aug_training.yaml b/configs/model/task/doppler_aug_training.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d1f1e7d175f2c5f30d9fde4d802170e3309d4e51 --- /dev/null +++ b/configs/model/task/doppler_aug_training.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 0.5 \ No newline at end of file diff --git a/configs/model/task/doppler_aug_training_with_image_dropout.yaml b/configs/model/task/doppler_aug_training_with_image_dropout.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9a858f6f38c51f1d11aea33c7d2704bda99a4a7f --- /dev/null +++ b/configs/model/task/doppler_aug_training_with_image_dropout.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 0.25 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 1 \ No newline at end of file diff --git a/configs/model/task/doppler_only.yaml b/configs/model/task/doppler_only.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b095eba177ff6809db0ff1e8f4fd3811fba2578b --- /dev/null +++ b/configs/model/task/doppler_only.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 0 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 1 \ No newline at end of file diff --git a/configs/model/task/images_only.yaml b/configs/model/task/images_only.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7a2dd5c90c2b44b68d4bcc4edc721963da14769b --- /dev/null +++ b/configs/model/task/images_only.yaml @@ -0,0 +1,29 @@ +defaults: + - default + +# Probability of Dropping out Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 0 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 1 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 0 +# Events Input Probability +event_imgs_prob: 0 \ No newline at end of file diff --git a/configs/model/task/mvs.yaml b/configs/model/task/mvs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8cd2850ba2cad9d925a268639d09a5b19158f323 --- /dev/null +++ b/configs/model/task/mvs.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 diff --git a/configs/model/task/mvs_non_metric.yaml b/configs/model/task/mvs_non_metric.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ba80d5b4694e1c2798191d4949dea22688187185 --- /dev/null +++ b/configs/model/task/mvs_non_metric.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 1 diff --git a/configs/model/task/mvs_training.yaml b/configs/model/task/mvs_training.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eb6b4db744cb16eb7ea351082ffde5a6edd434cb --- /dev/null +++ b/configs/model/task/mvs_training.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale quantities for the input metric high quality gt depth +# If 0, the metric scale quantities will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale quantities for the input metric pose +# If 0, the metric scale quantities will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0.05 diff --git a/configs/model/task/no_images_pass_through.yaml b/configs/model/task/no_images_pass_through.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2e5bf4bf2347c5149a11fe8e02d07fdade002c5b --- /dev/null +++ b/configs/model/task/no_images_pass_through.yaml @@ -0,0 +1,29 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 0 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 1 +# Events Input Probability +event_imgs_prob: 1 \ No newline at end of file diff --git a/configs/model/task/no_images_registration.yaml b/configs/model/task/no_images_registration.yaml new file mode 100644 index 0000000000000000000000000000000000000000..91b40c5b21defd74cd23f5719f7417d9d059a783 --- /dev/null +++ b/configs/model/task/no_images_registration.yaml @@ -0,0 +1,27 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 1 diff --git a/configs/model/task/non_metric_poses_metric_depth.yaml b/configs/model/task/non_metric_poses_metric_depth.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1bc224e89b0df9c44de0676bb4a085c1276abc4a --- /dev/null +++ b/configs/model/task/non_metric_poses_metric_depth.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 1 diff --git a/configs/model/task/non_metric_poses_metric_depth_sparse.yaml b/configs/model/task/non_metric_poses_metric_depth_sparse.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9c92ed95b830a78cf0224dacbb139fcc14e6b2f8 --- /dev/null +++ b/configs/model/task/non_metric_poses_metric_depth_sparse.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 1 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0.9 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 1 diff --git a/configs/model/task/non_metric_poses_non_metric_depth.yaml b/configs/model/task/non_metric_poses_non_metric_depth.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9a5e9902c59379e6e9c8e9b2384a363f58e510bf --- /dev/null +++ b/configs/model/task/non_metric_poses_non_metric_depth.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 1 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 1 diff --git a/configs/model/task/non_metric_poses_non_metric_depth_sparse.yaml b/configs/model/task/non_metric_poses_non_metric_depth_sparse.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5ee97881825a37347ad476f6c35b46639de7e296 --- /dev/null +++ b/configs/model/task/non_metric_poses_non_metric_depth_sparse.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 1 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0.9 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 1 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 1 diff --git a/configs/model/task/pass_through.yaml b/configs/model/task/pass_through.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3d0670997b322d310e03a8cd287eb01140521a5e --- /dev/null +++ b/configs/model/task/pass_through.yaml @@ -0,0 +1,29 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 1 +# Events Input Probability +event_imgs_prob: 1 \ No newline at end of file diff --git a/configs/model/task/posed_sfm.yaml b/configs/model/task/posed_sfm.yaml new file mode 100644 index 0000000000000000000000000000000000000000..df7a1fd65ead009f742a5b14ad6aa3f8383052dc --- /dev/null +++ b/configs/model/task/posed_sfm.yaml @@ -0,0 +1,25 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 0 \ No newline at end of file diff --git a/configs/model/task/posed_sfm_non_metric.yaml b/configs/model/task/posed_sfm_non_metric.yaml new file mode 100644 index 0000000000000000000000000000000000000000..62b540e366656503f984ed7053c453ad7bd1386c --- /dev/null +++ b/configs/model/task/posed_sfm_non_metric.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 0 +# Probability of Geometric Inputs with Depths +depth_prob: 0 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 1 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 1 diff --git a/configs/model/task/registration.yaml b/configs/model/task/registration.yaml new file mode 100644 index 0000000000000000000000000000000000000000..614aff79aa7876480adf8c6f3ce05be96cce3a7e --- /dev/null +++ b/configs/model/task/registration.yaml @@ -0,0 +1,29 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 0 +# Events Input Probability +event_imgs_prob: 0 \ No newline at end of file diff --git a/configs/model/task/registration_doppler.yaml b/configs/model/task/registration_doppler.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f4f8acbdc707efdc69238b63b5bc84fd51d54292 --- /dev/null +++ b/configs/model/task/registration_doppler.yaml @@ -0,0 +1,29 @@ +defaults: + - default + +# Probability of Image Inputs +images_prob: 1 +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 +# Doppler Input Probability +doppler_prob: 1 +# Events Input Probability +event_imgs_prob: 0 \ No newline at end of file diff --git a/configs/model/task/registration_sparse.yaml b/configs/model/task/registration_sparse.yaml new file mode 100644 index 0000000000000000000000000000000000000000..548d43b39ace8f8f70b81e98cc9fd46144a43bc6 --- /dev/null +++ b/configs/model/task/registration_sparse.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 1 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0.9 +# Probability for skipping input of the metric scale norm factor for the input metric high quality gt depth +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0 +# Probability for skipping input of the metric scale norm factor for the input metric pose +# If 0, the metric scale norm factor will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 diff --git a/configs/model/task/registration_training.yaml b/configs/model/task/registration_training.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0676132925de37e46560550a917a83a777662beb --- /dev/null +++ b/configs/model/task/registration_training.yaml @@ -0,0 +1,23 @@ +defaults: + - default + +# Overall Probability of Geometric Inputs +overall_prob: 1 +# Dropout Probability of Geometric Inputs (for each sample across batch size and number of views) +dropout_prob: 0 +# Probability of Geometric Inputs with Ray Directions +ray_dirs_prob: 1 +# Probability of Geometric Inputs with Depths +depth_prob: 1 +# Probability of Geometric Inputs with Camera Poses +cam_prob: 0 +# Probability of sparsely sampling the high quality gt depth +sparse_depth_prob: 0.5 +# Percentage of the valid depth to remove if the probability of using sparse depth is greater than 0 (Range: [0, 1]) +sparsification_removal_percent: 0.9 +# Probability for skipping input of the metric scale quantities for the input metric high quality gt depth +# If 0, the metric scale quantities will be provided as input to the model for all the metric scale conditionings +depth_scale_norm_all_prob: 0.05 +# Probability for skipping input of the metric scale quantities for the input metric pose +# If 0, the metric scale quantities will be provided as input to the model for all the metric scale conditionings +pose_scale_norm_all_prob: 0 diff --git a/configs/train.yaml b/configs/train.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bb2c0ea7e57a0ab3b37c14c0223e7e1b12df6e1e --- /dev/null +++ b/configs/train.yaml @@ -0,0 +1,11 @@ +defaults: + - machine: psc + - model: default + - _self_ + +output_dir: ${hydra:run.dir} +root_data_dir: ${machine.root_data_dir} +mapanything_dataset_metadata_dir: ${machine.mapanything_dataset_metadata_dir} +root_pretrained_checkpoints_dir: ${machine.root_pretrained_checkpoints_dir} +root_experiments_dir: ${machine.root_experiments_dir} +root_uniception_pretrained_checkpoints_dir: ${machine.root_uniception_pretrained_checkpoints_dir} diff --git a/examples/Bigfoot.jpg b/examples/Bigfoot.jpg new file mode 100644 index 0000000000000000000000000000000000000000..601c0c5ce17cac5766648468934dca40f82ed1dd --- /dev/null +++ b/examples/Bigfoot.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fb0c8686a970c694fa5f0e3474b818995a54ca23ff55fedfe3e7fdbbc0e8bc9b +size 414892 diff --git a/examples/Bigfoot.mp4 b/examples/Bigfoot.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..78e8f4ffb7acebfb14cdb045f41609c3a4c95aef --- /dev/null +++ b/examples/Bigfoot.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dd61a215065c55de733ab44270c56dee2ee9999c6dedd17f5ab709d3e33ddb95 +size 6563963 diff --git a/examples/Humanoid_Walking.jpg b/examples/Humanoid_Walking.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2a7967cfc1940f9b60121a8f2c8f93516db04253 --- /dev/null +++ b/examples/Humanoid_Walking.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d71e12340388385d3265ccf6069935e946555a761b1569a62c13ec2c8960e86b +size 92285 diff --git a/examples/Humanoid_Walking.mp4 b/examples/Humanoid_Walking.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..c4b355e73ac7c58af55d8980452dfb054074dee3 --- /dev/null +++ b/examples/Humanoid_Walking.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:179b1e573a61e8c7ff9d260418a55ffacf740dcb11fe765d3d2ae3bd8631747c +size 757643 diff --git a/examples/car.jpg b/examples/car.jpg new file mode 100644 index 0000000000000000000000000000000000000000..37428aee63a5ceee0794b8afa0a13bfe568456c4 --- /dev/null +++ b/examples/car.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4c17b878ec315f80621f9eacd9b77e2231f3ece01703da639b3d46601f33060b +size 86859 diff --git a/examples/car.mp4 b/examples/car.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..23914d4baf302c9688269b6595d02b330dcf1177 --- /dev/null +++ b/examples/car.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c7c5e4aeaa18abd4ef5a8c8551e4c9d2139719531dc64fb5e315569776000251 +size 563979 diff --git a/examples/cat_jump.jpg b/examples/cat_jump.jpg new file mode 100644 index 0000000000000000000000000000000000000000..20eef89b1435b100f9ffb8e6a991838d4e63997b --- /dev/null +++ b/examples/cat_jump.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:44b269d92cbd12ee12379cd976d17af524d86222b490cff872e3e667cc339b1d +size 64294 diff --git a/examples/cat_jump.mp4 b/examples/cat_jump.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..6275075750f9e2da4c41f47fc04eeedc8cd1714b --- /dev/null +++ b/examples/cat_jump.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:698d7abb32d48b57c044ed939cd43ae68c62d9479f641b431fbd325a584e54cd +size 306336 diff --git a/examples/climbing.jpg b/examples/climbing.jpg new file mode 100644 index 0000000000000000000000000000000000000000..bbbe206bc876d97a6236e672c1cbbac543c5781d --- /dev/null +++ b/examples/climbing.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4cb3dc93e3aba72de7f697dfe2f958b2b4a8152af28a417f26eb178d87024add +size 100154 diff --git a/examples/climbing.mp4 b/examples/climbing.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..5ce6e1e390dbd2300fc1193af5bbad425a66b732 --- /dev/null +++ b/examples/climbing.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c88888ba501b13fa575f26d488eed4e3b81771ca19ef3275728473c0e8fec8e1 +size 1051826 diff --git a/examples/humanoid_football.mp4 b/examples/humanoid_football.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..3ea3b35a0a7a845f8b1535b1ac8dd31dbb95bbdb --- /dev/null +++ b/examples/humanoid_football.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7a02a31db5a1bbab2c75241374f607e4e050b9da54188e1e9b44b7b30ac92a91 +size 186537 diff --git a/examples/humanoid_football.png b/examples/humanoid_football.png new file mode 100644 index 0000000000000000000000000000000000000000..d7bf423db91aa7f077ecabaf8a8c42cae827b5dd --- /dev/null +++ b/examples/humanoid_football.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d1048342f1fafd70969560737839773ac141cdbaef6b46658d8b5ae3512f05e7 +size 176950 diff --git a/examples/lady-running.jpg b/examples/lady-running.jpg new file mode 100644 index 0000000000000000000000000000000000000000..75fbbe49b64f5929c0236a1917936da7d0f52d3c --- /dev/null +++ b/examples/lady-running.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c927d4b44c3acc353b1e6137e31c93b2080b01e52fe89599148387dd0071d29a +size 116088 diff --git a/examples/lady-running.mp4 b/examples/lady-running.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..598a21ba6b2da2b548459029a2ab9c63f2064d0c --- /dev/null +++ b/examples/lady-running.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9aab52b274f749be3d0f7cb4d9182cbce0c339e0a14d3579432d8ccd73663581 +size 707603 diff --git a/examples/layup_fpv.jpg b/examples/layup_fpv.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6a9d90acfb06ca40fb4f1732f2edad1f4441902b --- /dev/null +++ b/examples/layup_fpv.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01ff9fca65fbec29e23dd448bb7e4bdcbf598ddd11bed4f8d38df8898125a210 +size 424392 diff --git a/examples/layup_fpv.mp4 b/examples/layup_fpv.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..f5988aa9d29e999ee1eb933ee0e72bf2ba6cfd51 --- /dev/null +++ b/examples/layup_fpv.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:25d0e099ae0143b52fa0ee2d5ad298c50e1f5bf4a010d03f245b16a90fbb69bf +size 1967467 diff --git a/examples/layup_fpv.png b/examples/layup_fpv.png new file mode 100644 index 0000000000000000000000000000000000000000..45edb756af77bf09fb2fd26a3b091fb87351ae40 --- /dev/null +++ b/examples/layup_fpv.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:38cd8a86364b07fe2da5f916fc930570183922a9c0fbf6361536d1fb6527c430 +size 1315754 diff --git a/examples/mobile_robot.jpg b/examples/mobile_robot.jpg new file mode 100644 index 0000000000000000000000000000000000000000..966ebb66d0140b3d6ade3f61d7463dc7661b0347 --- /dev/null +++ b/examples/mobile_robot.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:79989f79c02130836e72a533da386644211a4e5b7045bc9b40ec5b65fee0121a +size 130773 diff --git a/examples/mobile_robot.mp4 b/examples/mobile_robot.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..1c23e0ef52e66f7ee3329d46d92f8900c862e016 --- /dev/null +++ b/examples/mobile_robot.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:326e03e677eee8d416eda6991891c05b260467ea6589d3ee3ba7f434652a1791 +size 1549967 diff --git a/examples/motorbike_desert.jpg b/examples/motorbike_desert.jpg new file mode 100644 index 0000000000000000000000000000000000000000..abde9ba0f6525617c3ee96b3bee2151f87e4c463 --- /dev/null +++ b/examples/motorbike_desert.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:47ea7b1b9f4d7033f26c3ef045ac745810cc83ac73cd2b3563fc60403dbfbd75 +size 422770 diff --git a/examples/motorbike_desert.mp4 b/examples/motorbike_desert.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..e462441a54368d23614a3b870d1b5f1bf84f4758 --- /dev/null +++ b/examples/motorbike_desert.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d3a2e4abd01da27a3bf4ec333e10c2e56bf419fbaac2b32e66db088b52ec1fb9 +size 2186665 diff --git a/examples/quadruped_climbing.jpg b/examples/quadruped_climbing.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2ffa4c91bbb6737286cfac13f49586e2c27ab52f --- /dev/null +++ b/examples/quadruped_climbing.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:57028f54d8954fbf1ed23128f3b580990ae07da3c2385f5f7beb5e38f1a01bb0 +size 78385 diff --git a/examples/quadruped_climbing.mp4 b/examples/quadruped_climbing.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..7dcfd1933a4c446d2b70adbd65da8a6c49917739 --- /dev/null +++ b/examples/quadruped_climbing.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bd04308999bf688f4707f5369fd78b74433e5f498c44210dd7c7b7891fd2e95a +size 992633 diff --git a/examples/seagull.jpg b/examples/seagull.jpg new file mode 100644 index 0000000000000000000000000000000000000000..061f06e7e7f84d17fd961c61a1736d1b256b1f1e --- /dev/null +++ b/examples/seagull.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6099f8378d1204b5c3dd897d26a010912a5c9d7c47d680b89cd6042635b5eb9c +size 7992 diff --git a/examples/seagull.mp4 b/examples/seagull.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..bfe5a4ab3b883ce454cdcd022dc785a6dcf369cb --- /dev/null +++ b/examples/seagull.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a9fcaccd3eaa2d004ab83db1737dedcd1d5fa92b76d3534e8bbb201a79f74064 +size 67881 diff --git a/mapanything/__init__.py b/mapanything/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mapanything/__pycache__/__init__.cpython-312.pyc b/mapanything/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b59ba50b47f97f129108d805dc0b3248258f005 Binary files /dev/null and b/mapanything/__pycache__/__init__.cpython-312.pyc differ diff --git a/mapanything/models/__init__.py b/mapanything/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8240fba091bf86ec85ad26997c047bfa5e14ff57 --- /dev/null +++ b/mapanything/models/__init__.py @@ -0,0 +1,136 @@ +""" +Model Factory for MapAnything +""" + +import importlib.util +import logging +import warnings + +import numpy as np +from omegaconf import DictConfig, OmegaConf + +# Core models that are always available +from mapanything.models.any4d.models.model import Any4D + +# Suppress DINOv2 warnings +logging.getLogger("dinov2").setLevel(logging.WARNING) +warnings.filterwarnings("ignore", message="xFormers is available", category=UserWarning) +warnings.filterwarnings( + "ignore", message="xFormers is not available", category=UserWarning +) + + +def resolve_special_float(value): + if value == "inf": + return np.inf + elif value == "-inf": + return -np.inf + else: + raise ValueError(f"Unknown special float value: {value}") + + +def init_model( + model_str: str, model_config: DictConfig, torch_hub_force_reload: bool = False +): + """ + Initialize a model using OmegaConf configuration. + + Args: + model_str (str): Name of the model class to create. + model_config (DictConfig): OmegaConf model configuration. + torch_hub_force_reload (bool): Whether to force reload relevant parts of the model from torch hub. + """ + if not OmegaConf.has_resolver("special_float"): + OmegaConf.register_new_resolver("special_float", resolve_special_float) + model_dict = OmegaConf.to_container(model_config, resolve=True) + model = model_factory( + model_str, torch_hub_force_reload=torch_hub_force_reload, **model_dict + ) + + return model + + +# Define model configurations with import paths +MODEL_CONFIGS = { + "any4d": { + "class": Any4D, + }, +} + + +def check_module_exists(module_path): + """ + Check if a module can be imported without actually importing it. + + Args: + module_path (str): The path to the module to check. + + Returns: + bool: True if the module can be imported, False otherwise. + """ + return importlib.util.find_spec(module_path) is not None + + +def model_factory(model_str: str, **kwargs): + """ + Model factory for MapAnything. + + Args: + model_str (str): Name of the model to create. + **kwargs: Additional keyword arguments to pass to the model constructor. + + Returns: + nn.Module: An instance of the specified model. + """ + if model_str not in MODEL_CONFIGS: + raise ValueError( + f"Unknown model: {model_str}. Valid options are: {', '.join(MODEL_CONFIGS.keys())}" + ) + + model_config = MODEL_CONFIGS[model_str] + + # Handle core models directly + if "class" in model_config: + model_class = model_config["class"] + # Handle external models with dynamic imports + elif "module" in model_config: + module_path = model_config["module"] + class_name = model_config["class_name"] + + # Check if the module can be imported + if not check_module_exists(module_path): + raise ImportError( + f"Model '{model_str}' requires module '{module_path}' which is not installed. " + f"Please install the corresponding submodule or package." + ) + + # Dynamically import the module and get the class + try: + module = importlib.import_module(module_path) + model_class = getattr(module, class_name) + except (ImportError, AttributeError) as e: + raise ImportError( + f"Failed to import {class_name} from {module_path}: {str(e)}" + ) + else: + raise ValueError(f"Invalid model configuration for {model_str}") + + print(f"Initializing {model_class} with kwargs: {kwargs}") + if model_str != "org_dust3r": + return model_class(**kwargs) + else: + eval_str = kwargs.get("model_eval_str", None) + return eval(eval_str) + + +def get_available_models() -> list: + """ + Get a list of available models in MapAnything. + + Returns: + list: A list of available model names. + """ + return list(MODEL_CONFIGS.keys()) + + +__all__ = ["model_factory", "get_available_models"] diff --git a/mapanything/models/__pycache__/__init__.cpython-312.pyc b/mapanything/models/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a511ec9ea9aee9e79811692f18e94f6aebec2d7 Binary files /dev/null and b/mapanything/models/__pycache__/__init__.cpython-312.pyc differ diff --git a/mapanything/models/any4d/models/__pycache__/model.cpython-312.pyc b/mapanything/models/any4d/models/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df11a8a80b463b33b5d737a1db4c156b256d9ff5 Binary files /dev/null and b/mapanything/models/any4d/models/__pycache__/model.cpython-312.pyc differ diff --git a/mapanything/models/any4d/models/model.py b/mapanything/models/any4d/models/model.py new file mode 100644 index 0000000000000000000000000000000000000000..53760dee90c50fdaf9f0bda2934f9f890a8c53ce --- /dev/null +++ b/mapanything/models/any4d/models/model.py @@ -0,0 +1,2143 @@ +""" +MapAnything model class defined using UniCeption modules. +""" + +from functools import partial +from typing import Callable, Dict, Type, Union + +import torch +import torch.nn as nn + +from mapanything.utils.geometry import ( + apply_log_to_norm, + convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap, + normalize_depth_using_non_zero_pixels, + normalize_pose_translations, + normalize_multiple_pointclouds, + transform_pose_using_quats_and_trans_2_to_1, +) +from uniception.models.encoders import ( + encoder_factory, + EncoderGlobalRepInput, + ViTEncoderInput, + ViTEncoderNonImageInput, +) +from uniception.models.info_sharing.alternating_attention_transformer import ( + MultiViewAlternatingAttentionTransformer, + MultiViewAlternatingAttentionTransformerIFR, +) +from uniception.models.info_sharing.base import MultiViewTransformerInput +from uniception.models.info_sharing.cross_attention_transformer import ( + MultiViewCrossAttentionTransformer, + MultiViewCrossAttentionTransformerIFR, +) +from uniception.models.info_sharing.global_attention_transformer import ( + MultiViewGlobalAttentionTransformer, + MultiViewGlobalAttentionTransformerIFR, +) +from uniception.models.prediction_heads.adaptors import ( + CamTranslationPlusQuatsAdaptor, + PointMapAdaptor, + PointMapPlusRayDirectionsPlusDepthAdaptor, + PointMapPlusRayDirectionsPlusDepthWithConfidenceAdaptor, + PointMapPlusRayDirectionsPlusDepthWithConfidenceAndMaskAdaptor, + PointMapPlusRayDirectionsPlusDepthWithMaskAdaptor, + PointMapWithConfidenceAdaptor, + PointMapWithConfidenceAndMaskAdaptor, + PointMapWithMaskAdaptor, + RayDirectionsPlusDepthAdaptor, + RayDirectionsPlusDepthWithConfidenceAdaptor, + RayDirectionsPlusDepthWithConfidenceAndMaskAdaptor, + RayDirectionsPlusDepthWithMaskAdaptor, + RayMapPlusDepthAdaptor, + RayMapPlusDepthWithConfidenceAdaptor, + RayMapPlusDepthWithConfidenceAndMaskAdaptor, + RayMapPlusDepthWithMaskAdaptor, + ScaleAdaptor, + SceneFlowAdaptor, +) + +from uniception.models.prediction_heads.base import ( + AdaptorInput, + PredictionHeadInput, + PredictionHeadLayeredInput, + PredictionHeadTokenInput, +) +from uniception.models.prediction_heads.dpt import DPTFeature, DPTRegressionProcessor +from uniception.models.prediction_heads.linear import LinearFeature +from uniception.models.prediction_heads.mlp_head import MLPHead +from uniception.models.prediction_heads.pose_head import PoseHead + +# Enable TF32 precision if supported (for GPU >= Ampere and PyTorch >= 1.12) +if hasattr(torch.backends.cuda, "matmul") and hasattr( + torch.backends.cuda.matmul, "allow_tf32" +): + torch.backends.cuda.matmul.allow_tf32 = True + + +class Any4D(nn.Module): + "Modular Any4D model class that supports input of images & optional geometric modalities (multiple reconstruction tasks)." + + def __init__( + self, + name: str, + encoder_config: Dict, + info_sharing_config: Dict, + pred_head_config: Dict, + scene_flow_pred_head_config: Dict, + geometric_input_config: Dict, + fusion_norm_layer: Union[Type[nn.Module], Callable[..., nn.Module]] = partial( + nn.LayerNorm, eps=1e-6 + ), + pretrained_checkpoint_path: str = None, + load_specific_pretrained_submodules: bool = False, + specific_pretrained_submodules: list = None, + torch_hub_force_reload: bool = False, + ): + """ + Multi-view model containing an image encoder fused with optional geometric modalities followed by a multi-view attention transformer and respective downstream heads. + The goal is to output scene representation. + The multi-view attention transformer also takes as input a scale token to predict the metric scaling factor for the predicted scene representation. + + Args: + name (str): Name of the model. + encoder_config (Dict): Configuration for the encoder. + info_sharing_config (Dict): Configuration for the multi-view attention transformer. + pred_head_config (Dict): Configuration for the prediction heads. + geometric_input_config (Dict): Configuration for the input of optional geometric modalities. + fusion_norm_layer (Union[Type[nn.Module], Callable[..., nn.Module]]): Normalization layer to use after fusion (addition) of encoder and geometric modalities. (default: partial(nn.LayerNorm, eps=1e-6)) + pretrained_checkpoint_path (str): Path to pretrained checkpoint. (default: None) + load_specific_pretrained_submodules (bool): Whether to load specific pretrained submodules. (default: False) + specific_pretrained_submodules (list): List of specific pretrained submodules to load. Must be provided when load_specific_pretrained_submodules is True. (default: None) + torch_hub_force_reload (bool): Whether to force reload the encoder from torch hub. (default: False) + """ + super().__init__() + + # Initalize the attributes + self.name = name + self.encoder_config = encoder_config + self.info_sharing_config = info_sharing_config + self.pred_head_config = pred_head_config + self.scene_flow_pred_head_config = scene_flow_pred_head_config + self.geometric_input_config = geometric_input_config + self.pretrained_checkpoint_path = pretrained_checkpoint_path + self.load_specific_pretrained_submodules = load_specific_pretrained_submodules + self.specific_pretrained_submodules = specific_pretrained_submodules + self.torch_hub_force_reload = torch_hub_force_reload + self.class_init_args = { + "name": self.name, + "encoder_config": self.encoder_config, + "info_sharing_config": self.info_sharing_config, + "pred_head_config": self.pred_head_config, + "scene_flow_pred_head_config": self.scene_flow_pred_head_config, + "geometric_input_config": self.geometric_input_config, + "pretrained_checkpoint_path": self.pretrained_checkpoint_path, + "load_specific_pretrained_submodules": self.load_specific_pretrained_submodules, + "specific_pretrained_submodules": self.specific_pretrained_submodules, + "torch_hub_force_reload": self.torch_hub_force_reload, + } + + # Get relevant parameters from the configs + self.info_sharing_type = info_sharing_config["model_type"] + self.info_sharing_return_type = info_sharing_config["model_return_type"] + self.pred_head_type = pred_head_config["type"] + self.scene_flow_pred_head_type = scene_flow_pred_head_config["type"] + + # Initialize image encoder + if self.encoder_config["uses_torch_hub"]: + self.encoder_config["torch_hub_force_reload"] = torch_hub_force_reload + del self.encoder_config["uses_torch_hub"] + self.encoder = encoder_factory(**self.encoder_config) + + # Initialize the encoder for ray directions + ray_dirs_encoder_config = self.geometric_input_config["ray_dirs_encoder_config"] + ray_dirs_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim + ray_dirs_encoder_config["patch_size"] = self.encoder.patch_size + self.ray_dirs_encoder = encoder_factory(**ray_dirs_encoder_config) + + # Initialize the encoder for depth (normalized per view and values after normalization are scaled logarithmically) + depth_encoder_config = self.geometric_input_config["depth_encoder_config"] + depth_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim + depth_encoder_config["patch_size"] = self.encoder.patch_size + self.depth_encoder = encoder_factory(**depth_encoder_config) + + # Initialize the encoder for log scale factor of depth + depth_scale_encoder_config = self.geometric_input_config["scale_encoder_config"] + depth_scale_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim + self.depth_scale_encoder = encoder_factory(**depth_scale_encoder_config) + + # Initialize the encoder for camera rotation + cam_rot_encoder_config = self.geometric_input_config["cam_rot_encoder_config"] + cam_rot_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim + self.cam_rot_encoder = encoder_factory(**cam_rot_encoder_config) + + # Initialize the encoder for camera translation (normalized across all provided camera translations) + cam_trans_encoder_config = self.geometric_input_config[ + "cam_trans_encoder_config" + ] + cam_trans_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim + self.cam_trans_encoder = encoder_factory(**cam_trans_encoder_config) + + # Initialize the encoder for log scale factor of camera translation + cam_trans_scale_encoder_config = self.geometric_input_config[ + "scale_encoder_config" + ] + cam_trans_scale_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim + self.cam_trans_scale_encoder = encoder_factory(**cam_trans_scale_encoder_config) + + # Initialize the encoder for scene flow + scene_flow_encoder_config = self.geometric_input_config["scene_flow_encoder_config"] + scene_flow_encoder_config["enc_embed_dim"] = self.encoder.enc_embed_dim + self.scene_flow_encoder = encoder_factory(**scene_flow_encoder_config) + + # Initialize the fusion norm layer + self.fusion_norm_layer = fusion_norm_layer(self.encoder.enc_embed_dim) + + # Initialize the Scale Token + # Used to scale the final scene predictions to metric scale + # During inference extended to (B, C, T), where T is the number of tokens (i.e., 1) + self.scale_token = nn.Parameter(torch.zeros(self.encoder.enc_embed_dim)) + torch.nn.init.trunc_normal_(self.scale_token, std=0.02) + + # Initialize the info sharing module (multi-view transformer) + self._initialize_info_sharing(info_sharing_config) + + # Initialize the prediction heads + self._initialize_prediction_heads(pred_head_config, scene_flow_pred_head_config) + + # Initialize the final adaptors + self._initialize_adaptors(pred_head_config, scene_flow_pred_head_config) + + # Load pretrained weights + self._load_pretrained_weights() + + def _initialize_info_sharing(self, info_sharing_config): + """ + Initialize the information sharing module based on the configuration. + + This method sets up the custom positional encoding if specified and initializes + the appropriate multi-view transformer based on the configuration type. + + Args: + info_sharing_config (Dict): Configuration for the multi-view attention transformer. + Should contain 'custom_positional_encoding', 'model_type', and 'model_return_type'. + + Returns: + None + + Raises: + ValueError: If invalid configuration options are provided. + """ + # Initialize Custom Positional Encoding if required + custom_positional_encoding = info_sharing_config["custom_positional_encoding"] + if custom_positional_encoding is not None: + if isinstance(custom_positional_encoding, str): + print( + f"Using custom positional encoding for multi-view attention transformer: {custom_positional_encoding}" + ) + raise ValueError( + f"Invalid custom_positional_encoding: {custom_positional_encoding}. None implemented." + ) + elif isinstance(custom_positional_encoding, Callable): + print( + "Using callable function as custom positional encoding for multi-view attention transformer." + ) + self.custom_positional_encoding = custom_positional_encoding + else: + self.custom_positional_encoding = None + + # Add dependecies to info_sharing_config + info_sharing_config["module_args"]["input_embed_dim"] = ( + self.encoder.enc_embed_dim + ) + info_sharing_config["module_args"]["custom_positional_encoding"] = ( + self.custom_positional_encoding + ) + + # Initialize Multi-View Transformer + if self.info_sharing_return_type == "no_intermediate_features": + # Returns only normalized last layer features + # Intialize multi-view transformer based on type + if self.info_sharing_type == "cross_attention": + self.info_sharing = MultiViewCrossAttentionTransformer( + **info_sharing_config["module_args"] + ) + elif self.info_sharing_type == "global_attention": + self.info_sharing = MultiViewGlobalAttentionTransformer( + **info_sharing_config["module_args"] + ) + elif self.info_sharing_type == "alternating_attention": + self.info_sharing = MultiViewAlternatingAttentionTransformer( + **info_sharing_config["module_args"] + ) + else: + raise ValueError( + f"Invalid info_sharing_type: {self.info_sharing_type}. Valid options: ['cross_attention', 'global_attention', 'alternating_attention']" + ) + elif self.info_sharing_return_type == "intermediate_features": + # Returns intermediate features and normalized last layer features + # Initialize mulit-view transformer based on type + if self.info_sharing_type == "cross_attention": + self.info_sharing = MultiViewCrossAttentionTransformerIFR( + **info_sharing_config["module_args"] + ) + elif self.info_sharing_type == "global_attention": + self.info_sharing = MultiViewGlobalAttentionTransformerIFR( + **info_sharing_config["module_args"] + ) + elif self.info_sharing_type == "alternating_attention": + self.info_sharing = MultiViewAlternatingAttentionTransformerIFR( + **info_sharing_config["module_args"] + ) + else: + raise ValueError( + f"Invalid info_sharing_type: {self.info_sharing_type}. Valid options: ['cross_attention', 'global_attention', 'alternating_attention']" + ) + # Assess if the DPT needs to use encoder features + if len(self.info_sharing.indices) == 2: + self.use_encoder_features_for_dpt = True + elif len(self.info_sharing.indices) == 3: + self.use_encoder_features_for_dpt = False + else: + raise ValueError( + "Invalid number of indices provided for info sharing feature returner. Please provide 2 or 3 indices." + ) + else: + raise ValueError( + f"Invalid info_sharing_return_type: {self.info_sharing_return_type}. Valid options: ['no_intermediate_features', 'intermediate_features']" + ) + + def _initialize_prediction_heads(self, pred_head_config, scene_flow_pred_head_config): + """ + Initialize the prediction heads based on the prediction head configuration. + + This method configures and initializes the appropriate prediction heads based on the + specified prediction head type (linear, DPT, or DPT+pose). It sets up the necessary + dependencies and creates the required model components. + + Args: + pred_head_config (Dict): Configuration for the geometry prediction heads. + scene_flow_pred_head_config (Dict): Configuration for the scene flow prediction heads. + + Returns: + None + + Raises: + ValueError: If an invalid pred_head_type is provided. + """ + # Add dependencies to prediction head config + pred_head_config["feature_head"]["patch_size"] = self.encoder.patch_size + if self.pred_head_type == "linear": + pred_head_config["feature_head"]["input_feature_dim"] = ( + self.info_sharing.dim + ) + scene_flow_pred_head_config["feature_head"]["input_feature_dim"] = ( + self.info_sharing.dim + ) + elif "dpt" in self.pred_head_type: + # Add dependencies for DPT & Regressor head + if self.use_encoder_features_for_dpt: + pred_head_config["feature_head"]["input_feature_dims"] = [ + self.encoder.enc_embed_dim + ] + [self.info_sharing.dim] * 3 + scene_flow_pred_head_config["feature_head"]["input_feature_dims"] = [ + self.encoder.enc_embed_dim + ] + [self.info_sharing.dim] * 3 + else: + pred_head_config["feature_head"]["input_feature_dims"] = [ + self.info_sharing.dim + ] * 4 + scene_flow_pred_head_config["feature_head"]["input_feature_dims"] = [ + self.info_sharing.dim + ] * 4 + pred_head_config["regressor_head"]["input_feature_dim"] = pred_head_config[ + "feature_head" + ]["feature_dim"] + scene_flow_pred_head_config["regressor_head"]["input_feature_dim"] = scene_flow_pred_head_config[ + "feature_head" + ]["feature_dim"] + # Add dependencies for Pose head if required + if "pose" in self.pred_head_type: + pred_head_config["pose_head"]["patch_size"] = self.encoder.patch_size + # scene_flow_pred_head_config["pose_head"]["patch_size"] = self.encoder.patch_size + pred_head_config["pose_head"]["input_feature_dim"] = ( + self.info_sharing.dim + ) + # scene_flow_pred_head_config["pose_head"]["input_feature_dim"] = ( + # self.info_sharing.dim + # ) + else: + raise ValueError( + f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" + ) + pred_head_config["scale_head"]["input_feature_dim"] = self.info_sharing.dim + # scene_flow_pred_head_config["scale_head"]["input_feature_dim"] = self.info_sharing.dim + + # Initialize Prediction Heads + if self.pred_head_type == "linear": + # Initialize Dense Prediction Head for all views + self.dense_head = LinearFeature(**pred_head_config["feature_head"]) + self.scene_flow_dense_head = LinearFeature(**scene_flow_pred_head_config["feature_head"]) + elif "dpt" in self.pred_head_type: + # Initialze Dense Prediction Head for all views + self.dpt_feature_head = DPTFeature(**pred_head_config["feature_head"]) + self.dpt_regressor_head = DPTRegressionProcessor( + **pred_head_config["regressor_head"] + ) + self.dense_head = nn.Sequential( + self.dpt_feature_head, self.dpt_regressor_head + ) + + self.scene_flow_dpt_feature_head = DPTFeature(**scene_flow_pred_head_config["feature_head"]) + self.scene_flow_dpt_regressor_head = DPTRegressionProcessor( + **scene_flow_pred_head_config["regressor_head"] + ) + self.scene_flow_dense_head = nn.Sequential( + self.scene_flow_dpt_feature_head, self.scene_flow_dpt_regressor_head + ) + # Initialize Pose Head for all views if required + if "pose" in self.pred_head_type: + self.pose_head = PoseHead(**pred_head_config["pose_head"]) + else: + raise ValueError( + f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" + ) + self.scale_head = MLPHead(**pred_head_config["scale_head"]) + + + def _initialize_adaptors(self, pred_head_config, scene_flow_pred_head_config): + """ + Initialize the adaptors based on the prediction head configuration. + + This method sets up the appropriate adaptors for different scene representation types, + such as pointmaps, ray maps with depth, or ray directions with depth and pose. + + Args: + pred_head_config (Dict): Configuration for the prediction heads including adaptor type. + scene_flow_pred_head_config (Dict): Configuration for the scene flow prediction heads. + + Returns: + None + + Raises: + ValueError: If an invalid adaptor_type is provided. + AssertionError: If ray directions + depth + pose is used with an incompatible head type. + """ + if pred_head_config["adaptor_type"] == "pointmap": + self.dense_adaptor = PointMapAdaptor(**pred_head_config["adaptor"]) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "pointmap+scene_flow" + elif pred_head_config["adaptor_type"] == "pointmap+confidence": + self.dense_adaptor = PointMapWithConfidenceAdaptor( + **pred_head_config["adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "pointmap+scene_flow+confidence" + elif pred_head_config["adaptor_type"] == "pointmap+mask": + self.dense_adaptor = PointMapWithMaskAdaptor(**pred_head_config["adaptor"]) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "pointmap+scene_flow+mask" + elif pred_head_config["adaptor_type"] == "pointmap+confidence+mask": + self.dense_adaptor = PointMapWithConfidenceAndMaskAdaptor( + **pred_head_config["adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "pointmap+scene_flow+confidence+mask" + elif pred_head_config["adaptor_type"] == "raymap+depth": + self.dense_adaptor = RayMapPlusDepthAdaptor(**pred_head_config["adaptor"]) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "raymap+depth+scene_flow" + elif pred_head_config["adaptor_type"] == "raymap+depth+confidence": + self.dense_adaptor = RayMapPlusDepthWithConfidenceAdaptor( + **pred_head_config["adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "raymap+depth+scene_flow+confidence" + elif pred_head_config["adaptor_type"] == "raymap+depth+mask": + self.dense_adaptor = RayMapPlusDepthWithMaskAdaptor( + **pred_head_config["adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "raymap+depth+scene_flow+mask" + elif pred_head_config["adaptor_type"] == "raymap+depth+confidence+mask": + self.dense_adaptor = RayMapPlusDepthWithConfidenceAndMaskAdaptor( + **pred_head_config["adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.scene_rep_type = "raymap+depth+scene_flow+confidence+mask" + elif pred_head_config["adaptor_type"] == "raydirs+depth+pose": + assert self.pred_head_type == "dpt+pose", ( + "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = RayDirectionsPlusDepthAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "raydirs+depth+pose+scene_flow" + elif pred_head_config["adaptor_type"] == "raydirs+depth+pose+confidence": + assert self.pred_head_type == "dpt+pose", ( + "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = RayDirectionsPlusDepthWithConfidenceAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "raydirs+depth+pose+scene_flow+confidence" + elif pred_head_config["adaptor_type"] == "raydirs+depth+pose+mask": + assert self.pred_head_type == "dpt+pose", ( + "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = RayDirectionsPlusDepthWithMaskAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "raydirs+depth+pose+scene_flow+mask" + elif pred_head_config["adaptor_type"] == "raydirs+depth+pose+confidence+mask": + assert self.pred_head_type == "dpt+pose", ( + "Ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = RayDirectionsPlusDepthWithConfidenceAndMaskAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "raydirs+depth+pose+scene_flow+confidence+mask" + elif pred_head_config["adaptor_type"] == "campointmap+pose": + assert self.pred_head_type == "dpt+pose", ( + "Camera pointmap + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = PointMapAdaptor(**pred_head_config["dpt_adaptor"]) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "campointmap+pose+scene_flow" + elif pred_head_config["adaptor_type"] == "campointmap+pose+confidence": + assert self.pred_head_type == "dpt+pose", ( + "Camera pointmap + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = PointMapWithConfidenceAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "campointmap+pose+scene_flow+confidence" + elif pred_head_config["adaptor_type"] == "campointmap+pose+mask": + assert self.pred_head_type == "dpt+pose", ( + "Camera pointmap + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = PointMapWithMaskAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "campointmap+pose+scene_flow+mask" + elif pred_head_config["adaptor_type"] == "campointmap+pose+confidence+mask": + assert self.pred_head_type == "dpt+pose", ( + "Camera pointmap + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = PointMapWithConfidenceAndMaskAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "campointmap+pose+scene_flow+confidence+mask" + elif pred_head_config["adaptor_type"] == "pointmap+raydirs+depth+pose": + assert self.pred_head_type == "dpt+pose", ( + "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = PointMapPlusRayDirectionsPlusDepthAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "pointmap+raydirs+depth+pose+scene_flow" + elif ( + pred_head_config["adaptor_type"] == "pointmap+raydirs+depth+pose+confidence" + ): + assert self.pred_head_type == "dpt+pose", ( + "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = ( + PointMapPlusRayDirectionsPlusDepthWithConfidenceAdaptor( + **pred_head_config["dpt_adaptor"] + ) + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "pointmap+raydirs+depth+pose+scene_flow+confidence" + elif pred_head_config["adaptor_type"] == "pointmap+raydirs+depth+pose+mask": + assert self.pred_head_type == "dpt+pose", ( + "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = PointMapPlusRayDirectionsPlusDepthWithMaskAdaptor( + **pred_head_config["dpt_adaptor"] + ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "pointmap+raydirs+depth+pose+scene_flow+mask" + elif ( + pred_head_config["adaptor_type"] + == "pointmap+raydirs+depth+pose+confidence+mask" + ): + assert self.pred_head_type == "dpt+pose", ( + "Pointmap + ray directions + depth + pose can only be used as scene representation with dpt + pose head." + ) + self.dense_adaptor = ( + PointMapPlusRayDirectionsPlusDepthWithConfidenceAndMaskAdaptor( + **pred_head_config["dpt_adaptor"] + ) + ) + # self.dense_adaptor = ( + # RayDirectionsPlusDepthWithConfidenceAndMaskAdaptor( + # **pred_head_config["dpt_adaptor"] + # ) + # ) + self.scene_flow_dense_adaptor = SceneFlowAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + # self.scene_flow_dense_adaptor = SceneFlowPlusPointMapAdaptor(**scene_flow_pred_head_config["dpt_adaptor"]) + self.pose_adaptor = CamTranslationPlusQuatsAdaptor( + **pred_head_config["pose_adaptor"] + ) + self.scene_rep_type = "pointmap+raydirs+depth+pose+scene_flow+confidence+mask" + else: + raise ValueError( + f"Invalid adaptor_type: {pred_head_config['adaptor_type']}. \ + Valid options: ['pointmap', 'raymap+depth', 'raydirs+depth+pose', 'campointmap+pose', 'pointmap+raydirs+depth+pose' \ + 'pointmap+confidence', 'raymap+depth+confidence', 'raydirs+depth+pose+confidence', 'campointmap+pose+confidence', 'pointmap+raydirs+depth+pose+confidence' \ + 'pointmap+mask', 'raymap+depth+mask', 'raydirs+depth+pose+mask', 'campointmap+pose+mask', 'pointmap+raydirs+depth+pose+mask' \ + 'pointmap+confidence+mask', 'raymap+depth+confidence+mask', 'raydirs+depth+pose+confidence+mask', 'campointmap+pose+confidence+mask', 'pointmap+raydirs+depth+pose+confidence+mask']" + ) + self.scale_adaptor = ScaleAdaptor(**pred_head_config["scale_adaptor"]) + + + def _load_pretrained_weights(self): + """ + Load pretrained weights from a checkpoint file. + + If load_specific_pretrained_submodules is True, only loads weights for the specified submodules. + Otherwise, loads all weights from the checkpoint. + + Returns: + None + """ + if self.pretrained_checkpoint_path is not None: + if not self.load_specific_pretrained_submodules: + print( + f"Loading pretrained MapAnything weights from {self.pretrained_checkpoint_path} ..." + ) + # ckpt = torch.load(self.pretrained_checkpoint_path, weights_only=False) + ckpt = torch.load(self.pretrained_checkpoint_path, map_location="cpu", weights_only=False) + + # Remove the multi modal encoder weights if they are present in the checkpoint + remove_multimodal_encoders_from_ckpt = False + if remove_multimodal_encoders_from_ckpt: + keys_to_remove = [] + for key in ckpt["model"].keys(): + if "ray_dirs_encoder" in key or "depth_encoder" in key or "depth_scale_encoder" in key or "cam_rot_encoder" in key or "cam_trans_encoder" in key or "cam_trans_scale_encoder" in key or "scene_flow_encoder" in key: + # if key.contains("ray_dirs_encoder") or key.contains("depth_encoder") or key.contains("depth_scale_encoder") or key.contains("cam_rot_encoder") or key.contains("cam_trans_encoder") or key.contains("cam_trans_scale_encoder") or key.contains("scene_flow_encoder"): + print(f"Removing key from pretrained checkpoint: {key}") + keys_to_remove.append(key) + for key in keys_to_remove: + del ckpt["model"][key] + + print(self.load_state_dict(ckpt["model"], strict=False)) + + init_scene_flow_encoder_with_depth_encoder = False + if init_scene_flow_encoder_with_depth_encoder: + print("Initializing scene flow encoder with depth encoder weights...") + depth_encoder_state = self.depth_encoder.state_dict() + scene_flow_encoder_state = self.scene_flow_encoder.state_dict() + + # Copy weights that match in both name and shape + for k in scene_flow_encoder_state.keys(): + if k in depth_encoder_state and depth_encoder_state[k].shape == scene_flow_encoder_state[k].shape: + scene_flow_encoder_state[k] = depth_encoder_state[k] + print(f"Copied weight: {k}") + else: + print(f"Skipping weight: {k}, shape mismatch or not found in depth encoder.") + + # Load the updated weights (non-strict so any differing layers can be skipped) + print(self.scene_flow_encoder.load_state_dict(scene_flow_encoder_state, strict=False)) + print("✅ Scene flow encoder successfully initialized from depth encoder weights.") + + else: + print( + f"Loading pretrained MapAnything weights from {self.pretrained_checkpoint_path} for specific submodules: {self.specific_pretrained_submodules} ..." + ) + assert self.pred_head_type is not None, ( + "Specific submodules to load cannot be None." + ) + ckpt = torch.load(self.pretrained_checkpoint_path, map_location="cpu", weights_only=False) + filtered_ckpt = {} + for ckpt_key, ckpt_value in ckpt["model"].items(): + for submodule in self.specific_pretrained_submodules: + if ckpt_key.startswith(submodule): + filtered_ckpt[ckpt_key] = ckpt_value + print(self.load_state_dict(filtered_ckpt, strict=False)) + + initialize_scene_flow_dense_head_with_dense_head = False + if initialize_scene_flow_dense_head_with_dense_head: + # Copy dense_head to scene_flow_dense_head weights + print("Copying dense_head weights into scene_flow_dense_head (except final layer)...") + dense_state = self.dense_head.state_dict() + scene_flow_state = self.scene_flow_dense_head.state_dict() + + # Copy weights that match in both name and shape + for k in scene_flow_state.keys(): + if k in dense_state and dense_state[k].shape == scene_flow_state[k].shape: + scene_flow_state[k] = dense_state[k] + else: + print(f"Skipping weight: {k}, shape mismatch or not found in dense_head.") + + # Load the updated weights (non-strict so final layer can differ) + print(self.scene_flow_dense_head.load_state_dict(scene_flow_state, strict=False)) + print("✅ Scene flow dense head successfully initialized from dense head weights.") + + + # Free unused memory + import gc + torch.cuda.empty_cache() + gc.collect() + + # def _encode_n_views(self, views): + # """ + # Encode all the input views (batch of images) in a single forward pass. + # Assumes all the input views have the same image shape, batch size, and data normalization type. + + # Args: + # views (List[dict]): List of dictionaries containing the input views' images and instance information. + + # Returns: + # List[torch.Tensor]: A list containing the encoded features for all N views. + # """ + # num_views = len(views) + # data_norm_type = views[0]["data_norm_type"][0] + # imgs_list = [view["img"] for view in views] + # all_imgs_across_views = torch.cat(imgs_list, dim=0) + # encoder_input = ViTEncoderInput( + # image=all_imgs_across_views, data_norm_type=data_norm_type + # ) + # encoder_output = self.encoder(encoder_input) + # all_encoder_features_across_views = encoder_output.features.chunk( + # num_views, dim=0 + # ) + + # return all_encoder_features_across_views + + + def _encode_n_views(self, views): + """ + Encode all the input views (batch of images) in a single forward pass. + Assumes all the input views have the same image shape, batch size, and data normalization type. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + + Returns: + List[torch.Tensor]: A list containing the encoded features for all N views. + """ + # num_views = len(views) + # data_norm_type = views[0]["data_norm_type"][0] + # imgs_list = [view["img"] for view in views] + # all_imgs_across_views = torch.cat(imgs_list, dim=0) + # encoder_input = ViTEncoderInput( + # image=all_imgs_across_views, data_norm_type=data_norm_type + # ) + # encoder_output = self.encoder(encoder_input) + # all_encoder_features_across_views = encoder_output.features.chunk( + # num_views, dim=0 + # ) + + # Option-1 + num_views = len(views) + batch_size_per_view = views[0]["img"].shape[0] + data_norm_type = views[0]["data_norm_type"][0] + device = views[0]["img"].device + + # Create image masking mask + per_sample_image_mask = ( + torch.rand(batch_size_per_view, device=device) + < self.geometric_input_config["images_prob"] + ) + + # Expand to all views + per_sample_image_mask = per_sample_image_mask.repeat(num_views) + + # Prepare images with masking + imgs_list = [] + for view_idx, view in enumerate(views): + img = view["img"].clone() + view_mask = per_sample_image_mask[ + view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view + ] + # Zero out images where the mask is False + img[~view_mask] = 0.0 + imgs_list.append(img) + + all_imgs_across_views = torch.cat(imgs_list, dim=0) + + encoder_input = ViTEncoderInput( + image=all_imgs_across_views, data_norm_type=data_norm_type + ) + encoder_output = self.encoder(encoder_input) + all_encoder_features_across_views = encoder_output.features.chunk( + num_views, dim=0 + ) + # Store the image mask for use in geometric encoding + self._current_image_mask = per_sample_image_mask + + return all_encoder_features_across_views + + + def _compute_pose_quats_and_trans_for_across_views_in_ref_view( + self, + views, + num_views, + device, + dtype, + batch_size_per_view, + per_sample_cam_input_mask, + ): + """ + Compute the pose quats and trans for all the views in the frame of the reference view 0. + Returns identity pose for views where the camera input mask is False or the pose is not provided. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + num_views (int): Number of views. + device (torch.device): Device to use for the computation. + dtype (torch.dtype): Data type to use for the computation. + per_sample_cam_input_mask (torch.Tensor): Tensor containing the per sample camera input mask. + + Returns: + torch.Tensor: A tensor containing the pose quats for all the views in the frame of the reference view 0. (batch_size_per_view * view, 4) + torch.Tensor: A tensor containing the pose trans for all the views in the frame of the reference view 0. (batch_size_per_view * view, 3) + torch.Tensor: A tensor containing the per sample camera input mask. + """ + # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 + pose_quats_non_ref_views = [] + pose_trans_non_ref_views = [] + pose_quats_ref_view_0 = [] + pose_trans_ref_view_0 = [] + for view_idx in range(num_views): + per_sample_cam_input_mask_for_curr_view = per_sample_cam_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view + ] + if ( + "camera_pose_quats" in views[view_idx] + and "camera_pose_trans" in views[view_idx] + and per_sample_cam_input_mask_for_curr_view.any() + ): + # Get the camera pose quats and trans for the current view + cam_pose_quats = views[view_idx]["camera_pose_quats"][ + per_sample_cam_input_mask_for_curr_view + ] + cam_pose_trans = views[view_idx]["camera_pose_trans"][ + per_sample_cam_input_mask_for_curr_view + ] + # Append to the list + pose_quats_non_ref_views.append(cam_pose_quats) + pose_trans_non_ref_views.append(cam_pose_trans) + # Get the camera pose quats and trans for the reference view 0 + cam_pose_quats = views[0]["camera_pose_quats"][ + per_sample_cam_input_mask_for_curr_view + ] + cam_pose_trans = views[0]["camera_pose_trans"][ + per_sample_cam_input_mask_for_curr_view + ] + # Append to the list + pose_quats_ref_view_0.append(cam_pose_quats) + pose_trans_ref_view_0.append(cam_pose_trans) + else: + per_sample_cam_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) + * batch_size_per_view + ] = False + + # Initialize the pose quats and trans for all views as identity + pose_quats_across_views = torch.tensor( + [0.0, 0.0, 0.0, 1.0], dtype=dtype, device=device + ).repeat(batch_size_per_view * num_views, 1) # (q_x, q_y, q_z, q_w) + pose_trans_across_views = torch.zeros( + (batch_size_per_view * num_views, 3), dtype=dtype, device=device + ) + + # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 + if len(pose_quats_non_ref_views) > 0: + # Stack the pose quats and trans for all the non-reference views and reference view 0 + pose_quats_non_ref_views = torch.cat(pose_quats_non_ref_views, dim=0) + pose_trans_non_ref_views = torch.cat(pose_trans_non_ref_views, dim=0) + pose_quats_ref_view_0 = torch.cat(pose_quats_ref_view_0, dim=0) + pose_trans_ref_view_0 = torch.cat(pose_trans_ref_view_0, dim=0) + + # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 + ( + pose_quats_non_ref_views_in_ref_view_0, + pose_trans_non_ref_views_in_ref_view_0, + ) = transform_pose_using_quats_and_trans_2_to_1( + pose_quats_ref_view_0, + pose_trans_ref_view_0, + pose_quats_non_ref_views, + pose_trans_non_ref_views, + ) + + # Update the pose quats and trans for all the non-reference views + pose_quats_across_views[per_sample_cam_input_mask] = ( + pose_quats_non_ref_views_in_ref_view_0.to(dtype=dtype) + ) + pose_trans_across_views[per_sample_cam_input_mask] = ( + pose_trans_non_ref_views_in_ref_view_0.to(dtype=dtype) + ) + + return ( + pose_quats_across_views, + pose_trans_across_views, + per_sample_cam_input_mask, + ) + + def _encode_and_fuse_ray_dirs( + self, + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + per_sample_ray_dirs_input_mask, + ): + """ + Encode the ray directions for all the views and fuse it with the other encoder features in a single forward pass. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + num_views (int): Number of views. + batch_size_per_view (int): Batch size per view. + all_encoder_features_across_views (torch.Tensor): Tensor containing the encoded features for all N views. + per_sample_ray_dirs_input_mask (torch.Tensor): Tensor containing the per sample ray direction input mask. + + Returns: + torch.Tensor: A tensor containing the encoded features for all the views. + """ + # Get the height and width of the images + _, _, height, width = views[0]["img"].shape + + # Get the ray directions for all the views where info is provided and the ray direction input mask is True + ray_dirs_list = [] + for view_idx in range(num_views): + per_sample_ray_dirs_input_mask_for_curr_view = ( + per_sample_ray_dirs_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) + * batch_size_per_view + ] + ) + ray_dirs_for_curr_view = torch.zeros( + (batch_size_per_view, height, width, 3), + dtype=all_encoder_features_across_views.dtype, + device=all_encoder_features_across_views.device, + ) + if ( + "ray_directions_cam" in views[view_idx] + and per_sample_ray_dirs_input_mask_for_curr_view.any() + ): + ray_dirs_for_curr_view[per_sample_ray_dirs_input_mask_for_curr_view] = ( + views[view_idx]["ray_directions_cam"][ + per_sample_ray_dirs_input_mask_for_curr_view + ] + ) + else: + per_sample_ray_dirs_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) + * batch_size_per_view + ] = False + ray_dirs_list.append(ray_dirs_for_curr_view) + + # Stack the ray directions for all the views and permute to (B * V, C, H, W) + ray_dirs = torch.cat(ray_dirs_list, dim=0) # (B * V, H, W, 3) + ray_dirs = ray_dirs.permute(0, 3, 1, 2).contiguous() # (B * V, 3, H, W) + + # Encode the ray directions + ray_dirs_features_across_views = self.ray_dirs_encoder( + ViTEncoderNonImageInput(data=ray_dirs) + ).features + + # Fuse the ray direction features with the other encoder features (zero out the features where the ray direction input mask is False) + ray_dirs_features_across_views = ( + ray_dirs_features_across_views + * per_sample_ray_dirs_input_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + ) + all_encoder_features_across_views = ( + all_encoder_features_across_views + ray_dirs_features_across_views + ) + + return all_encoder_features_across_views + + def _encode_and_fuse_depths( + self, + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + per_sample_depth_input_mask, + ): + """ + Encode the z depths for all the views and fuse it with the other encoder features in a single forward pass. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + num_views (int): Number of views. + batch_size_per_view (int): Batch size per view. + all_encoder_features_across_views (torch.Tensor): Tensor containing the encoded features for all N views. + per_sample_depth_input_mask (torch.Tensor): Tensor containing the per sample depth input mask. + + Returns: + torch.Tensor: A tensor containing the encoded features for all the views. + """ + # Get the device and height and width of the images + device = all_encoder_features_across_views.device + _, _, height, width = views[0]["img"].shape + + # Decide to use randomly sampled sparse depth or dense depth + if torch.rand(1) < self.geometric_input_config["sparse_depth_prob"]: + use_sparse_depth = True + else: + use_sparse_depth = False + + # Get the depths for all the views + depth_list = [] + depth_norm_factors_list = [] + metric_scale_depth_mask_list = [] + for view_idx in range(num_views): + # Get the input mask for current view + per_sample_depth_input_mask_for_curr_view = per_sample_depth_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view + ] + depth_for_curr_view = torch.zeros( + (batch_size_per_view, height, width, 1), + dtype=all_encoder_features_across_views.dtype, + device=device, + ) + depth_norm_factor_for_curr_view = torch.zeros( + (batch_size_per_view), + dtype=all_encoder_features_across_views.dtype, + device=device, + ) + metric_scale_mask_for_curr_view = torch.zeros( + (batch_size_per_view), + dtype=torch.bool, + device=device, + ) + if ( + "depth_along_ray" in views[view_idx] + ) and per_sample_depth_input_mask_for_curr_view.any(): + # Get depth for current view + depth_for_curr_view_input = views[view_idx]["depth_along_ray"][ + per_sample_depth_input_mask_for_curr_view + ] + # Get the metric scale mask + if "is_metric_scale" in views[view_idx]: + metric_scale_mask = views[view_idx]["is_metric_scale"][ + per_sample_depth_input_mask_for_curr_view + ] + else: + metric_scale_mask = torch.zeros( + depth_for_curr_view_input.shape[0], + dtype=torch.bool, + device=device, + ) + # Turn off indication of metric scale samples based on the depth_scale_norm_all_prob + depth_scale_norm_all_mask = ( + torch.rand(metric_scale_mask.shape[0]) + < self.geometric_input_config["depth_scale_norm_all_prob"] + ) + if depth_scale_norm_all_mask.any(): + metric_scale_mask[depth_scale_norm_all_mask] = False + # Assign the metric scale mask to the respective indices + metric_scale_mask_for_curr_view[ + per_sample_depth_input_mask_for_curr_view + ] = metric_scale_mask + # Sparsely sample the depth if required + if use_sparse_depth: + # Create a mask of ones + sparsification_mask = torch.ones_like( + depth_for_curr_view_input, device=device + ) + # Create a mask for valid pixels (depth > 0) + valid_pixel_mask = depth_for_curr_view_input > 0 + # Calculate the number of valid pixels + num_valid_pixels = valid_pixel_mask.sum().item() + # Calculate the number of valid pixels to set to zero + num_to_zero = int( + num_valid_pixels + * self.geometric_input_config["sparsification_removal_percent"] + ) + if num_to_zero > 0: + # Get the indices of valid pixels + valid_indices = valid_pixel_mask.nonzero(as_tuple=True) + # Randomly select indices to zero out + indices_to_zero = torch.randperm(num_valid_pixels)[:num_to_zero] + # Set selected valid indices to zero in the mask + sparsification_mask[ + valid_indices[0][indices_to_zero], + valid_indices[1][indices_to_zero], + valid_indices[2][indices_to_zero], + valid_indices[3][indices_to_zero], + ] = 0 + # Apply the mask on the depth + depth_for_curr_view_input = ( + depth_for_curr_view_input * sparsification_mask + ) + # Normalize the depth + scaled_depth_for_curr_view_input, depth_norm_factor = ( + normalize_depth_using_non_zero_pixels( + depth_for_curr_view_input, return_norm_factor=True + ) + ) + # Assign the depth and depth norm factor to the respective indices + depth_for_curr_view[per_sample_depth_input_mask_for_curr_view] = ( + scaled_depth_for_curr_view_input + ) + depth_norm_factor_for_curr_view[ + per_sample_depth_input_mask_for_curr_view + ] = depth_norm_factor + else: + per_sample_depth_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) + * batch_size_per_view + ] = False + # Append the depths, depth norm factor and metric scale mask for the current view + depth_list.append(depth_for_curr_view) + depth_norm_factors_list.append(depth_norm_factor_for_curr_view) + metric_scale_depth_mask_list.append(metric_scale_mask_for_curr_view) + + # Stack the depths for all the views and permute to (B * V, C, H, W) + depths = torch.cat(depth_list, dim=0) # (B * V, H, W, 1) + depths = apply_log_to_norm( + depths + ) # Scale logarithimically (norm is computed along last dim) + depths = depths.permute(0, 3, 1, 2).contiguous() # (B * V, 1, H, W) + # Encode the depths using the depth encoder + depth_features_across_views = self.depth_encoder( + ViTEncoderNonImageInput(data=depths) + ).features + # Zero out the depth features where the depth input mask is False + depth_features_across_views = ( + depth_features_across_views + * per_sample_depth_input_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + ) + + # Stack the depth norm factors for all the views + depth_norm_factors = torch.cat(depth_norm_factors_list, dim=0) # (B * V, ) + # Encode the depth norm factors using the log scale encoder for depth + log_depth_norm_factors = torch.log(depth_norm_factors + 1e-8) # (B * V, ) + depth_scale_features_across_views = self.depth_scale_encoder( + EncoderGlobalRepInput(data=log_depth_norm_factors.unsqueeze(-1)) + ).features + # Zero out the depth scale features where the depth input mask is False + depth_scale_features_across_views = ( + depth_scale_features_across_views + * per_sample_depth_input_mask.unsqueeze(-1) + ) + # Stack the metric scale mask for all the views + metric_scale_depth_mask = torch.cat( + metric_scale_depth_mask_list, dim=0 + ) # (B * V, ) + # Zero out the depth scale features where the metric scale mask is False + # Scale encoding is only provided for metric scale samples + depth_scale_features_across_views = ( + depth_scale_features_across_views * metric_scale_depth_mask.unsqueeze(-1) + ) + + # Fuse the depth features & depth scale features with the other encoder features + all_encoder_features_across_views = ( + all_encoder_features_across_views + + depth_features_across_views + + depth_scale_features_across_views.unsqueeze(-1).unsqueeze(-1) + ) + + return all_encoder_features_across_views + + def _encode_and_fuse_cam_quats_and_trans( + self, + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + pose_quats_across_views, + pose_trans_across_views, + per_sample_cam_input_mask, + ): + """ + Encode the camera quats and trans for all the views and fuse it with the other encoder features in a single forward pass. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + num_views (int): Number of views. + batch_size_per_view (int): Batch size per view. + all_encoder_features_across_views (torch.Tensor): Tensor containing the encoded features for all N views. + pose_quats_across_views (torch.Tensor): Tensor containing the pose quats for all the views in the frame of the reference view 0. (batch_size_per_view * view, 4) + pose_trans_across_views (torch.Tensor): Tensor containing the pose trans for all the views in the frame of the reference view 0. (batch_size_per_view * view, 3) + per_sample_cam_input_mask (torch.Tensor): Tensor containing the per sample camera input mask. + + Returns: + torch.Tensor: A tensor containing the encoded features for all the views. + """ + # Encode the pose quats + pose_quats_features_across_views = self.cam_rot_encoder( + EncoderGlobalRepInput(data=pose_quats_across_views) + ).features + # Zero out the pose quat features where the camera input mask is False + pose_quats_features_across_views = ( + pose_quats_features_across_views * per_sample_cam_input_mask.unsqueeze(-1) + ) + + # Get the metric scale mask for all samples + device = all_encoder_features_across_views.device + metric_scale_pose_trans_mask = torch.zeros( + (batch_size_per_view * num_views), dtype=torch.bool, device=device + ) + for view_idx in range(num_views): + if "is_metric_scale" in views[view_idx]: + # Get the metric scale mask for the input pose priors + metric_scale_mask = views[view_idx]["is_metric_scale"] + else: + metric_scale_mask = torch.zeros( + batch_size_per_view, dtype=torch.bool, device=device + ) + metric_scale_pose_trans_mask[ + view_idx * batch_size_per_view : (view_idx + 1) * batch_size_per_view + ] = metric_scale_mask + + # Turn off indication of metric scale samples based on the pose_scale_norm_all_prob + pose_norm_all_mask = ( + torch.rand(batch_size_per_view * num_views) + < self.geometric_input_config["pose_scale_norm_all_prob"] + ) + if pose_norm_all_mask.any(): + metric_scale_pose_trans_mask[pose_norm_all_mask] = False + + # Get the scale norm factor for all the samples and scale the pose translations + pose_trans_across_views = torch.split( + pose_trans_across_views, batch_size_per_view, dim=0 + ) # Split into num_views chunks + pose_trans_across_views = torch.stack( + pose_trans_across_views, dim=1 + ) # Stack the views along a new dimension (batch_size_per_view, num_views, 3) + scaled_pose_trans_across_views, pose_trans_norm_factors = ( + normalize_pose_translations( + pose_trans_across_views, return_norm_factor=True + ) + ) + + # Resize the pose translation back to (batch_size_per_view * num_views, 3) and extend the norm factor to (batch_size_per_view * num_views, 1) + scaled_pose_trans_across_views = scaled_pose_trans_across_views.unbind( + dim=1 + ) # Convert back to list of views, where each view has batch_size_per_view tensor + scaled_pose_trans_across_views = torch.cat( + scaled_pose_trans_across_views, dim=0 + ) # Concatenate back to (batch_size_per_view * num_views, 3) + pose_trans_norm_factors_across_views = pose_trans_norm_factors.unsqueeze( + -1 + ).repeat(num_views, 1) # (B, ) -> (B * V, 1) + + # Encode the pose trans + pose_trans_features_across_views = self.cam_trans_encoder( + EncoderGlobalRepInput(data=scaled_pose_trans_across_views) + ).features + # Zero out the pose trans features where the camera input mask is False + pose_trans_features_across_views = ( + pose_trans_features_across_views * per_sample_cam_input_mask.unsqueeze(-1) + ) + + # Encode the pose translation norm factors using the log scale encoder for pose trans + log_pose_trans_norm_factors_across_views = torch.log( + pose_trans_norm_factors_across_views + 1e-8 + ) + pose_trans_scale_features_across_views = self.cam_trans_scale_encoder( + EncoderGlobalRepInput(data=log_pose_trans_norm_factors_across_views) + ).features + # Zero out the pose trans scale features where the camera input mask is False + pose_trans_scale_features_across_views = ( + pose_trans_scale_features_across_views + * per_sample_cam_input_mask.unsqueeze(-1) + ) + # Zero out the pose trans scale features where the metric scale mask is False + # Scale encoding is only provided for metric scale samples + pose_trans_scale_features_across_views = ( + pose_trans_scale_features_across_views + * metric_scale_pose_trans_mask.unsqueeze(-1) + ) + + # Fuse the pose quat features, pose trans features, pose trans scale features and pose trans type PE features with the other encoder features + all_encoder_features_across_views = ( + all_encoder_features_across_views + + pose_quats_features_across_views.unsqueeze(-1).unsqueeze(-1) + + pose_trans_features_across_views.unsqueeze(-1).unsqueeze(-1) + + pose_trans_scale_features_across_views.unsqueeze(-1).unsqueeze(-1) + ) + + return all_encoder_features_across_views + + def _encode_and_fuse_scene_flow( + self, + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + per_sample_scene_flow_input_mask, + ): + """ + Get doppler-style radial scene flow for all the views and fuse it with the other encoder features in a single forward pass. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + num_views (int): Number of views. + batch_size_per_view (int): Batch size per view. + all_encoder_features_across_views (torch.Tensor): Tensor containing the encoded features for all N views. + per_sample_scene_flow_input_mask (torch.Tensor): Tensor containing the per sample scene flow input mask. + + Returns: + torch.Tensor: A tensor containing the encoded features for all the views. + """ + + # Get the device and height and width of the images + device = all_encoder_features_across_views.device + _, _, height, width = views[0]["img"].shape + + # Get pts3d_cam for 0th view and compute normalization factor - for normalizing scene flows of views + if per_sample_scene_flow_input_mask.any(): + pts3d_cam_0 = views[0]["pts3d_cam"] + cam_0_valid_mask = views[0]["valid_mask"] + _, norm_factor = normalize_multiple_pointclouds([pts3d_cam_0], valid_masks=[cam_0_valid_mask], ret_factor= True) + + # Get ego-scene flow for all the views where info is provided and the scene flow input mask is True + scene_flow_list = [] + for view_idx in range(num_views): + per_sample_scene_flow_input_mask_for_curr_view = ( + per_sample_scene_flow_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) + * batch_size_per_view + ] + ) + scene_flow_for_curr_view = torch.zeros( + (batch_size_per_view, height, width, 3), + dtype=all_encoder_features_across_views.dtype, + device=all_encoder_features_across_views.device, + ) + if ( + "radial_scene_flow" in views[view_idx] + and per_sample_scene_flow_input_mask_for_curr_view.any() + ): + scene_flow_for_curr_view[per_sample_scene_flow_input_mask_for_curr_view] = ( + views[view_idx]["radial_scene_flow"][ + per_sample_scene_flow_input_mask_for_curr_view + ] + ) + + # Normalize the scene flow using the normalization factor of 0th view + scene_flow_for_curr_view = scene_flow_for_curr_view / (norm_factor + 1e-8) + else: + per_sample_scene_flow_input_mask[ + view_idx * batch_size_per_view : (view_idx + 1) + * batch_size_per_view + ] = False + scene_flow_list.append(scene_flow_for_curr_view) + + # Stack the scene flows for all the views and permute to (B * V, C, H, W) + scene_flows = torch.cat(scene_flow_list, dim=0) # (B * V, H, W, 3) + scene_flows = scene_flows.permute(0, 3, 1, 2).contiguous() # (B * V, 3, H, W) + + # Encode the scene flows + scene_flow_features_across_views = self.scene_flow_encoder( + ViTEncoderNonImageInput(data=scene_flows) + ).features + # Fuse the scene flow features with the other encoder features (zero out the features where the scene flow input mask is False) + scene_flow_features_across_views = ( + scene_flow_features_across_views + * per_sample_scene_flow_input_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + ) + all_encoder_features_across_views = ( + all_encoder_features_across_views + scene_flow_features_across_views + ) + + return all_encoder_features_across_views + + def _encode_and_fuse_optional_geometric_inputs( + self, views, all_encoder_features_across_views_list + ): + """ + Encode all the input optional geometric modalities and fuses it with the image encoder features in a single forward pass. + Assumes all the input views have the same shape and batch size. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + all_encoder_features_across_views (List[torch.Tensor]): List of tensors containing the encoded image features for all N views. + + Returns: + List[torch.Tensor]: A list containing the encoded features for all N views. + """ + num_views = len(views) + batch_size_per_view, _, _, _ = views[0]["img"].shape + device = all_encoder_features_across_views_list[0].device + dtype = all_encoder_features_across_views_list[0].dtype + all_encoder_features_across_views = torch.cat( + all_encoder_features_across_views_list, dim=0 + ) + # Get the overall input mask for all the views + overall_geometric_input_mask = ( + torch.rand(batch_size_per_view, device=device) + < self.geometric_input_config["overall_prob"] + ) + overall_geometric_input_mask = overall_geometric_input_mask.repeat(num_views) + + # Get the per sample input mask after dropout + # Per sample input mask is in view-major order so that index v*B + b in each mask corresponds to sample b of view v: (B * V) + per_sample_geometric_input_mask = torch.rand( + batch_size_per_view * num_views, device=device + ) < (1 - self.geometric_input_config["dropout_prob"]) + per_sample_geometric_input_mask = ( + per_sample_geometric_input_mask & overall_geometric_input_mask + ) + + # Get the ray direction input mask + per_sample_ray_dirs_input_mask = ( + torch.rand(batch_size_per_view, device=device) + < self.geometric_input_config["ray_dirs_prob"] + ) + per_sample_ray_dirs_input_mask = per_sample_ray_dirs_input_mask.repeat( + num_views + ) + per_sample_ray_dirs_input_mask = ( + per_sample_ray_dirs_input_mask & per_sample_geometric_input_mask + ) + + # Get the depth input mask + per_sample_depth_input_mask = ( + torch.rand(batch_size_per_view, device=device) + < self.geometric_input_config["depth_prob"] + ) + per_sample_depth_input_mask = per_sample_depth_input_mask.repeat(num_views) + per_sample_depth_input_mask = ( + per_sample_depth_input_mask & per_sample_geometric_input_mask + ) + + # Get the camera input mask + per_sample_cam_input_mask = ( + torch.rand(batch_size_per_view, device=device) + < self.geometric_input_config["cam_prob"] + ) + per_sample_cam_input_mask = per_sample_cam_input_mask.repeat(num_views) + per_sample_cam_input_mask = ( + per_sample_cam_input_mask & per_sample_geometric_input_mask + ) + + # Get the scene flow input mask + per_sample_scene_flow_input_mask = ( + torch.rand(batch_size_per_view, device=device) + < self.geometric_input_config["doppler_prob"] + ) + + per_sample_scene_flow_input_mask = per_sample_scene_flow_input_mask.repeat(num_views) + per_sample_scene_flow_input_mask = ( + per_sample_scene_flow_input_mask & per_sample_geometric_input_mask + ) + + # Compute the pose quats and trans for all the non-reference views in the frame of the reference view 0 + # Returned pose quats and trans represent identity pose for views/samples where the camera input mask is False + pose_quats_across_views, pose_trans_across_views, per_sample_cam_input_mask = ( + self._compute_pose_quats_and_trans_for_across_views_in_ref_view( + views, + num_views, + device, + dtype, + batch_size_per_view, + per_sample_cam_input_mask, + ) + ) + + # Encode the ray directions and fuse with the image encoder features + all_encoder_features_across_views = self._encode_and_fuse_ray_dirs( + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + per_sample_ray_dirs_input_mask, + ) + + # Encode the depths and fuse with the image encoder features + all_encoder_features_across_views = self._encode_and_fuse_depths( + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + per_sample_depth_input_mask, + ) + + # Encode the cam quat and trans and fuse with the image encoder features + all_encoder_features_across_views = self._encode_and_fuse_cam_quats_and_trans( + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + pose_quats_across_views, + pose_trans_across_views, + per_sample_cam_input_mask, + ) + + # Encode the scene flows and fuse with the image encoder features + all_encoder_features_across_views = self._encode_and_fuse_scene_flow( + views, + num_views, + batch_size_per_view, + all_encoder_features_across_views, + per_sample_scene_flow_input_mask, + ) + + # Normalize the fused features (permute -> normalize -> permute) + all_encoder_features_across_views = all_encoder_features_across_views.permute( + 0, 2, 3, 1 + ).contiguous() + all_encoder_features_across_views = self.fusion_norm_layer( + all_encoder_features_across_views + ) + all_encoder_features_across_views = all_encoder_features_across_views.permute( + 0, 3, 1, 2 + ).contiguous() + + # Split the batched views into individual views + fused_all_encoder_features_across_views = ( + all_encoder_features_across_views.chunk(num_views, dim=0) + ) + + return fused_all_encoder_features_across_views + + def forward(self, views): + """ + Forward pass performing the following operations: + 1. Encodes the N input views (images). + 2. Encodes the optional geometric inputs (ray directions, depths, camera rotations, camera translations). + 3. Fuses the encoded features from the N input views and the optional geometric inputs using addition and normalization. + 4. Information sharing across the encoded features and a scale token using a multi-view attention transformer. + 5. Passes the final features from transformer through the prediction heads. + 6. Returns the processed final outputs for N views. + + Assumption: + - All the input views and dense geometric inputs have the same image shape. + + Args: + views (List[dict]): List of dictionaries containing the input views' images and instance information. + Each dictionary should contain the following keys: + "img" (tensor): Image tensor of shape (B, C, H, W). Input images must be normalized based on the data norm type of image encoder. + "data_norm_type" (list): [model.encoder.data_norm_type] + Optionally, each dictionary can also contain the following keys for the respective optional geometric inputs: + "ray_directions_cam" (tensor): Ray directions in the local camera frame. Tensor of shape (B, H, W, 3). + "depth_along_ray" (tensor): Depth along the ray. Tensor of shape (B, H, W, 1). + "camera_pose_quats" (tensor): Camera pose quaternions. Tensor of shape (B, 4). Camera pose is opencv (RDF) cam2world transformation. + "camera_pose_trans" (tensor): Camera pose translations. Tensor of shape (B, 3). Camera pose is opencv (RDF) cam2world transformation. + "is_metric_scale" (tensor): Boolean tensor indicating whether the geometric inputs are in metric scale or not. Tensor of shape (B, 1). + + Returns: + List[dict]: A list containing the final outputs for all N views. + """ + # Get input shape of the images, number of views, and batch size per view + batch_size_per_view, _, height, width = views[0]["img"].shape + img_shape = (int(height), int(width)) + num_views = len(views) + + # Run the image encoder on all the input views + all_encoder_features_across_views = self._encode_n_views(views) + + # Encode the optional geometric inputs and fuse with the encoded features from the N input views + # Use high precision to prevent NaN values after layer norm in dense representation encoder (due to high variance in last dim of features) + with torch.autocast("cuda", enabled=False): + all_encoder_features_across_views = ( + self._encode_and_fuse_optional_geometric_inputs( + views, all_encoder_features_across_views + ) + ) + + # Clean up the stored mask + if hasattr(self, '_current_image_mask'): + delattr(self, '_current_image_mask') + + # Expand the scale token to match the batch size + input_scale_token = ( + self.scale_token.unsqueeze(0) + .unsqueeze(-1) + .repeat(batch_size_per_view, 1, 1) + ) # (B, C, 1) + + # Combine all images into view-centric representation + # Output is a list containing the encoded features for all N views after information sharing. + info_sharing_input = MultiViewTransformerInput( + features=all_encoder_features_across_views, + additional_input_tokens=input_scale_token, + ) + if self.info_sharing_return_type == "no_intermediate_features": + final_info_sharing_multi_view_feat = self.info_sharing(info_sharing_input) + elif self.info_sharing_return_type == "intermediate_features": + ( + final_info_sharing_multi_view_feat, + intermediate_info_sharing_multi_view_feat, + ) = self.info_sharing(info_sharing_input) + + if self.pred_head_type == "linear": + # Stack the features for all views + dense_head_inputs = torch.cat( + final_info_sharing_multi_view_feat.features, dim=0 + ) + elif self.pred_head_type in ["dpt", "dpt+pose"]: + # Get the list of features for all views + dense_head_inputs_list = [] + if self.use_encoder_features_for_dpt: + # Stack all the image encoder features for all views + stacked_encoder_features = torch.cat( + all_encoder_features_across_views, dim=0 + ) + dense_head_inputs_list.append(stacked_encoder_features) + # Stack the first intermediate features for all views + stacked_intermediate_features_1 = torch.cat( + intermediate_info_sharing_multi_view_feat[0].features, dim=0 + ) + dense_head_inputs_list.append(stacked_intermediate_features_1) + # Stack the second intermediate features for all views + stacked_intermediate_features_2 = torch.cat( + intermediate_info_sharing_multi_view_feat[1].features, dim=0 + ) + dense_head_inputs_list.append(stacked_intermediate_features_2) + # Stack the last layer features for all views + stacked_final_features = torch.cat( + final_info_sharing_multi_view_feat.features, dim=0 + ) + dense_head_inputs_list.append(stacked_final_features) + else: + # Stack the first intermediate features for all views + stacked_intermediate_features_1 = torch.cat( + intermediate_info_sharing_multi_view_feat[0].features, dim=0 + ) + dense_head_inputs_list.append(stacked_intermediate_features_1) + # Stack the second intermediate features for all views + stacked_intermediate_features_2 = torch.cat( + intermediate_info_sharing_multi_view_feat[1].features, dim=0 + ) + dense_head_inputs_list.append(stacked_intermediate_features_2) + # Stack the third intermediate features for all views + stacked_intermediate_features_3 = torch.cat( + intermediate_info_sharing_multi_view_feat[2].features, dim=0 + ) + dense_head_inputs_list.append(stacked_intermediate_features_3) + # Stack the last layer + stacked_final_features = torch.cat( + final_info_sharing_multi_view_feat.features, dim=0 + ) + dense_head_inputs_list.append(stacked_final_features) + else: + raise ValueError( + f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" + ) + + # Downstream task prediction + with torch.autocast("cuda", enabled=False): + # Run Prediction Heads & Post-Process Outputs + if self.pred_head_type == "linear": + dense_head_outputs = self.dense_head( + PredictionHeadInput(last_feature=dense_head_inputs) + ) + dense_final_outputs = self.dense_adaptor( + AdaptorInput( + adaptor_feature=dense_head_outputs.decoded_channels, + output_shape_hw=img_shape, + ) + ) + scene_flow_dense_head_outputs = self.scene_flow_dense_head( + PredictionHeadInput(last_feature=dense_head_inputs) + ) + scene_flow_dense_final_outputs = self.scene_flow_dense_adaptor( + AdaptorInput( + adaptor_feature=scene_flow_dense_head_outputs.decoded_channels, + output_shape_hw=img_shape, + ) + ) + elif self.pred_head_type == "dpt": + dense_head_outputs = self.dense_head( + PredictionHeadLayeredInput( + list_features=dense_head_inputs_list, + target_output_shape=img_shape, + ) + ) + dense_final_outputs = self.dense_adaptor( + AdaptorInput( + adaptor_feature=dense_head_outputs.decoded_channels, + output_shape_hw=img_shape, + ) + ) + scene_flow_dense_head_outputs = self.scene_flow_dense_head( + PredictionHeadLayeredInput( + list_features=dense_head_inputs_list, + target_output_shape=img_shape, + ) + ) + scene_flow_dense_final_outputs = self.scene_flow_dense_adaptor( + AdaptorInput( + adaptor_feature=scene_flow_dense_head_outputs.decoded_channels, + output_shape_hw=img_shape, + ) + ) + elif self.pred_head_type == "dpt+pose": + dense_head_outputs = self.dense_head( + PredictionHeadLayeredInput( + list_features=dense_head_inputs_list, + target_output_shape=img_shape, + ) + ) + dense_final_outputs = self.dense_adaptor( + AdaptorInput( + adaptor_feature=dense_head_outputs.decoded_channels, + output_shape_hw=img_shape, + ) + ) + scene_flow_dense_head_outputs = self.scene_flow_dense_head( + PredictionHeadLayeredInput( + list_features=dense_head_inputs_list, + target_output_shape=img_shape, + ) + ) + scene_flow_dense_final_outputs = self.scene_flow_dense_adaptor( + AdaptorInput( + adaptor_feature=scene_flow_dense_head_outputs.decoded_channels, + output_shape_hw=img_shape, + ) + ) + pose_head_outputs = self.pose_head( + PredictionHeadInput(last_feature=dense_head_inputs_list[-1]) + ) + pose_final_outputs = self.pose_adaptor( + AdaptorInput( + adaptor_feature=pose_head_outputs.decoded_channels, + output_shape_hw=img_shape, + ) + ) + else: + raise ValueError( + f"Invalid pred_head_type: {self.pred_head_type}. Valid options: ['linear', 'dpt', 'dpt+pose']" + ) + scale_head_output = self.scale_head( + PredictionHeadTokenInput( + last_feature=final_info_sharing_multi_view_feat.additional_token_features + ) + ) + scale_final_output = self.scale_adaptor( + AdaptorInput( + adaptor_feature=scale_head_output.decoded_channels, + output_shape_hw=img_shape, + ) + ) + scale_final_output = scale_final_output.value.squeeze( + -1 + ) # (B, 1, 1) -> (B, 1) + + # Prepare the final scene representation for all views + if self.scene_rep_type in [ + "pointmap", + "pointmap+confidence", + "pointmap+mask", + "pointmap+confidence+mask", + ]: + output_pts3d = dense_final_outputs.value + # Reshape final scene representation to (B * V, H, W, C) + output_pts3d = output_pts3d.permute(0, 2, 3, 1).contiguous() + # Split the predicted pointmaps back to their respective views + output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) + # Pack the output as a list of dictionaries + res = [] + for i in range(num_views): + res.append( + { + "pts3d": output_pts3d_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "metric_scaling_factor": scale_final_output, + } + ) + elif self.scene_rep_type in [ + "raymap+depth", + "raymap+depth+confidence", + "raymap+depth+mask", + "raymap+depth+confidence+mask", + ]: + # Reshape final scene representation to (B * V, H, W, C) + output_scene_rep = dense_final_outputs.value.permute( + 0, 2, 3, 1 + ).contiguous() + # Get the predicted ray origins, directions, and depths along rays + output_ray_origins, output_ray_directions, output_depth_along_ray = ( + output_scene_rep.split([3, 3, 1], dim=-1) + ) + # Get the predicted pointmaps + output_pts3d = ( + output_ray_origins + output_ray_directions * output_depth_along_ray + ) + # Split the predicted quantities back to their respective views + output_ray_origins_per_view = output_ray_origins.chunk(num_views, dim=0) + output_ray_directions_per_view = output_ray_directions.chunk( + num_views, dim=0 + ) + output_depth_along_ray_per_view = output_depth_along_ray.chunk( + num_views, dim=0 + ) + output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) + # Pack the output as a list of dictionaries + res = [] + for i in range(num_views): + res.append( + { + "pts3d": output_pts3d_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "ray_origins": output_ray_origins_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "ray_directions": output_ray_directions_per_view[i], + "depth_along_ray": output_depth_along_ray_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "metric_scaling_factor": scale_final_output, + } + ) + elif self.scene_rep_type in [ + "raydirs+depth+pose", + "raydirs+depth+pose+confidence", + "raydirs+depth+pose+mask", + "raydirs+depth+pose+confidence+mask", + ]: + # Reshape output dense rep to (B * V, H, W, C) + output_dense_rep = dense_final_outputs.value.permute( + 0, 2, 3, 1 + ).contiguous() + # Get the predicted ray directions and depths along rays + output_ray_directions, output_depth_along_ray = output_dense_rep.split( + [3, 1], dim=-1 + ) + # Get the predicted camera translations and quaternions + output_cam_translations, output_cam_quats = ( + pose_final_outputs.value.split([3, 4], dim=-1) + ) + # Get the predicted pointmaps in world frame and camera frame + output_pts3d = ( + convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( + output_ray_directions, + output_depth_along_ray, + output_cam_translations, + output_cam_quats, + ) + ) + output_pts3d_cam = output_ray_directions * output_depth_along_ray + # Split the predicted quantities back to their respective views + output_ray_directions_per_view = output_ray_directions.chunk( + num_views, dim=0 + ) + output_depth_along_ray_per_view = output_depth_along_ray.chunk( + num_views, dim=0 + ) + output_cam_translations_per_view = output_cam_translations.chunk( + num_views, dim=0 + ) + output_cam_quats_per_view = output_cam_quats.chunk(num_views, dim=0) + output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) + output_pts3d_cam_per_view = output_pts3d_cam.chunk(num_views, dim=0) + # Pack the output as a list of dictionaries + res = [] + for i in range(num_views): + res.append( + { + "pts3d": output_pts3d_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "pts3d_cam": output_pts3d_cam_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "ray_directions": output_ray_directions_per_view[i], + "depth_along_ray": output_depth_along_ray_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "cam_trans": output_cam_translations_per_view[i] + * scale_final_output, + "cam_quats": output_cam_quats_per_view[i], + "metric_scaling_factor": scale_final_output, + } + ) + elif self.scene_rep_type in [ + "campointmap+pose", + "campointmap+pose+confidence", + "campointmap+pose+mask", + "campointmap+pose+confidence+mask", + ]: + # Get the predicted camera frame pointmaps + output_pts3d_cam = dense_final_outputs.value + # Reshape final scene representation to (B * V, H, W, C) + output_pts3d_cam = output_pts3d_cam.permute(0, 2, 3, 1).contiguous() + # Get the predicted camera translations and quaternions + output_cam_translations, output_cam_quats = ( + pose_final_outputs.value.split([3, 4], dim=-1) + ) + # Get the ray directions and depths along rays + output_depth_along_ray = torch.norm( + output_pts3d_cam, dim=-1, keepdim=True + ) + output_ray_directions = output_pts3d_cam / output_depth_along_ray + # Get the predicted pointmaps in world frame + output_pts3d = ( + convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( + output_ray_directions, + output_depth_along_ray, + output_cam_translations, + output_cam_quats, + ) + ) + # Split the predicted quantities back to their respective views + output_ray_directions_per_view = output_ray_directions.chunk( + num_views, dim=0 + ) + output_depth_along_ray_per_view = output_depth_along_ray.chunk( + num_views, dim=0 + ) + output_cam_translations_per_view = output_cam_translations.chunk( + num_views, dim=0 + ) + output_cam_quats_per_view = output_cam_quats.chunk(num_views, dim=0) + output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) + output_pts3d_cam_per_view = output_pts3d_cam.chunk(num_views, dim=0) + # Pack the output as a list of dictionaries + res = [] + for i in range(num_views): + res.append( + { + "pts3d": output_pts3d_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "pts3d_cam": output_pts3d_cam_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "ray_directions": output_ray_directions_per_view[i], + "depth_along_ray": output_depth_along_ray_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "cam_trans": output_cam_translations_per_view[i] + * scale_final_output, + "cam_quats": output_cam_quats_per_view[i], + "metric_scaling_factor": scale_final_output, + } + ) + elif self.scene_rep_type in [ + "raydirs+depth+pose+scene_flow", + "raydirs+depth+pose+scene_flow+confidence", + "raydirs+depth+pose+scene_flow+mask", + "raydirs+depth+pose+scene_flow+confidence+mask", + "raydirs+depth+pose+scene_flow+confidence+motion_mask", + ]: + # Reshape output dense rep to (B * V, H, W, C) + dense_rep = dense_final_outputs.value.permute( + 0, 2, 3, 1 + ).contiguous() + # Get the predicted ray directions and depths along rays + output_ray_directions, output_depth_along_ray = dense_rep.split( + [3, 1], dim=-1 + ) + # Get the predicted scene flow + scene_flow_dense_rep = scene_flow_dense_final_outputs.value.permute(0, 2, 3, 1).contiguous() + output_scene_flow = scene_flow_dense_rep + # Get the predicted camera translations and quaternions + output_cam_translations, output_cam_quats = ( + pose_final_outputs.value.split([3, 4], dim=-1) + ) + # Get the predicted pointmaps in world frame and camera frame + output_pts3d = ( + convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( + output_ray_directions, + output_depth_along_ray, + output_cam_translations, + output_cam_quats, + ) + ) + output_pts3d_cam = output_ray_directions * output_depth_along_ray + # Split the predicted quantities back to their respective views + output_ray_directions_per_view = output_ray_directions.chunk( + num_views, dim=0 + ) + output_depth_along_ray_per_view = output_depth_along_ray.chunk( + num_views, dim=0 + ) + output_cam_translations_per_view = output_cam_translations.chunk( + num_views, dim=0 + ) + output_cam_quats_per_view = output_cam_quats.chunk(num_views, dim=0) + output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) + output_pts3d_cam_per_view = output_pts3d_cam.chunk(num_views, dim=0) + output_scene_flow_per_view = output_scene_flow.chunk(num_views, dim=0) + # Pack the output as a list of dictionaries + res = [] + for i in range(num_views): + res.append( + { + "pts3d": output_pts3d_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "pts3d_cam": output_pts3d_cam_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "ray_directions": output_ray_directions_per_view[i], + "depth_along_ray": output_depth_along_ray_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "cam_trans": output_cam_translations_per_view[i] + * scale_final_output, + "cam_quats": output_cam_quats_per_view[i], + "scene_flow": output_scene_flow_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "metric_scaling_factor": scale_final_output, + } + ) + + elif self.scene_rep_type in [ + "pointmap+raydirs+depth+pose+scene_flow", + "pointmap+raydirs+depth+pose+scene_flow+confidence", + "pointmap+raydirs+depth+pose+scene_flow+mask", + "pointmap+raydirs+depth+pose+scene_flow+confidence+mask", + "pointmap+raydirs+depth+pose+scene_flow+confidence+motion_mask", + ]: + # Reshape output dense rep to (B * V, H, W, C) + dense_rep = dense_final_outputs.value.permute( + 0, 2, 3, 1 + ).contiguous() + # Get the predicted ray directions and depths along rays + output_pts3d, output_ray_directions, output_depth_along_ray = dense_rep.split( + [3, 3, 1], dim=-1 + ) + # Get the predicted scene flow + scene_flow_dense_rep = scene_flow_dense_final_outputs.value.permute(0, 2, 3, 1).contiguous() + output_scene_flow = scene_flow_dense_rep + # Get the predicted camera translations and quaternions + output_cam_translations, output_cam_quats = ( + pose_final_outputs.value.split([3, 4], dim=-1) + ) + # Replace the predicted world-frame pointmaps if required + if self.pred_head_config["adaptor_config"][ + "use_factored_predictions_for_global_pointmaps" + ]: + output_pts3d = ( + convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( + output_ray_directions, + output_depth_along_ray, + output_cam_translations, + output_cam_quats, + ) + ) + output_pts3d_cam = output_ray_directions * output_depth_along_ray + # Split the predicted quantities back to their respective views + output_ray_directions_per_view = output_ray_directions.chunk( + num_views, dim=0 + ) + output_depth_along_ray_per_view = output_depth_along_ray.chunk( + num_views, dim=0 + ) + output_cam_translations_per_view = output_cam_translations.chunk( + num_views, dim=0 + ) + output_cam_quats_per_view = output_cam_quats.chunk(num_views, dim=0) + output_pts3d_per_view = output_pts3d.chunk(num_views, dim=0) + output_pts3d_cam_per_view = output_pts3d_cam.chunk(num_views, dim=0) + output_scene_flow_per_view = output_scene_flow.chunk(num_views, dim=0) + # Pack the output as a list of dictionaries + res = [] + for i in range(num_views): + res.append( + { + "pts3d": output_pts3d_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "pts3d_cam": output_pts3d_cam_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "ray_directions": output_ray_directions_per_view[i], + "depth_along_ray": output_depth_along_ray_per_view[i] + * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "cam_trans": output_cam_translations_per_view[i] + * scale_final_output, + "cam_quats": output_cam_quats_per_view[i], + "scene_flow": output_scene_flow_per_view[i] * scale_final_output.unsqueeze(-1).unsqueeze(-1), + "metric_scaling_factor": scale_final_output, + } + ) + + + else: + raise ValueError( + f"Invalid scene_rep_type: {self.scene_rep_type}. \ + Valid options: ['pointmap', 'raymap+depth', 'raydirs+depth+pose', 'campointmap+pose', 'raydirs+depth+scene_flow+pose', \ + 'pointmap+confidence', 'raymap+depth+confidence', 'raydirs+depth+pose+confidence', 'campointmap+pose+confidence', 'raydirs+depth+scene_flow+pose+confidence', \ + 'pointmap+mask', 'raymap+depth+mask', 'raydirs+depth+pose+mask', 'campointmap+pose+mask', 'raydirs+depth+scene_flow+pose+mask', \ + 'pointmap+confidence+mask', 'raymap+depth+confidence+mask', 'raydirs+depth+pose+confidence+mask', 'campointmap+pose+confidence+mask', 'raydirs+depth+scene_flow+pose+confidence+mask']" + ) + + # Get the output confidences for all views (if available) and add them to the result + if "confidence" in self.scene_rep_type: + output_confidences = dense_final_outputs.confidence + # Reshape confidences to (B * V, H, W) + output_confidences = ( + output_confidences.permute(0, 2, 3, 1).squeeze(-1).contiguous() + ) + # Split the predicted confidences back to their respective views + output_confidences_per_view = output_confidences.chunk(num_views, dim=0) + # Add the confidences to the result + for i in range(num_views): + res[i]["conf"] = output_confidences_per_view[i] + + # Get the output masks (and logits) for all views (if available) and add them to the result + if "mask" in self.scene_rep_type: + # Get the output masks + output_masks = dense_final_outputs.mask + # Reshape masks to (B * V, H, W) + output_masks = output_masks.permute(0, 2, 3, 1).squeeze(-1).contiguous() + # Threshold the masks at 0.5 to get binary masks (0: ambiguous, 1: non-ambiguous) + output_masks = output_masks > 0.5 + # Split the predicted masks back to their respective views + output_masks_per_view = output_masks.chunk(num_views, dim=0) + # Get the output mask logits (for loss) + output_mask_logits = dense_final_outputs.logits + # Reshape mask logits to (B * V, H, W) + output_mask_logits = ( + output_mask_logits.permute(0, 2, 3, 1).squeeze(-1).contiguous() + ) + # Split the predicted mask logits back to their respective views + output_mask_logits_per_view = output_mask_logits.chunk(num_views, dim=0) + # Add the masks and logits to the result + for i in range(num_views): + res[i]["non_ambiguous_mask"] = output_masks_per_view[i] + res[i]["non_ambiguous_mask_logits"] = output_mask_logits_per_view[i] + + # Get the output motion masks (and logits) for all views (if available) and add them to the result + if "motion_mask" in self.scene_rep_type: + # Get the output motion masks + output_motion_masks = dense_final_outputs.mask + # Reshape motion masks to (B * V, H, W) + output_motion_masks = output_motion_masks.permute(0, 2, 3, 1).squeeze(-1).contiguous() + # Threshold the motion masks at 0.5 to get binary masks (0: static, 1: moving) + output_motion_masks = output_motion_masks > 0.5 + # Split the predicted motion masks back to their respective views + output_motion_masks_per_view = output_motion_masks.chunk(num_views, dim=0) + # Get the output motion mask logits (for loss) + output_motion_mask_logits = dense_final_outputs.logits + # Reshape motion mask logits to (B * V, H, W) + output_motion_mask_logits = output_motion_mask_logits.permute(0, 2, 3, 1).squeeze(-1).contiguous() + # Split the predicted motion mask logits back to their respective views + output_motion_mask_logits_per_view = output_motion_mask_logits.chunk(num_views, dim=0) + # Add the motion masks and logits to the result + for i in range(num_views): + res[i]["motion_mask"] = output_motion_masks_per_view[i] + res[i]["motion_mask_logits"] = output_motion_mask_logits_per_view[i] + + return res diff --git a/mapanything/utils/__init__.py b/mapanything/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mapanything/utils/__pycache__/__init__.cpython-312.pyc b/mapanything/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a7d43030efbd5578f109afee5764e22c99b82c3 Binary files /dev/null and b/mapanything/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/cropping.cpython-312.pyc b/mapanything/utils/__pycache__/cropping.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..28494165204bd54b31ce004f4080872fa94987fc Binary files /dev/null and b/mapanything/utils/__pycache__/cropping.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/geometry.cpython-312.pyc b/mapanything/utils/__pycache__/geometry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2832965b8fc0ddcebdbfe9929b24857f339fe76a Binary files /dev/null and b/mapanything/utils/__pycache__/geometry.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/image.cpython-312.pyc b/mapanything/utils/__pycache__/image.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f586ab838e385d4d21b52806e971aad6c681eef6 Binary files /dev/null and b/mapanything/utils/__pycache__/image.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/inference.cpython-312.pyc b/mapanything/utils/__pycache__/inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..967b3c953e3eff3bad674253a60e987f68f5abd0 Binary files /dev/null and b/mapanything/utils/__pycache__/inference.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/misc.cpython-312.pyc b/mapanything/utils/__pycache__/misc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..445dc2d4f5d86ddf938aea3f9065ce6e58d5b872 Binary files /dev/null and b/mapanything/utils/__pycache__/misc.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/moge_inference.cpython-312.pyc b/mapanything/utils/__pycache__/moge_inference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b821c0024ecb2d8d625c5d7cb009343a080ee2fe Binary files /dev/null and b/mapanything/utils/__pycache__/moge_inference.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/viz.cpython-312.pyc b/mapanything/utils/__pycache__/viz.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c53f441843c93958118c631e5ff23d7b5ceb2ac Binary files /dev/null and b/mapanything/utils/__pycache__/viz.cpython-312.pyc differ diff --git a/mapanything/utils/__pycache__/warnings.cpython-312.pyc b/mapanything/utils/__pycache__/warnings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe6332a67f444f39885c07934a04a0da0ac2e11b Binary files /dev/null and b/mapanything/utils/__pycache__/warnings.cpython-312.pyc differ diff --git a/mapanything/utils/cropping.py b/mapanything/utils/cropping.py new file mode 100644 index 0000000000000000000000000000000000000000..709ff34f56276e80dc842d6689a88bc84363679b --- /dev/null +++ b/mapanything/utils/cropping.py @@ -0,0 +1,462 @@ +""" +Utility functions for cropping and resizing data while maintaining proper cameras. + +References: DUSt3R +""" + +import cv2 +import numpy as np +import PIL.Image + +try: + lanczos = PIL.Image.Resampling.LANCZOS + bicubic = PIL.Image.Resampling.BICUBIC +except AttributeError: + lanczos = PIL.Image.LANCZOS + bicubic = PIL.Image.BICUBIC + +from mapanything.utils.geometry import ( + colmap_to_opencv_intrinsics, + opencv_to_colmap_intrinsics, +) + + +class ImageList: + """ + Convenience class to apply the same operation to a whole set of images. + + This class wraps a list of PIL.Image objects and provides methods to perform + operations on all images simultaneously. + """ + + def __init__(self, images): + if not isinstance(images, (tuple, list, set)): + images = [images] + self.images = [] + for image in images: + if not isinstance(image, PIL.Image.Image): + image = PIL.Image.fromarray(image) + self.images.append(image) + + def __len__(self): + """Return the number of images in the list.""" + return len(self.images) + + def to_pil(self): + """ + Convert ImageList back to PIL Image(s). + + Returns: + PIL.Image.Image or tuple: Single PIL Image if list contains one image, + or tuple of PIL Images if multiple images + """ + return tuple(self.images) if len(self.images) > 1 else self.images[0] + + @property + def size(self): + """ + Get the size of images in the list. + + Returns: + tuple: (width, height) of the images + + Raises: + AssertionError: If images have different sizes + """ + sizes = [im.size for im in self.images] + assert all(sizes[0] == s for s in sizes), "All images must have the same size" + return sizes[0] + + def resize(self, *args, **kwargs): + """ + Resize all images with the same parameters. + + Args: + *args, **kwargs: Arguments passed to PIL.Image.resize() + + Returns: + ImageList: New ImageList containing resized images + """ + return ImageList(self._dispatch("resize", *args, **kwargs)) + + def crop(self, *args, **kwargs): + """ + Crop all images with the same parameters. + + Args: + *args, **kwargs: Arguments passed to PIL.Image.crop() + + Returns: + ImageList: New ImageList containing cropped images + """ + return ImageList(self._dispatch("crop", *args, **kwargs)) + + def _dispatch(self, func, *args, **kwargs): + """ + Apply a PIL.Image method to all images in the list. + + Args: + func (str): Name of the PIL.Image method to call + *args, **kwargs: Arguments to pass to the method + + Returns: + list: List of results from applying the method to each image + """ + return [getattr(im, func)(*args, **kwargs) for im in self.images] + + +def resize_with_nearest_interpolation_to_match_aspect_ratio(input_data, img_h, img_w): + """ + Resize input map to match the aspect ratio of an image while ensuring + the input resolution never increases beyond the original. + Uses nearest interpolation for resizing. + + Args: + input_data (np.ndarray): The input map to resize + img_h (int): Height of the target image + img_w (int): Width of the target image + + Returns: + tuple: (resized_input, target_h, target_w) + - resized_input: The resized input map + - target_h: The target height used for resizing + - target_w: The target width used for resizing + """ + # Get the dimensions of the input map + input_h, input_w = input_data.shape[:2] + + # Calculate aspect ratios + img_aspect = img_w / img_h + + # Option 1: Keep input_w fixed and calculate new height + option1_h = int(input_w / img_aspect) + # Option 2: Keep input_h fixed and calculate new width + option2_w = int(input_h * img_aspect) + + # Check if either option would increase a dimension + option1_increases = option1_h > input_h + option2_increases = option2_w > input_w + + if option1_increases and option2_increases: + # Both options would increase a dimension, so we need to scale down both dimensions + # Find the scaling factor that preserves aspect ratio and ensures no dimension increases + scale_h = input_h / img_h + scale_w = input_w / img_w + scale = min(scale_h, scale_w) + + target_input_h = int(img_h * scale) + target_input_w = int(img_w * scale) + elif option1_increases: + # Option 1 would increase height, so use option 2 + target_input_h = input_h + target_input_w = option2_w + elif option2_increases: + # Option 2 would increase width, so use option 1 + target_input_w = input_w + target_input_h = option1_h + else: + # Neither option increases dimensions, choose the one that maintains resolution better + if abs(input_h * input_w - input_w * option1_h) < abs( + input_h * input_w - option2_w * input_h + ): + # Option 1 is better: keep width fixed, adjust height + target_input_w = input_w + target_input_h = option1_h + else: + # Option 2 is better: keep height fixed, adjust width + target_input_h = input_h + target_input_w = option2_w + + # Resize input using nearest interpolation to maintain input values + if target_input_h != input_h or target_input_w != input_w: + resized_input = cv2.resize( + input_data, + (target_input_w, target_input_h), + interpolation=cv2.INTER_NEAREST, + ) + else: + resized_input = input_data + + return resized_input, target_input_h, target_input_w + + +def rescale_image_and_other_optional_info( + image, + output_resolution, + depthmap=None, + camera_intrinsics=None, + force=True, + additional_quantities_to_be_resized_with_nearest=None, +): + """ + Rescale the image and depthmap to the output resolution. + If the image is larger than the output resolution, it is rescaled with lanczos interpolation. + If force is false and the image is smaller than the output resolution, it is not rescaled. + If force is true and the image is smaller than the output resolution, it is rescaled with bicubic interpolation. + Depth and other quantities are rescaled with nearest interpolation. + + Args: + image (PIL.Image.Image or np.ndarray): The input image to be rescaled. + output_resolution (tuple): The desired output resolution as a tuple (width, height). + depthmap (np.ndarray, optional): The depth map associated with the image. Defaults to None. + camera_intrinsics (np.ndarray, optional): The camera intrinsics matrix. Defaults to None. + force (bool, optional): If True, force rescaling even if the image is smaller than the output resolution. Defaults to True. + additional_quantities_to_be_resized_with_nearest (list of np.ndarray, optional): Additional quantities to be rescaled using nearest interpolation. Defaults to None. + + Returns: + tuple: A tuple containing: + - The rescaled image (PIL.Image.Image) + - The rescaled depthmap (numpy.ndarray or None) + - The updated camera intrinsics (numpy.ndarray or None) + - The list of rescaled additional quantities (list of numpy.ndarray or None) + """ + image = ImageList(image) + input_resolution = np.array(image.size) # (W, H) + output_resolution = np.array(output_resolution) + if depthmap is not None: + assert tuple(depthmap.shape[:2]) == image.size[::-1] + if additional_quantities_to_be_resized_with_nearest is not None: + assert all( + tuple(additional_quantity.shape[:2]) == image.size[::-1] + for additional_quantity in additional_quantities_to_be_resized_with_nearest + ) + + # Define output resolution + assert output_resolution.shape == (2,) + scale_final = max(output_resolution / image.size) + 1e-8 + if scale_final >= 1 and not force: # image is already smaller than what is asked + output = ( + image.to_pil(), + depthmap, + camera_intrinsics, + additional_quantities_to_be_resized_with_nearest, + ) + return output + output_resolution = np.floor(input_resolution * scale_final).astype(int) + + # First rescale the image so that it contains the crop + image = image.resize( + tuple(output_resolution), resample=lanczos if scale_final < 1 else bicubic + ) + if depthmap is not None: + depthmap = cv2.resize( + depthmap, + output_resolution, + fx=scale_final, + fy=scale_final, + interpolation=cv2.INTER_NEAREST, + ) + if additional_quantities_to_be_resized_with_nearest is not None: + resized_additional_quantities = [] + for quantity in additional_quantities_to_be_resized_with_nearest: + resized_additional_quantities.append( + cv2.resize( + quantity, + output_resolution, + fx=scale_final, + fy=scale_final, + interpolation=cv2.INTER_NEAREST, + ) + ) + additional_quantities_to_be_resized_with_nearest = resized_additional_quantities + + # No offset here; simple rescaling + if camera_intrinsics is not None: + camera_intrinsics = camera_matrix_of_crop( + camera_intrinsics, input_resolution, output_resolution, scaling=scale_final + ) + + # Return + return ( + image.to_pil(), + depthmap, + camera_intrinsics, + additional_quantities_to_be_resized_with_nearest, + ) + + +def camera_matrix_of_crop( + input_camera_matrix, + input_resolution, + output_resolution, + scaling=1, + offset_factor=0.5, + offset=None, +): + """ + Calculate the camera matrix for a cropped image. + + Args: + input_camera_matrix (numpy.ndarray): Original camera intrinsics matrix + input_resolution (tuple or numpy.ndarray): Original image resolution as (width, height) + output_resolution (tuple or numpy.ndarray): Target image resolution as (width, height) + scaling (float, optional): Scaling factor for the image. Defaults to 1. + offset_factor (float, optional): Factor to determine crop offset. Defaults to 0.5 (centered). + offset (tuple or numpy.ndarray, optional): Explicit offset to use. If None, calculated from offset_factor. + + Returns: + numpy.ndarray: Updated camera matrix for the cropped image + """ + # Margins to offset the origin + margins = np.asarray(input_resolution) * scaling - output_resolution + assert np.all(margins >= 0.0) + if offset is None: + offset = offset_factor * margins + + # Generate new camera parameters + output_camera_matrix_colmap = opencv_to_colmap_intrinsics(input_camera_matrix) + output_camera_matrix_colmap[:2, :] *= scaling + output_camera_matrix_colmap[:2, 2] -= offset + output_camera_matrix = colmap_to_opencv_intrinsics(output_camera_matrix_colmap) + + return output_camera_matrix + + +def crop_image_and_other_optional_info( + image, + crop_bbox, + depthmap=None, + camera_intrinsics=None, + additional_quantities=None, +): + """ + Return a crop of the input view and associated data. + + Args: + image (PIL.Image.Image or numpy.ndarray): The input image to be cropped + crop_bbox (tuple): Crop bounding box as (left, top, right, bottom) + depthmap (numpy.ndarray, optional): Depth map associated with the image + camera_intrinsics (numpy.ndarray, optional): Camera intrinsics matrix + additional_quantities (list of numpy.ndarray, optional): Additional data arrays to crop + + Returns: + tuple: A tuple containing: + - The cropped image + - The cropped depth map (if provided or None) + - Updated camera intrinsics (if provided or None) + - List of cropped additional quantities (if provided or None) + """ + image = ImageList(image) + left, top, right, bottom = crop_bbox + + image = image.crop((left, top, right, bottom)) + if depthmap is not None: + depthmap = depthmap[top:bottom, left:right] + if additional_quantities is not None: + additional_quantities = [ + quantity[top:bottom, left:right] for quantity in additional_quantities + ] + + if camera_intrinsics is not None: + camera_intrinsics = camera_intrinsics.copy() + camera_intrinsics[0, 2] -= left + camera_intrinsics[1, 2] -= top + + return (image.to_pil(), depthmap, camera_intrinsics, additional_quantities) + + +def bbox_from_intrinsics_in_out( + input_camera_matrix, output_camera_matrix, output_resolution +): + """ + Calculate the bounding box for cropping based on input and output camera intrinsics. + + Args: + input_camera_matrix (numpy.ndarray): Original camera intrinsics matrix + output_camera_matrix (numpy.ndarray): Target camera intrinsics matrix + output_resolution (tuple): Target resolution as (width, height) + + Returns: + tuple: Crop bounding box as (left, top, right, bottom) + """ + out_width, out_height = output_resolution + left, top = np.int32( + np.round(input_camera_matrix[:2, 2] - output_camera_matrix[:2, 2]) + ) + crop_bbox = (left, top, left + out_width, top + out_height) + return crop_bbox + + +def crop_resize_if_necessary( + image, + resolution, + depthmap=None, + intrinsics=None, + additional_quantities=None, +): + """ + First downsample image using LANCZOS and then crop if necessary to achieve target resolution. + + This function performs high-quality downsampling followed by cropping to achieve the + desired output resolution while maintaining proper camera intrinsics. + + Args: + image (PIL.Image.Image or numpy.ndarray): The input image to be processed + resolution (tuple): Target resolution as (width, height) + depthmap (numpy.ndarray, optional): Depth map associated with the image + intrinsics (numpy.ndarray, optional): Camera intrinsics matrix + additional_quantities (list of numpy.ndarray, optional): Additional data arrays to process + + Returns: + tuple: A tuple containing the processed image and any provided additional data + (depthmap, intrinsics, additional_quantities) that have been similarly processed + """ + # Convert image to PIL.Image.Image if necessary + if not isinstance(image, PIL.Image.Image): + image = PIL.Image.fromarray(image) + + # Get width and height of image + original_width, original_height = image.size + + # High-quality Lanczos down-scaling + target_rescale_resolution = np.array(resolution) + image, depthmap, intrinsics, additional_quantities = ( + rescale_image_and_other_optional_info( + image=image, + output_resolution=target_rescale_resolution, + depthmap=depthmap, + camera_intrinsics=intrinsics, + additional_quantities_to_be_resized_with_nearest=additional_quantities, + ) + ) + + # Actual cropping (if necessary) + if intrinsics is not None: + new_intrinsics = camera_matrix_of_crop( + input_camera_matrix=intrinsics, + input_resolution=image.size, + output_resolution=resolution, + offset_factor=0.5, + ) + crop_bbox = bbox_from_intrinsics_in_out( + input_camera_matrix=intrinsics, + output_camera_matrix=new_intrinsics, + output_resolution=resolution, + ) + else: + # Create a centered crop if no intrinsics are available + w, h = image.size + target_w, target_h = resolution + left = (w - target_w) // 2 + top = (h - target_h) // 2 + crop_bbox = (left, top, left + target_w, top + target_h) + + image, depthmap, new_intrinsics, additional_quantities = ( + crop_image_and_other_optional_info( + image=image, + crop_bbox=crop_bbox, + depthmap=depthmap, + camera_intrinsics=intrinsics, + additional_quantities=additional_quantities, + ) + ) + + # Return the output + output = (image,) + if depthmap is not None: + output += (depthmap,) + if new_intrinsics is not None: + output += (new_intrinsics,) + if additional_quantities is not None: + output += (additional_quantities,) + return output diff --git a/mapanything/utils/device.py b/mapanything/utils/device.py new file mode 100644 index 0000000000000000000000000000000000000000..dc9ca662767457f89d8801bdfba786bb99c6ec2a --- /dev/null +++ b/mapanything/utils/device.py @@ -0,0 +1,83 @@ +""" +Utility functions for managing computation device +""" + +import numpy as np +import torch + + +def to_device(batch, device, callback=None, non_blocking=False): + """ + Transfer data to another device (i.e. GPU, CPU:torch, CPU:numpy). + + This function recursively processes nested data structures (lists, tuples, dicts) + and transfers each tensor to the specified device. + + Args: + batch: Data to transfer (list, tuple, dict of tensors or other objects) + device: Target device - pytorch device (e.g., 'cuda', 'cpu') or 'numpy' + callback: Optional function that would be called on every element before processing + non_blocking: If True, allows asynchronous copy to GPU (may be faster) + + Returns: + Data with the same structure as input but with tensors transferred to target device + """ + if callback: + batch = callback(batch) + + if isinstance(batch, dict): + return { + k: to_device(v, device, non_blocking=non_blocking) for k, v in batch.items() + } + + if isinstance(batch, (tuple, list)): + return type(batch)( + to_device(x, device, non_blocking=non_blocking) for x in batch + ) + + x = batch + if device == "numpy": + if isinstance(x, torch.Tensor): + x = x.detach().cpu().numpy() + elif x is not None: + if isinstance(x, np.ndarray): + x = torch.from_numpy(x) + if torch.is_tensor(x): + x = x.to(device, non_blocking=non_blocking) + return x + + +def to_numpy(x): + """Convert data to numpy arrays. + + Args: + x: Input data (can be tensor, array, or nested structure) + + Returns: + Data with the same structure but with tensors converted to numpy arrays + """ + return to_device(x, "numpy") + + +def to_cpu(x): + """Transfer data to CPU. + + Args: + x: Input data (can be tensor, array, or nested structure) + + Returns: + Data with the same structure but with tensors moved to CPU + """ + return to_device(x, "cpu") + + +def to_cuda(x): + """Transfer data to CUDA device (GPU). + + Args: + x: Input data (can be tensor, array, or nested structure) + + Returns: + Data with the same structure but with tensors moved to GPU + """ + return to_device(x, "cuda") diff --git a/mapanything/utils/geometry.py b/mapanything/utils/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..77e87b6f83612412e70487581111754368d15cbc --- /dev/null +++ b/mapanything/utils/geometry.py @@ -0,0 +1,2188 @@ +""" +Utilities for geometry operations. + +References: DUSt3R, MoGe +""" + +from numbers import Number +from typing import Tuple, Union + +import einops as ein +import numpy as np +import torch +import torch.nn.functional as F + +from mapanything.utils.misc import invalid_to_zeros +from mapanything.utils.warnings import no_warnings + + +def depthmap_to_camera_frame(depthmap, intrinsics): + """ + Convert depth image to a pointcloud in camera frame. + + Args: + - depthmap: HxW or BxHxW torch tensor + - intrinsics: 3x3 or Bx3x3 torch tensor + + Returns: + pointmap in camera frame (HxWx3 or BxHxWx3 tensor), and a mask specifying valid pixels. + """ + # Add batch dimension if not present + if depthmap.dim() == 2: + depthmap = depthmap.unsqueeze(0) + intrinsics = intrinsics.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + batch_size, height, width = depthmap.shape + device = depthmap.device + + # Compute 3D point in camera frame associated with each pixel + x_grid, y_grid = torch.meshgrid( + torch.arange(width, device=device).float(), + torch.arange(height, device=device).float(), + indexing="xy", + ) + x_grid = x_grid.unsqueeze(0).expand(batch_size, -1, -1) + y_grid = y_grid.unsqueeze(0).expand(batch_size, -1, -1) + + fx = intrinsics[:, 0, 0].view(-1, 1, 1) + fy = intrinsics[:, 1, 1].view(-1, 1, 1) + cx = intrinsics[:, 0, 2].view(-1, 1, 1) + cy = intrinsics[:, 1, 2].view(-1, 1, 1) + + depth_z = depthmap + xx = (x_grid - cx) * depth_z / fx + yy = (y_grid - cy) * depth_z / fy + pts3d_cam = torch.stack((xx, yy, depth_z), dim=-1) + + # Compute mask of valid non-zero depth pixels + valid_mask = depthmap > 0.0 + + # Remove batch dimension if it was added + if squeeze_batch_dim: + pts3d_cam = pts3d_cam.squeeze(0) + valid_mask = valid_mask.squeeze(0) + + return pts3d_cam, valid_mask + + +def depthmap_to_world_frame(depthmap, intrinsics, camera_pose=None): + """ + Convert depth image to a pointcloud in world frame. + + Args: + - depthmap: HxW or BxHxW torch tensor + - intrinsics: 3x3 or Bx3x3 torch tensor + - camera_pose: 4x4 or Bx4x4 torch tensor + + Returns: + pointmap in world frame (HxWx3 or BxHxWx3 tensor), and a mask specifying valid pixels. + """ + pts3d_cam, valid_mask = depthmap_to_camera_frame(depthmap, intrinsics) + + if camera_pose is not None: + # Add batch dimension if not present + if camera_pose.dim() == 2: + camera_pose = camera_pose.unsqueeze(0) + pts3d_cam = pts3d_cam.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Convert points from camera frame to world frame + pts3d_cam_homo = torch.cat( + [pts3d_cam, torch.ones_like(pts3d_cam[..., :1])], dim=-1 + ) + pts3d_world = ein.einsum( + camera_pose, pts3d_cam_homo, "b i k, b h w k -> b h w i" + ) + pts3d_world = pts3d_world[..., :3] + + # Remove batch dimension if it was added + if squeeze_batch_dim: + pts3d_world = pts3d_world.squeeze(0) + else: + pts3d_world = pts3d_cam + + return pts3d_world, valid_mask + + +def transform_pts3d(pts3d, transformation): + """ + Transform 3D points using a 4x4 transformation matrix. + + Args: + - pts3d: HxWx3 or BxHxWx3 torch tensor + - transformation: 4x4 or Bx4x4 torch tensor + + Returns: + transformed points (HxWx3 or BxHxWx3 tensor) + """ + # Add batch dimension if not present + if pts3d.dim() == 3: + pts3d = pts3d.unsqueeze(0) + transformation = transformation.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Convert points to homogeneous coordinates + pts3d_homo = torch.cat([pts3d, torch.ones_like(pts3d[..., :1])], dim=-1) + + # Transform points + transformed_pts3d = ein.einsum( + transformation, pts3d_homo, "b i k, b h w k -> b h w i" + ) + transformed_pts3d = transformed_pts3d[..., :3] + + # Remove batch dimension if it was added + if squeeze_batch_dim: + transformed_pts3d = transformed_pts3d.squeeze(0) + + return transformed_pts3d + + +def project_pts3d_to_image(pts3d, intrinsics, return_z_dim): + """ + Project 3D points to image plane (assumes pinhole camera model with no distortion). + + Args: + - pts3d: HxWx3 or BxHxWx3 torch tensor + - intrinsics: 3x3 or Bx3x3 torch tensor + - return_z_dim: bool, whether to return the third dimension of the projected points + + Returns: + projected points (HxWx2) + """ + if pts3d.dim() == 3: + pts3d = pts3d.unsqueeze(0) + intrinsics = intrinsics.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Project points to image plane + projected_pts2d = ein.einsum(intrinsics, pts3d, "b i k, b h w k -> b h w i") + projected_pts2d[..., :2] /= projected_pts2d[..., 2].unsqueeze(-1).clamp(min=1e-6) + + # Remove the z dimension if not required + if not return_z_dim: + projected_pts2d = projected_pts2d[..., :2] + + # Remove batch dimension if it was added + if squeeze_batch_dim: + projected_pts2d = projected_pts2d.squeeze(0) + + return projected_pts2d + + +def get_rays_in_camera_frame(intrinsics, height, width, normalize_to_unit_sphere): + """ + Convert camera intrinsics to a raymap (ray origins + directions) in camera frame. + Note: Currently only supports pinhole camera model. + + Args: + - intrinsics: 3x3 or Bx3x3 torch tensor + - height: int + - width: int + - normalize_to_unit_sphere: bool + + Returns: + - ray_origins: (HxWx3 or BxHxWx3) tensor + - ray_directions: (HxWx3 or BxHxWx3) tensor + """ + # Add batch dimension if not present + if intrinsics.dim() == 2: + intrinsics = intrinsics.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + batch_size = intrinsics.shape[0] + device = intrinsics.device + + # Compute rays in camera frame associated with each pixel + x_grid, y_grid = torch.meshgrid( + torch.arange(width, device=device).float(), + torch.arange(height, device=device).float(), + indexing="xy", + ) + x_grid = x_grid.unsqueeze(0).expand(batch_size, -1, -1) + y_grid = y_grid.unsqueeze(0).expand(batch_size, -1, -1) + + fx = intrinsics[:, 0, 0].view(-1, 1, 1) + fy = intrinsics[:, 1, 1].view(-1, 1, 1) + cx = intrinsics[:, 0, 2].view(-1, 1, 1) + cy = intrinsics[:, 1, 2].view(-1, 1, 1) + + ray_origins = torch.zeros((batch_size, height, width, 3), device=device) + xx = (x_grid - cx) / fx + yy = (y_grid - cy) / fy + ray_directions = torch.stack((xx, yy, torch.ones_like(xx)), dim=-1) + + # Normalize ray directions to unit sphere if required (else rays will lie on unit plane) + if normalize_to_unit_sphere: + ray_directions = ray_directions / torch.norm( + ray_directions, dim=-1, keepdim=True + ) + + # Remove batch dimension if it was added + if squeeze_batch_dim: + ray_origins = ray_origins.squeeze(0) + ray_directions = ray_directions.squeeze(0) + + return ray_origins, ray_directions + + +def get_rays_in_world_frame( + intrinsics, height, width, normalize_to_unit_sphere, camera_pose=None +): + """ + Convert camera intrinsics & camera_pose (if provided) to a raymap (ray origins + directions) in camera or world frame (if camera_pose is provided). + Note: Currently only supports pinhole camera model. + + Args: + - intrinsics: 3x3 or Bx3x3 torch tensor + - height: int + - width: int + - normalize_to_unit_sphere: bool + - camera_pose: 4x4 or Bx4x4 torch tensor + + Returns: + - ray_origins: (HxWx3 or BxHxWx3) tensor + - ray_directions: (HxWx3 or BxHxWx3) tensor + """ + # Get rays in camera frame + ray_origins, ray_directions = get_rays_in_camera_frame( + intrinsics, height, width, normalize_to_unit_sphere + ) + + if camera_pose is not None: + # Add batch dimension if not present + if camera_pose.dim() == 2: + camera_pose = camera_pose.unsqueeze(0) + ray_origins = ray_origins.unsqueeze(0) + ray_directions = ray_directions.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Convert rays from camera frame to world frame + ray_origins_homo = torch.cat( + [ray_origins, torch.ones_like(ray_origins[..., :1])], dim=-1 + ) + ray_directions_homo = torch.cat( + [ray_directions, torch.zeros_like(ray_directions[..., :1])], dim=-1 + ) + ray_origins_world = ein.einsum( + camera_pose, ray_origins_homo, "b i k, b h w k -> b h w i" + ) + ray_directions_world = ein.einsum( + camera_pose, ray_directions_homo, "b i k, b h w k -> b h w i" + ) + ray_origins_world = ray_origins_world[..., :3] + ray_directions_world = ray_directions_world[..., :3] + + # Remove batch dimension if it was added + if squeeze_batch_dim: + ray_origins_world = ray_origins_world.squeeze(0) + ray_directions_world = ray_directions_world.squeeze(0) + else: + ray_origins_world = ray_origins + ray_directions_world = ray_directions + + return ray_origins_world, ray_directions_world + + +def recover_pinhole_intrinsics_from_ray_directions( + ray_directions, use_geometric_calculation=False +): + """ + Recover pinhole camera intrinsics from ray directions, supporting both batched and non-batched inputs. + + Args: + ray_directions: Tensor of shape [H, W, 3] or [B, H, W, 3] containing unit normalized ray directions + + Returns: + Dictionary containing camera intrinsics (fx, fy, cx, cy) as tensors + """ + # Add batch dimension if not present + if ray_directions.dim() == 3: # [H, W, 3] + ray_directions = ray_directions.unsqueeze(0) # [1, H, W, 3] + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + batch_size, height, width, _ = ray_directions.shape + device = ray_directions.device + + # Create pixel coordinate grid + x_grid, y_grid = torch.meshgrid( + torch.arange(width, device=device).float(), + torch.arange(height, device=device).float(), + indexing="xy", + ) + + # Expand grid for all batches + x_grid = x_grid.unsqueeze(0).expand(batch_size, -1, -1) # [B, H, W] + y_grid = y_grid.unsqueeze(0).expand(batch_size, -1, -1) # [B, H, W] + + # Determine if high resolution or not + is_high_res = height * width > 1000000 + + if is_high_res or use_geometric_calculation: + # For high-resolution cases, use direct geometric calculation + # Define key points + center_h, center_w = height // 2, width // 2 + quarter_w, three_quarter_w = width // 4, 3 * width // 4 + quarter_h, three_quarter_h = height // 4, 3 * height // 4 + + # Get rays at key points + center_rays = ray_directions[:, center_h, center_w, :].clone() # [B, 3] + left_rays = ray_directions[:, center_h, quarter_w, :].clone() # [B, 3] + right_rays = ray_directions[:, center_h, three_quarter_w, :].clone() # [B, 3] + top_rays = ray_directions[:, quarter_h, center_w, :].clone() # [B, 3] + bottom_rays = ray_directions[:, three_quarter_h, center_w, :].clone() # [B, 3] + + # Normalize rays to have dz = 1 + center_rays = center_rays / center_rays[:, 2].unsqueeze(1) # [B, 3] + left_rays = left_rays / left_rays[:, 2].unsqueeze(1) # [B, 3] + right_rays = right_rays / right_rays[:, 2].unsqueeze(1) # [B, 3] + top_rays = top_rays / top_rays[:, 2].unsqueeze(1) # [B, 3] + bottom_rays = bottom_rays / bottom_rays[:, 2].unsqueeze(1) # [B, 3] + + # Calculate fx directly (vectorized across batch) + fx_left = (quarter_w - center_w) / (left_rays[:, 0] - center_rays[:, 0]) + fx_right = (three_quarter_w - center_w) / (right_rays[:, 0] - center_rays[:, 0]) + fx = (fx_left + fx_right) / 2 # Average for robustness + + # Calculate cx + cx = center_w - fx * center_rays[:, 0] + + # Calculate fy and cy + fy_top = (quarter_h - center_h) / (top_rays[:, 1] - center_rays[:, 1]) + fy_bottom = (three_quarter_h - center_h) / ( + bottom_rays[:, 1] - center_rays[:, 1] + ) + fy = (fy_top + fy_bottom) / 2 + + cy = center_h - fy * center_rays[:, 1] + else: + # For standard resolution, use regression with sampling for efficiency + # Sample a grid of points (but more dense than for high-res) + step_h = max(1, height // 50) + step_w = max(1, width // 50) + + h_indices = torch.arange(0, height, step_h, device=device) + w_indices = torch.arange(0, width, step_w, device=device) + + # Extract subset of coordinates + x_sampled = x_grid[:, h_indices[:, None], w_indices[None, :]] # [B, H', W'] + y_sampled = y_grid[:, h_indices[:, None], w_indices[None, :]] # [B, H', W'] + rays_sampled = ray_directions[ + :, h_indices[:, None], w_indices[None, :], : + ] # [B, H', W', 3] + + # Reshape for linear regression + x_flat = x_sampled.reshape(batch_size, -1) # [B, N] + y_flat = y_sampled.reshape(batch_size, -1) # [B, N] + + # Extract ray direction components + dx = rays_sampled[..., 0].reshape(batch_size, -1) # [B, N] + dy = rays_sampled[..., 1].reshape(batch_size, -1) # [B, N] + dz = rays_sampled[..., 2].reshape(batch_size, -1) # [B, N] + + # Compute ratios for linear regression + ratio_x = dx / dz # [B, N] + ratio_y = dy / dz # [B, N] + + # Since torch.linalg.lstsq doesn't support batched input, we'll use a different approach + # For x-direction: x = cx + fx * (dx/dz) + # We can solve this using normal equations: A^T A x = A^T b + # Create design matrices + ones = torch.ones_like(x_flat) # [B, N] + A_x = torch.stack([ones, ratio_x], dim=2) # [B, N, 2] + b_x = x_flat.unsqueeze(2) # [B, N, 1] + + # Compute A^T A and A^T b for each batch + ATA_x = torch.bmm(A_x.transpose(1, 2), A_x) # [B, 2, 2] + ATb_x = torch.bmm(A_x.transpose(1, 2), b_x) # [B, 2, 1] + + # Solve the system for each batch + solution_x = torch.linalg.solve(ATA_x, ATb_x).squeeze(2) # [B, 2] + cx, fx = solution_x[:, 0], solution_x[:, 1] + + # Repeat for y-direction + A_y = torch.stack([ones, ratio_y], dim=2) # [B, N, 2] + b_y = y_flat.unsqueeze(2) # [B, N, 1] + + ATA_y = torch.bmm(A_y.transpose(1, 2), A_y) # [B, 2, 2] + ATb_y = torch.bmm(A_y.transpose(1, 2), b_y) # [B, 2, 1] + + solution_y = torch.linalg.solve(ATA_y, ATb_y).squeeze(2) # [B, 2] + cy, fy = solution_y[:, 0], solution_y[:, 1] + + # Create intrinsics matrices + batch_size = fx.shape[0] + intrinsics = torch.zeros(batch_size, 3, 3, device=ray_directions.device) + + # Fill in the intrinsics matrices + intrinsics[:, 0, 0] = fx # focal length x + intrinsics[:, 1, 1] = fy # focal length y + intrinsics[:, 0, 2] = cx # principal point x + intrinsics[:, 1, 2] = cy # principal point y + intrinsics[:, 2, 2] = 1.0 # bottom-right element is always 1 + + # Remove batch dimension if it was added + if squeeze_batch_dim: + intrinsics = intrinsics.squeeze(0) + + return intrinsics + + +def transform_rays(ray_origins, ray_directions, transformation): + """ + Transform 6D rays (ray origins and ray directions) using a 4x4 transformation matrix. + + Args: + - ray_origins: HxWx3 or BxHxWx3 torch tensor + - ray_directions: HxWx3 or BxHxWx3 torch tensor + - transformation: 4x4 or Bx4x4 torch tensor + - normalize_to_unit_sphere: bool, whether to normalize the transformed ray directions to unit length + + Returns: + transformed ray_origins (HxWx3 or BxHxWx3 tensor) and ray_directions (HxWx3 or BxHxWx3 tensor) + """ + # Add batch dimension if not present + if ray_origins.dim() == 3: + ray_origins = ray_origins.unsqueeze(0) + ray_directions = ray_directions.unsqueeze(0) + transformation = transformation.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Transform ray origins and directions + ray_origins_homo = torch.cat( + [ray_origins, torch.ones_like(ray_origins[..., :1])], dim=-1 + ) + ray_directions_homo = torch.cat( + [ray_directions, torch.zeros_like(ray_directions[..., :1])], dim=-1 + ) + transformed_ray_origins = ein.einsum( + transformation, ray_origins_homo, "b i k, b h w k -> b h w i" + ) + transformed_ray_directions = ein.einsum( + transformation, ray_directions_homo, "b i k, b h w k -> b h w i" + ) + transformed_ray_origins = transformed_ray_origins[..., :3] + transformed_ray_directions = transformed_ray_directions[..., :3] + + # Remove batch dimension if it was added + if squeeze_batch_dim: + transformed_ray_origins = transformed_ray_origins.squeeze(0) + transformed_ray_directions = transformed_ray_directions.squeeze(0) + + return transformed_ray_origins, transformed_ray_directions + + +def convert_z_depth_to_depth_along_ray(z_depth, intrinsics): + """ + Convert z-depth image to depth along camera rays. + + Args: + - z_depth: HxW or BxHxW torch tensor + - intrinsics: 3x3 or Bx3x3 torch tensor + + Returns: + - depth_along_ray: HxW or BxHxW torch tensor + """ + # Add batch dimension if not present + if z_depth.dim() == 2: + z_depth = z_depth.unsqueeze(0) + intrinsics = intrinsics.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Get rays in camera frame + batch_size, height, width = z_depth.shape + _, ray_directions = get_rays_in_camera_frame( + intrinsics, height, width, normalize_to_unit_sphere=False + ) + + # Compute depth along ray + pts3d_cam = z_depth[..., None] * ray_directions + depth_along_ray = torch.norm(pts3d_cam, dim=-1) + + # Remove batch dimension if it was added + if squeeze_batch_dim: + depth_along_ray = depth_along_ray.squeeze(0) + + return depth_along_ray + + +def convert_raymap_z_depth_quats_to_pointmap(ray_origins, ray_directions, depth, quats): + """ + Convert raymap (ray origins + directions on unit plane), z-depth and + unit quaternions (representing rotation) to a pointmap in world frame. + + Args: + - ray_origins: (HxWx3 or BxHxWx3) torch tensor + - ray_directions: (HxWx3 or BxHxWx3) torch tensor + - depth: (HxWx1 or BxHxWx1) torch tensor + - quats: (HxWx4 or BxHxWx4) torch tensor (unit quaternions and notation is (x, y, z, w)) + + Returns: + - pointmap: (HxWx3 or BxHxWx3) torch tensor + """ + # Add batch dimension if not present + if ray_origins.dim() == 3: + ray_origins = ray_origins.unsqueeze(0) + ray_directions = ray_directions.unsqueeze(0) + depth = depth.unsqueeze(0) + quats = quats.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + batch_size, height, width, _ = depth.shape + device = depth.device + + # Normalize the quaternions to ensure they are unit quaternions + quats = quats / torch.norm(quats, dim=-1, keepdim=True) + + # Convert quaternions to pixel-wise rotation matrices + qx, qy, qz, qw = quats[..., 0], quats[..., 1], quats[..., 2], quats[..., 3] + rot_mat = ( + torch.stack( + [ + qw**2 + qx**2 - qy**2 - qz**2, + 2 * (qx * qy - qw * qz), + 2 * (qw * qy + qx * qz), + 2 * (qw * qz + qx * qy), + qw**2 - qx**2 + qy**2 - qz**2, + 2 * (qy * qz - qw * qx), + 2 * (qx * qz - qw * qy), + 2 * (qw * qx + qy * qz), + qw**2 - qx**2 - qy**2 + qz**2, + ], + dim=-1, + ) + .reshape(batch_size, height, width, 3, 3) + .to(device) + ) + + # Compute 3D points in local camera frame + pts3d_local = depth * ray_directions + + # Rotate the local points using the quaternions + rotated_pts3d_local = ein.einsum( + rot_mat, pts3d_local, "b h w i k, b h w k -> b h w i" + ) + + # Compute 3D point in world frame associated with each pixel + pts3d = ray_origins + rotated_pts3d_local + + # Remove batch dimension if it was added + if squeeze_batch_dim: + pts3d = pts3d.squeeze(0) + + return pts3d + + +def quaternion_to_rotation_matrix(quat): + """ + Convert a quaternion into a 3x3 rotation matrix. + + Args: + - quat: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + + Returns: + - rot_matrix: 3x3 or Bx3x3 torch tensor + """ + if quat.dim() == 1: + quat = quat.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Ensure the quaternion is normalized + quat = quat / quat.norm(dim=1, keepdim=True) + x, y, z, w = quat.unbind(dim=1) + + # Compute the rotation matrix elements + xx = x * x + yy = y * y + zz = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + + # Construct the rotation matrix + rot_matrix = torch.stack( + [ + 1 - 2 * (yy + zz), + 2 * (xy - wz), + 2 * (xz + wy), + 2 * (xy + wz), + 1 - 2 * (xx + zz), + 2 * (yz - wx), + 2 * (xz - wy), + 2 * (yz + wx), + 1 - 2 * (xx + yy), + ], + dim=1, + ).view(-1, 3, 3) + + # Squeeze batch dimension if it was unsqueezed + if squeeze_batch_dim: + rot_matrix = rot_matrix.squeeze(0) + + return rot_matrix + + +def rotation_matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part last, as tensor of shape (..., 4). + Quaternion Order: XYZW or say ijkr, scalar-last + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + out = quat_candidates[ + F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : + ].reshape(batch_dim + (4,)) + + # Convert from rijk to ijkr + out = out[..., [1, 2, 3, 0]] + + out = standardize_quaternion(out) + + return out + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + if torch.is_grad_enabled(): + ret[positive_mask] = torch.sqrt(x[positive_mask]) + else: + ret = torch.where(positive_mask, torch.sqrt(x), ret) + return ret + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) + + +def quaternion_inverse(quat): + """ + Compute the inverse of a quaternion. + + Args: + - quat: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + + Returns: + - inv_quat: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + """ + # Unsqueeze batch dimension if not present + if quat.dim() == 1: + quat = quat.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Compute the inverse + quat_conj = quat.clone() + quat_conj[:, :3] = -quat_conj[:, :3] + quat_norm = torch.sum(quat * quat, dim=1, keepdim=True) + inv_quat = quat_conj / quat_norm + + # Squeeze batch dimension if it was unsqueezed + if squeeze_batch_dim: + inv_quat = inv_quat.squeeze(0) + + return inv_quat + + +def quaternion_multiply(q1, q2): + """ + Multiply two quaternions. + + Args: + - q1: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + - q2: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + + Returns: + - qm: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + """ + # Unsqueeze batch dimension if not present + if q1.dim() == 1: + q1 = q1.unsqueeze(0) + q2 = q2.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Unbind the quaternions + x1, y1, z1, w1 = q1.unbind(dim=1) + x2, y2, z2, w2 = q2.unbind(dim=1) + + # Compute the product + x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 + y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 + z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 + w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 + + # Stack the components + qm = torch.stack([x, y, z, w], dim=1) + + # Squeeze batch dimension if it was unsqueezed + if squeeze_batch_dim: + qm = qm.squeeze(0) + + return qm + + +def transform_pose_using_quats_and_trans_2_to_1(quats1, trans1, quats2, trans2): + """ + Transform quats and translation of pose2 from absolute frame (pose2 to world) to relative frame (pose2 to pose1). + + Args: + - quats1: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + - trans1: 3 or Bx3 torch tensor + - quats2: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + - trans2: 3 or Bx3 torch tensor + + Returns: + - quats: 4 or Bx4 torch tensor (unit quaternions and notation is (x, y, z, w)) + - trans: 3 or Bx3 torch tensor + """ + # Unsqueeze batch dimension if not present + if quats1.dim() == 1: + quats1 = quats1.unsqueeze(0) + trans1 = trans1.unsqueeze(0) + quats2 = quats2.unsqueeze(0) + trans2 = trans2.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + # Compute the inverse of view1's pose + inv_quats1 = quaternion_inverse(quats1) + R1_inv = quaternion_to_rotation_matrix(inv_quats1) + t1_inv = -1 * ein.einsum(R1_inv, trans1, "b i j, b j -> b i") + + # Transform view2's pose to view1's frame + quats = quaternion_multiply(inv_quats1, quats2) + trans = ein.einsum(R1_inv, trans2, "b i j, b j -> b i") + t1_inv + + # Squeeze batch dimension if it was unsqueezed + if squeeze_batch_dim: + quats = quats.squeeze(0) + trans = trans.squeeze(0) + + return quats, trans + + +def convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap( + ray_directions, depth_along_ray, pose_trans, pose_quats +): + """ + Convert ray directions, depth along ray, pose translation, and + unit quaternions (representing pose rotation) to a pointmap in world frame. + + Args: + - ray_directions: (HxWx3 or BxHxWx3) torch tensor + - depth_along_ray: (HxWx1 or BxHxWx1) torch tensor + - pose_trans: (3 or Bx3) torch tensor + - pose_quats: (4 or Bx4) torch tensor (unit quaternions and notation is (x, y, z, w)) + + Returns: + - pointmap: (HxWx3 or BxHxWx3) torch tensor + """ + # Add batch dimension if not present + if ray_directions.dim() == 3: + ray_directions = ray_directions.unsqueeze(0) + depth_along_ray = depth_along_ray.unsqueeze(0) + pose_trans = pose_trans.unsqueeze(0) + pose_quats = pose_quats.unsqueeze(0) + squeeze_batch_dim = True + else: + squeeze_batch_dim = False + + batch_size, height, width, _ = depth_along_ray.shape + device = depth_along_ray.device + + # Normalize the quaternions to ensure they are unit quaternions + pose_quats = pose_quats / torch.norm(pose_quats, dim=-1, keepdim=True) + + # Convert quaternions to rotation matrices (B x 3 x 3) + rot_mat = quaternion_to_rotation_matrix(pose_quats) + + # Get pose matrix (B x 4 x 4) + pose_mat = torch.eye(4, device=device).unsqueeze(0).repeat(batch_size, 1, 1) + pose_mat[:, :3, :3] = rot_mat + pose_mat[:, :3, 3] = pose_trans + + # Compute 3D points in local camera frame + pts3d_local = depth_along_ray * ray_directions + + # Compute 3D points in world frame + pts3d_homo = torch.cat([pts3d_local, torch.ones_like(pts3d_local[..., :1])], dim=-1) + pts3d_world = ein.einsum(pose_mat, pts3d_homo, "b i k, b h w k -> b h w i") + pts3d_world = pts3d_world[..., :3] + + # Remove batch dimension if it was added + if squeeze_batch_dim: + pts3d_world = pts3d_world.squeeze(0) + + return pts3d_world + + +def xy_grid( + W, + H, + device=None, + origin=(0, 0), + unsqueeze=None, + cat_dim=-1, + homogeneous=False, + **arange_kw, +): + """ + Generate a coordinate grid of shape (H,W,2) or (H,W,3) if homogeneous=True. + + Args: + W (int): Width of the grid + H (int): Height of the grid + device (torch.device, optional): Device to place the grid on. If None, uses numpy arrays + origin (tuple, optional): Origin coordinates (x,y) for the grid. Default is (0,0) + unsqueeze (int, optional): Dimension to unsqueeze in the output tensors + cat_dim (int, optional): Dimension to concatenate the x,y coordinates. If None, returns tuple + homogeneous (bool, optional): If True, adds a third dimension of ones to make homogeneous coordinates + **arange_kw: Additional keyword arguments passed to np.arange or torch.arange + + Returns: + numpy.ndarray or torch.Tensor: Coordinate grid where: + - output[j,i,0] = i + origin[0] (x-coordinate) + - output[j,i,1] = j + origin[1] (y-coordinate) + - output[j,i,2] = 1 (if homogeneous=True) + """ + if device is None: + # numpy + arange, meshgrid, stack, ones = np.arange, np.meshgrid, np.stack, np.ones + else: + # torch + def arange(*a, **kw): + return torch.arange(*a, device=device, **kw) + + meshgrid, stack = torch.meshgrid, torch.stack + + def ones(*a): + return torch.ones(*a, device=device) + + tw, th = [arange(o, o + s, **arange_kw) for s, o in zip((W, H), origin)] + grid = meshgrid(tw, th, indexing="xy") + if homogeneous: + grid = grid + (ones((H, W)),) + if unsqueeze is not None: + grid = (grid[0].unsqueeze(unsqueeze), grid[1].unsqueeze(unsqueeze)) + if cat_dim is not None: + grid = stack(grid, cat_dim) + + return grid + + +def geotrf(Trf, pts, ncol=None, norm=False): + """ + Apply a geometric transformation to a set of 3-D points. + + Args: + Trf: 3x3 or 4x4 projection matrix (typically a Homography) or batch of matrices + with shape (B, 3, 3) or (B, 4, 4) + pts: numpy/torch/tuple of coordinates with shape (..., 2) or (..., 3) + ncol: int, number of columns of the result (2 or 3) + norm: float, if not 0, the result is projected on the z=norm plane + (homogeneous normalization) + + Returns: + Array or tensor of projected points with the same type as input and shape (..., ncol) + """ + assert Trf.ndim >= 2 + if isinstance(Trf, np.ndarray): + pts = np.asarray(pts) + elif isinstance(Trf, torch.Tensor): + pts = torch.as_tensor(pts, dtype=Trf.dtype) + + # Adapt shape if necessary + output_reshape = pts.shape[:-1] + ncol = ncol or pts.shape[-1] + + # Optimized code + if ( + isinstance(Trf, torch.Tensor) + and isinstance(pts, torch.Tensor) + and Trf.ndim == 3 + and pts.ndim == 4 + ): + d = pts.shape[3] + if Trf.shape[-1] == d: + pts = torch.einsum("bij, bhwj -> bhwi", Trf, pts) + elif Trf.shape[-1] == d + 1: + pts = ( + torch.einsum("bij, bhwj -> bhwi", Trf[:, :d, :d], pts) + + Trf[:, None, None, :d, d] + ) + else: + raise ValueError(f"bad shape, not ending with 3 or 4, for {pts.shape=}") + else: + if Trf.ndim >= 3: + n = Trf.ndim - 2 + assert Trf.shape[:n] == pts.shape[:n], "batch size does not match" + Trf = Trf.reshape(-1, Trf.shape[-2], Trf.shape[-1]) + + if pts.ndim > Trf.ndim: + # Trf == (B,d,d) & pts == (B,H,W,d) --> (B, H*W, d) + pts = pts.reshape(Trf.shape[0], -1, pts.shape[-1]) + elif pts.ndim == 2: + # Trf == (B,d,d) & pts == (B,d) --> (B, 1, d) + pts = pts[:, None, :] + + if pts.shape[-1] + 1 == Trf.shape[-1]: + Trf = Trf.swapaxes(-1, -2) # transpose Trf + pts = pts @ Trf[..., :-1, :] + Trf[..., -1:, :] + elif pts.shape[-1] == Trf.shape[-1]: + Trf = Trf.swapaxes(-1, -2) # transpose Trf + pts = pts @ Trf + else: + pts = Trf @ pts.T + if pts.ndim >= 2: + pts = pts.swapaxes(-1, -2) + + if norm: + pts = pts / pts[..., -1:] # DONT DO /=, it will lead to a bug + if norm != 1: + pts *= norm + + res = pts[..., :ncol].reshape(*output_reshape, ncol) + + return res + + +def inv(mat): + """ + Invert a torch or numpy matrix + """ + if isinstance(mat, torch.Tensor): + return torch.linalg.inv(mat) + if isinstance(mat, np.ndarray): + return np.linalg.inv(mat) + raise ValueError(f"bad matrix type = {type(mat)}") + + +def closed_form_pose_inverse( + pose_matrices, rotation_matrices=None, translation_vectors=None +): + """ + Compute the inverse of each 4x4 (or 3x4) SE3 pose matrices in a batch. + + If `rotation_matrices` and `translation_vectors` are provided, they must correspond to the rotation and translation + components of `pose_matrices`. Otherwise, they will be extracted from `pose_matrices`. + + Args: + pose_matrices: Nx4x4 or Nx3x4 array or tensor of SE3 matrices. + rotation_matrices (optional): Nx3x3 array or tensor of rotation matrices. + translation_vectors (optional): Nx3x1 array or tensor of translation vectors. + + Returns: + Inverted SE3 matrices with the same type and device as input `pose_matrices`. + + Shapes: + pose_matrices: (N, 4, 4) + rotation_matrices: (N, 3, 3) + translation_vectors: (N, 3, 1) + """ + # Check if pose_matrices is a numpy array or a torch tensor + is_numpy = isinstance(pose_matrices, np.ndarray) + + # Validate shapes + if pose_matrices.shape[-2:] != (4, 4) and pose_matrices.shape[-2:] != (3, 4): + raise ValueError( + f"pose_matrices must be of shape (N,4,4), got {pose_matrices.shape}." + ) + + # Extract rotation_matrices and translation_vectors if not provided + if rotation_matrices is None: + rotation_matrices = pose_matrices[:, :3, :3] + if translation_vectors is None: + translation_vectors = pose_matrices[:, :3, 3:] + + # Compute the inverse of input SE3 matrices + if is_numpy: + rotation_transposed = np.transpose(rotation_matrices, (0, 2, 1)) + new_translation = -np.matmul(rotation_transposed, translation_vectors) + inverted_matrix = np.tile(np.eye(4), (len(rotation_matrices), 1, 1)) + else: + rotation_transposed = rotation_matrices.transpose(1, 2) + new_translation = -torch.bmm(rotation_transposed, translation_vectors) + inverted_matrix = torch.eye(4, 4)[None].repeat(len(rotation_matrices), 1, 1) + inverted_matrix = inverted_matrix.to(rotation_matrices.dtype).to( + rotation_matrices.device + ) + inverted_matrix[:, :3, :3] = rotation_transposed + inverted_matrix[:, :3, 3:] = new_translation + + return inverted_matrix + + +def relative_pose_transformation(trans_01, trans_02): + r""" + Function that computes the relative homogenous transformation from a + reference transformation :math:`T_1^{0} = \begin{bmatrix} R_1 & t_1 \\ + \mathbf{0} & 1 \end{bmatrix}` to destination :math:`T_2^{0} = + \begin{bmatrix} R_2 & t_2 \\ \mathbf{0} & 1 \end{bmatrix}`. + + The relative transformation is computed as follows: + + .. math:: + + T_1^{2} = (T_0^{1})^{-1} \cdot T_0^{2} + + Arguments: + trans_01 (torch.Tensor): reference transformation tensor of shape + :math:`(N, 4, 4)` or :math:`(4, 4)`. + trans_02 (torch.Tensor): destination transformation tensor of shape + :math:`(N, 4, 4)` or :math:`(4, 4)`. + + Shape: + - Output: :math:`(N, 4, 4)` or :math:`(4, 4)`. + + Returns: + torch.Tensor: the relative transformation between the transformations. + + Example:: + >>> trans_01 = torch.eye(4) # 4x4 + >>> trans_02 = torch.eye(4) # 4x4 + >>> trans_12 = relative_transformation(trans_01, trans_02) # 4x4 + """ + if not torch.is_tensor(trans_01): + raise TypeError( + "Input trans_01 type is not a torch.Tensor. Got {}".format(type(trans_01)) + ) + if not torch.is_tensor(trans_02): + raise TypeError( + "Input trans_02 type is not a torch.Tensor. Got {}".format(type(trans_02)) + ) + if trans_01.dim() not in (2, 3) and trans_01.shape[-2:] == (4, 4): + raise ValueError( + "Input must be a of the shape Nx4x4 or 4x4. Got {}".format(trans_01.shape) + ) + if trans_02.dim() not in (2, 3) and trans_02.shape[-2:] == (4, 4): + raise ValueError( + "Input must be a of the shape Nx4x4 or 4x4. Got {}".format(trans_02.shape) + ) + if not trans_01.dim() == trans_02.dim(): + raise ValueError( + "Input number of dims must match. Got {} and {}".format( + trans_01.dim(), trans_02.dim() + ) + ) + + # Convert to Nx4x4 if inputs are 4x4 + squeeze_batch_dim = False + if trans_01.dim() == 2: + trans_01 = trans_01.unsqueeze(0) + trans_02 = trans_02.unsqueeze(0) + squeeze_batch_dim = True + + # Compute inverse of trans_01 using closed form + trans_10 = closed_form_pose_inverse(trans_01) + + # Compose transformations using matrix multiplication + trans_12 = torch.matmul(trans_10, trans_02) + + # Remove batch dimension if it was added + if squeeze_batch_dim: + trans_12 = trans_12.squeeze(0) + + return trans_12 + + +def depthmap_to_pts3d(depth, pseudo_focal, pp=None, **_): + """ + Args: + - depthmap (BxHxW array): + - pseudo_focal: [B,H,W] ; [B,2,H,W] or [B,1,H,W] + Returns: + pointmap of absolute coordinates (BxHxWx3 array) + """ + + if len(depth.shape) == 4: + B, H, W, n = depth.shape + else: + B, H, W = depth.shape + n = None + + if len(pseudo_focal.shape) == 3: # [B,H,W] + pseudo_focalx = pseudo_focaly = pseudo_focal + elif len(pseudo_focal.shape) == 4: # [B,2,H,W] or [B,1,H,W] + pseudo_focalx = pseudo_focal[:, 0] + if pseudo_focal.shape[1] == 2: + pseudo_focaly = pseudo_focal[:, 1] + else: + pseudo_focaly = pseudo_focalx + else: + raise NotImplementedError("Error, unknown input focal shape format.") + + assert pseudo_focalx.shape == depth.shape[:3] + assert pseudo_focaly.shape == depth.shape[:3] + grid_x, grid_y = xy_grid(W, H, cat_dim=0, device=depth.device)[:, None] + + # set principal point + if pp is None: + grid_x = grid_x - (W - 1) / 2 + grid_y = grid_y - (H - 1) / 2 + else: + grid_x = grid_x.expand(B, -1, -1) - pp[:, 0, None, None] + grid_y = grid_y.expand(B, -1, -1) - pp[:, 1, None, None] + + if n is None: + pts3d = torch.empty((B, H, W, 3), device=depth.device) + pts3d[..., 0] = depth * grid_x / pseudo_focalx + pts3d[..., 1] = depth * grid_y / pseudo_focaly + pts3d[..., 2] = depth + else: + pts3d = torch.empty((B, H, W, 3, n), device=depth.device) + pts3d[..., 0, :] = depth * (grid_x / pseudo_focalx)[..., None] + pts3d[..., 1, :] = depth * (grid_y / pseudo_focaly)[..., None] + pts3d[..., 2, :] = depth + return pts3d + + +def depthmap_to_camera_coordinates(depthmap, camera_intrinsics, pseudo_focal=None, flow2d=None): + """ + Args: + - depthmap (HxW array): + - camera_intrinsics: a 3x3 matrix + Returns: + pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels. + """ + camera_intrinsics = np.float32(camera_intrinsics) + H, W = depthmap.shape + + # Compute 3D ray associated with each pixel + # Strong assumption: there are no skew terms + assert camera_intrinsics[0, 1] == 0.0 + assert camera_intrinsics[1, 0] == 0.0 + if pseudo_focal is None: + fu = camera_intrinsics[0, 0] + fv = camera_intrinsics[1, 1] + else: + assert pseudo_focal.shape == (H, W) + fu = fv = pseudo_focal + cu = camera_intrinsics[0, 2] + cv = camera_intrinsics[1, 2] + + u, v = np.meshgrid(np.arange(W), np.arange(H)) + + if flow2d is not None: + u = u + flow2d[..., 0] + v = v + flow2d[..., 1] + + z_cam = depthmap + x_cam = (u - cu) * z_cam / fu + y_cam = (v - cv) * z_cam / fv + X_cam = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32) + + # Mask for valid coordinates + valid_mask = depthmap > 0.0 + + return X_cam, valid_mask + + +def depthmap_to_absolute_camera_coordinates( + depthmap, camera_intrinsics, camera_pose, flow2d=None, **kw +): + """ + Args: + - depthmap (HxW array): + - camera_intrinsics: a 3x3 matrix + - camera_pose: a 4x3 or 4x4 cam2world matrix + Returns: + pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels. + """ + X_cam, valid_mask = depthmap_to_camera_coordinates(depthmap, camera_intrinsics, flow2d=flow2d) + + X_world = X_cam # default + if camera_pose is not None: + # R_cam2world = np.float32(camera_params["R_cam2world"]) + # t_cam2world = np.float32(camera_params["t_cam2world"]).squeeze() + R_cam2world = camera_pose[:3, :3] + t_cam2world = camera_pose[:3, 3] + + # Express in absolute coordinates (invalid depth values) + X_world = ( + np.einsum("ik, vuk -> vui", R_cam2world, X_cam) + t_cam2world[None, None, :] + ) + + return X_world, valid_mask + + +def get_absolute_pointmaps_and_rays_info( + depthmap, camera_intrinsics, camera_pose, **kw +): + """ + Args: + - depthmap (HxW array): + - camera_intrinsics: a 3x3 matrix + - camera_pose: a 4x3 or 4x4 cam2world matrix + Returns: + pointmap of absolute coordinates (HxWx3 array), + a mask specifying valid pixels, + ray origins of absolute coordinates (HxWx3 array), + ray directions of absolute coordinates (HxWx3 array), + depth along ray (HxWx1 array), + ray directions of camera/local coordinates (HxWx3 array), + pointmap of camera/local coordinates (HxWx3 array). + """ + camera_intrinsics = np.float32(camera_intrinsics) + H, W = depthmap.shape + + # Compute 3D ray associated with each pixel + # Strong assumption: pinhole & there are no skew terms + assert camera_intrinsics[0, 1] == 0.0 + assert camera_intrinsics[1, 0] == 0.0 + fu = camera_intrinsics[0, 0] + fv = camera_intrinsics[1, 1] + cu = camera_intrinsics[0, 2] + cv = camera_intrinsics[1, 2] + + # Get the rays on the unit plane + u, v = np.meshgrid(np.arange(W), np.arange(H)) + x_cam = (u - cu) / fu + y_cam = (v - cv) / fv + z_cam = np.ones_like(x_cam) + ray_dirs_cam_on_unit_plane = np.stack((x_cam, y_cam, z_cam), axis=-1).astype( + np.float32 + ) + + # Compute the 3d points in the local camera coordinate system + pts_cam = depthmap[..., None] * ray_dirs_cam_on_unit_plane + + # Get the depth along the ray and compute the ray directions on the unit sphere + depth_along_ray = np.linalg.norm(pts_cam, axis=-1, keepdims=True) + ray_directions_cam = ray_dirs_cam_on_unit_plane / np.linalg.norm( + ray_dirs_cam_on_unit_plane, axis=-1, keepdims=True + ) + + # Mask for valid coordinates + valid_mask = depthmap > 0.0 + + # Get the ray origins in absolute coordinates and the ray directions in absolute coordinates + ray_origins_world = np.zeros_like(ray_directions_cam) + ray_directions_world = ray_directions_cam + pts_world = pts_cam + if camera_pose is not None: + R_cam2world = camera_pose[:3, :3] + t_cam2world = camera_pose[:3, 3] + + # Express in absolute coordinates + ray_origins_world = ray_origins_world + t_cam2world[None, None, :] + ray_directions_world = np.einsum( + "ik, vuk -> vui", R_cam2world, ray_directions_cam + ) + pts_world = ray_origins_world + ray_directions_world * depth_along_ray + + return ( + pts_world, + valid_mask, + ray_origins_world, + ray_directions_world, + depth_along_ray, + ray_directions_cam, + pts_cam, + ) + + +def adjust_camera_params_for_rotation(camera_params, original_size, k): + """ + Adjust camera parameters for rotation. + + Args: + camera_params: Camera parameters [fx, fy, cx, cy, ...] + original_size: Original image size as (width, height) + k: Number of 90-degree rotations counter-clockwise (k=3 means 90 degrees clockwise) + + Returns: + Adjusted camera parameters + """ + fx, fy, cx, cy = camera_params[:4] + width, height = original_size + + if k % 4 == 1: # 90 degrees counter-clockwise + new_fx, new_fy = fy, fx + new_cx, new_cy = height - cy, cx + elif k % 4 == 2: # 180 degrees + new_fx, new_fy = fx, fy + new_cx, new_cy = width - cx, height - cy + elif k % 4 == 3: # 90 degrees clockwise (270 counter-clockwise) + new_fx, new_fy = fy, fx + new_cx, new_cy = cy, width - cx + else: # No rotation + return camera_params + + adjusted_params = [new_fx, new_fy, new_cx, new_cy] + if len(camera_params) > 4: + adjusted_params.extend(camera_params[4:]) + + return adjusted_params + + +def adjust_pose_for_rotation(pose, k): + """ + Adjust camera pose for rotation. + + Args: + pose: 4x4 camera pose matrix (camera-to-world, OpenCV convention - X right, Y down, Z forward) + k: Number of 90-degree rotations counter-clockwise (k=3 means 90 degrees clockwise) + + Returns: + Adjusted 4x4 camera pose matrix + """ + # Create rotation matrices for different rotations + if k % 4 == 1: # 90 degrees counter-clockwise + rot_transform = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]]) + elif k % 4 == 2: # 180 degrees + rot_transform = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]) + elif k % 4 == 3: # 90 degrees clockwise (270 counter-clockwise) + rot_transform = np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]]) + else: # No rotation + return pose + + # Apply the transformation to the pose + adjusted_pose = pose + adjusted_pose[:3, :3] = adjusted_pose[:3, :3] @ rot_transform.T + + return adjusted_pose + + +def crop_to_aspect_ratio(image, depth, camera_params, target_ratio=1.5): + """ + Crop image and depth to the largest possible target aspect ratio while + keeping the left side if aspect ratio is wider and the bottom of image if the aspect ratio is taller. + + Args: + image: PIL image + depth: Depth map as numpy array + camera_params: Camera parameters [fx, fy, cx, cy, ...] + target_ratio: Target width/height ratio + + Returns: + Cropped image, cropped depth, adjusted camera parameters + """ + width, height = image.size + fx, fy, cx, cy = camera_params[:4] + current_ratio = width / height + + if abs(current_ratio - target_ratio) < 1e-6: + # Already at target ratio + return image, depth, camera_params + + if current_ratio > target_ratio: + # Image is wider than target ratio, crop width + new_width = int(height * target_ratio) + left = 0 + right = new_width + + # Crop image + cropped_image = image.crop((left, 0, right, height)) + + # Crop depth + if len(depth.shape) == 3: + cropped_depth = depth[:, left:right, :] + else: + cropped_depth = depth[:, left:right] + + # Adjust camera parameters + new_cx = cx - left + adjusted_params = [fx, fy, new_cx, cy] + list(camera_params[4:]) + + else: + # Image is taller than target ratio, crop height + new_height = int(width / target_ratio) + top = max(0, height - new_height) + bottom = height + + # Crop image + cropped_image = image.crop((0, top, width, bottom)) + + # Crop depth + if len(depth.shape) == 3: + cropped_depth = depth[top:bottom, :, :] + else: + cropped_depth = depth[top:bottom, :] + + # Adjust camera parameters + new_cy = cy - top + adjusted_params = [fx, fy, cx, new_cy] + list(camera_params[4:]) + + return cropped_image, cropped_depth, adjusted_params + + +def colmap_to_opencv_intrinsics(K): + """ + Modify camera intrinsics to follow a different convention. + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] -= 0.5 + K[1, 2] -= 0.5 + + return K + + +def opencv_to_colmap_intrinsics(K): + """ + Modify camera intrinsics to follow a different convention. + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] += 0.5 + K[1, 2] += 0.5 + + return K + + +def normalize_depth_using_non_zero_pixels(depth, return_norm_factor=False): + """ + Normalize the depth by the average depth of non-zero depth pixels. + + Args: + depth (torch.Tensor): Depth tensor of size [B, H, W, 1]. + Returns: + normalized_depth (torch.Tensor): Normalized depth tensor. + norm_factor (torch.Tensor): Norm factor tensor of size B. + """ + assert depth.ndim == 4 and depth.shape[3] == 1 + # Calculate the sum and count of non-zero depth pixels for each batch + valid_depth_mask = depth > 0 + valid_sum = torch.sum(depth * valid_depth_mask, dim=(1, 2, 3)) + valid_count = torch.sum(valid_depth_mask, dim=(1, 2, 3)) + + # Calculate the norm factor + norm_factor = valid_sum / (valid_count + 1e-8) + while norm_factor.ndim < depth.ndim: + norm_factor.unsqueeze_(-1) + + # Normalize the depth by the norm factor + norm_factor = norm_factor.clip(min=1e-8) + normalized_depth = depth / norm_factor + + # Create the output tuple + output = ( + (normalized_depth, norm_factor.squeeze(-1).squeeze(-1).squeeze(-1)) + if return_norm_factor + else normalized_depth + ) + + return output + + +def normalize_pose_translations(pose_translations, return_norm_factor=False): + """ + Normalize the pose translations by the average norm of the non-zero pose translations. + + Args: + pose_translations (torch.Tensor): Pose translations tensor of size [B, V, 3]. B is the batch size, V is the number of views. + Returns: + normalized_pose_translations (torch.Tensor): Normalized pose translations tensor of size [B, V, 3]. + norm_factor (torch.Tensor): Norm factor tensor of size B. + """ + assert pose_translations.ndim == 3 and pose_translations.shape[2] == 3 + # Compute distance of all pose translations to origin + pose_translations_dis = pose_translations.norm(dim=-1) # [B, V] + non_zero_pose_translations_dis = pose_translations_dis > 0 # [B, V] + + # Calculate the average norm of the translations across all views (considering only views with non-zero translations) + sum_of_all_views_pose_translations = pose_translations_dis.sum(dim=1) # [B] + count_of_all_views_with_non_zero_pose_translations = ( + non_zero_pose_translations_dis.sum(dim=1) + ) # [B] + norm_factor = sum_of_all_views_pose_translations / ( + count_of_all_views_with_non_zero_pose_translations + 1e-8 + ) # [B] + + # Normalize the pose translations by the norm factor + norm_factor = norm_factor.clip(min=1e-8) + normalized_pose_translations = pose_translations / norm_factor.unsqueeze( + -1 + ).unsqueeze(-1) + + # Create the output tuple + output = ( + (normalized_pose_translations, norm_factor) + if return_norm_factor + else normalized_pose_translations + ) + + return output + + +def normalize_multiple_pointclouds( + pts_list, valid_masks=None, norm_mode="avg_dis", ret_factor=False +): + """ + Normalize multiple point clouds using a joint normalization strategy. + + Args: + pts_list: List of point clouds, each with shape (..., H, W, 3) or (B, H, W, 3) + valid_masks: Optional list of masks indicating valid points in each point cloud + norm_mode: String in format "{norm}_{dis}" where: + - norm: Normalization strategy (currently only "avg" is supported) + - dis: Distance transformation ("dis" for raw distance, "log1p" for log(1+distance), + "warp-log1p" to warp points using log distance) + ret_factor: If True, return the normalization factor as the last element in the result list + + Returns: + List of normalized point clouds with the same shapes as inputs. + If ret_factor is True, the last element is the normalization factor. + """ + assert all(pts.ndim >= 3 and pts.shape[-1] == 3 for pts in pts_list) + if valid_masks is not None: + assert len(pts_list) == len(valid_masks) + + norm_mode, dis_mode = norm_mode.split("_") + + # Gather all points together (joint normalization) + nan_pts_list = [ + invalid_to_zeros(pts, valid_masks[i], ndim=3) + if valid_masks + else invalid_to_zeros(pts, None, ndim=3) + for i, pts in enumerate(pts_list) + ] + all_pts = torch.cat([nan_pts for nan_pts, _ in nan_pts_list], dim=1) + nnz_list = [nnz for _, nnz in nan_pts_list] + + # Compute distance to origin + all_dis = all_pts.norm(dim=-1) + if dis_mode == "dis": + pass # do nothing + elif dis_mode == "log1p": + all_dis = torch.log1p(all_dis) + elif dis_mode == "warp-log1p": + # Warp input points before normalizing them + log_dis = torch.log1p(all_dis) + warp_factor = log_dis / all_dis.clip(min=1e-8) + for i, pts in enumerate(pts_list): + H, W = pts.shape[1:-1] + pts_list[i] = pts * warp_factor[:, i * (H * W) : (i + 1) * (H * W)].view( + -1, H, W, 1 + ) + all_dis = log_dis + else: + raise ValueError(f"bad {dis_mode=}") + + # Compute normalization factor + norm_factor = all_dis.sum(dim=1) / (sum(nnz_list) + 1e-8) + norm_factor = norm_factor.clip(min=1e-8) + while norm_factor.ndim < pts_list[0].ndim: + norm_factor.unsqueeze_(-1) + + # Normalize points + res = [pts / norm_factor for pts in pts_list] + if ret_factor: + res.append(norm_factor) + + return res + + +def apply_log_to_norm(input_data): + """ + Normalize the input data and apply a logarithmic transformation based on the normalization factor. + + Args: + input_data (torch.Tensor): The input tensor to be normalized and transformed. + + Returns: + torch.Tensor: The transformed tensor after normalization and logarithmic scaling. + """ + org_d = input_data.norm(dim=-1, keepdim=True) + input_data = input_data / org_d.clip(min=1e-8) + input_data = input_data * torch.log1p(org_d) + return input_data + + +def angle_diff_vec3(v1, v2, eps=1e-12): + """ + Compute angle difference between 3D vectors. + + Args: + v1: torch.Tensor of shape (..., 3) + v2: torch.Tensor of shape (..., 3) + eps: Small epsilon value for numerical stability + + Returns: + torch.Tensor: Angle differences in radians + """ + cross_norm = torch.cross(v1, v2, dim=-1).norm(dim=-1) + eps + dot_prod = (v1 * v2).sum(dim=-1) + return torch.atan2(cross_norm, dot_prod) + + +def angle_diff_vec3_numpy(v1: np.ndarray, v2: np.ndarray, eps: float = 1e-12): + """ + Compute angle difference between 3D vectors using NumPy. + + Args: + v1 (np.ndarray): First vector of shape (..., 3) + v2 (np.ndarray): Second vector of shape (..., 3) + eps (float, optional): Small epsilon value for numerical stability. Defaults to 1e-12. + + Returns: + np.ndarray: Angle differences in radians + """ + return np.arctan2( + np.linalg.norm(np.cross(v1, v2, axis=-1), axis=-1) + eps, (v1 * v2).sum(axis=-1) + ) + + +@no_warnings(category=RuntimeWarning) +def points_to_normals( + point: np.ndarray, mask: np.ndarray = None, edge_threshold: float = None +) -> np.ndarray: + """ + Calculate normal map from point map. Value range is [-1, 1]. + + Args: + point (np.ndarray): shape (height, width, 3), point map + mask (optional, np.ndarray): shape (height, width), dtype=bool. Mask of valid depth pixels. Defaults to None. + edge_threshold (optional, float): threshold for the angle (in degrees) between the normal and the view direction. Defaults to None. + + Returns: + normal (np.ndarray): shape (height, width, 3), normal map. + """ + height, width = point.shape[-3:-1] + has_mask = mask is not None + + if mask is None: + mask = np.ones_like(point[..., 0], dtype=bool) + mask_pad = np.zeros((height + 2, width + 2), dtype=bool) + mask_pad[1:-1, 1:-1] = mask + mask = mask_pad + + pts = np.zeros((height + 2, width + 2, 3), dtype=point.dtype) + pts[1:-1, 1:-1, :] = point + up = pts[:-2, 1:-1, :] - pts[1:-1, 1:-1, :] + left = pts[1:-1, :-2, :] - pts[1:-1, 1:-1, :] + down = pts[2:, 1:-1, :] - pts[1:-1, 1:-1, :] + right = pts[1:-1, 2:, :] - pts[1:-1, 1:-1, :] + normal = np.stack( + [ + np.cross(up, left, axis=-1), + np.cross(left, down, axis=-1), + np.cross(down, right, axis=-1), + np.cross(right, up, axis=-1), + ] + ) + normal = normal / (np.linalg.norm(normal, axis=-1, keepdims=True) + 1e-12) + + valid = ( + np.stack( + [ + mask[:-2, 1:-1] & mask[1:-1, :-2], + mask[1:-1, :-2] & mask[2:, 1:-1], + mask[2:, 1:-1] & mask[1:-1, 2:], + mask[1:-1, 2:] & mask[:-2, 1:-1], + ] + ) + & mask[None, 1:-1, 1:-1] + ) + if edge_threshold is not None: + view_angle = angle_diff_vec3_numpy(pts[None, 1:-1, 1:-1, :], normal) + view_angle = np.minimum(view_angle, np.pi - view_angle) + valid = valid & (view_angle < np.deg2rad(edge_threshold)) + + normal = (normal * valid[..., None]).sum(axis=0) + normal = normal / (np.linalg.norm(normal, axis=-1, keepdims=True) + 1e-12) + + if has_mask: + normal_mask = valid.any(axis=0) + normal = np.where(normal_mask[..., None], normal, 0) + return normal, normal_mask + else: + return normal + + +def sliding_window_1d(x: np.ndarray, window_size: int, stride: int, axis: int = -1): + """ + Create a sliding window view of the input array along a specified axis. + + This function creates a memory-efficient view of the input array with sliding windows + of the specified size and stride. The window dimension is appended to the end of the + output array's shape. This is useful for operations like convolution, pooling, or + any analysis that requires examining local neighborhoods in the data. + + Args: + x (np.ndarray): Input array with shape (..., axis_size, ...) + window_size (int): Size of the sliding window + stride (int): Stride of the sliding window (step size between consecutive windows) + axis (int, optional): Axis to perform sliding window over. Defaults to -1 (last axis) + + Returns: + np.ndarray: View of the input array with shape (..., n_windows, ..., window_size), + where n_windows = (axis_size - window_size + 1) // stride + + Raises: + AssertionError: If window_size is larger than the size of the specified axis + + Example: + >>> x = np.array([1, 2, 3, 4, 5, 6]) + >>> sliding_window_1d(x, window_size=3, stride=2) + array([[1, 2, 3], + [3, 4, 5]]) + """ + assert x.shape[axis] >= window_size, ( + f"kernel_size ({window_size}) is larger than axis_size ({x.shape[axis]})" + ) + axis = axis % x.ndim + shape = ( + *x.shape[:axis], + (x.shape[axis] - window_size + 1) // stride, + *x.shape[axis + 1 :], + window_size, + ) + strides = ( + *x.strides[:axis], + stride * x.strides[axis], + *x.strides[axis + 1 :], + x.strides[axis], + ) + x_sliding = np.lib.stride_tricks.as_strided(x, shape=shape, strides=strides) + return x_sliding + + +def sliding_window_nd( + x: np.ndarray, + window_size: Tuple[int, ...], + stride: Tuple[int, ...], + axis: Tuple[int, ...], +) -> np.ndarray: + """ + Create sliding windows along multiple dimensions of the input array. + + This function applies sliding_window_1d sequentially along multiple axes to create + N-dimensional sliding windows. This is useful for operations that need to examine + local neighborhoods in multiple dimensions simultaneously. + + Args: + x (np.ndarray): Input array + window_size (Tuple[int, ...]): Size of the sliding window for each axis + stride (Tuple[int, ...]): Stride of the sliding window for each axis + axis (Tuple[int, ...]): Axes to perform sliding window over + + Returns: + np.ndarray: Array with sliding windows along the specified dimensions. + The window dimensions are appended to the end of the shape. + + Note: + The length of window_size, stride, and axis tuples must be equal. + + Example: + >>> x = np.random.rand(10, 10) + >>> windows = sliding_window_nd(x, window_size=(3, 3), stride=(2, 2), axis=(-2, -1)) + >>> # Creates 3x3 sliding windows with stride 2 in both dimensions + """ + axis = [axis[i] % x.ndim for i in range(len(axis))] + for i in range(len(axis)): + x = sliding_window_1d(x, window_size[i], stride[i], axis[i]) + return x + + +def sliding_window_2d( + x: np.ndarray, + window_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]], + axis: Tuple[int, int] = (-2, -1), +) -> np.ndarray: + """ + Create 2D sliding windows over the input array. + + Convenience function for creating 2D sliding windows, commonly used for image + processing operations like convolution, pooling, or patch extraction. + + Args: + x (np.ndarray): Input array + window_size (Union[int, Tuple[int, int]]): Size of the 2D sliding window. + If int, same size is used for both dimensions. + stride (Union[int, Tuple[int, int]]): Stride of the 2D sliding window. + If int, same stride is used for both dimensions. + axis (Tuple[int, int], optional): Two axes to perform sliding window over. + Defaults to (-2, -1) (last two dimensions). + + Returns: + np.ndarray: Array with 2D sliding windows. The window dimensions (height, width) + are appended to the end of the shape. + + Example: + >>> image = np.random.rand(100, 100) + >>> patches = sliding_window_2d(image, window_size=8, stride=4) + >>> # Creates 8x8 patches with stride 4 from the image + """ + if isinstance(window_size, int): + window_size = (window_size, window_size) + if isinstance(stride, int): + stride = (stride, stride) + return sliding_window_nd(x, window_size, stride, axis) + + +def max_pool_1d( + x: np.ndarray, kernel_size: int, stride: int, padding: int = 0, axis: int = -1 +): + """ + Perform 1D max pooling on the input array. + + Max pooling reduces the dimensionality of the input by taking the maximum value + within each sliding window. This is commonly used in neural networks and signal + processing for downsampling and feature extraction. + + Args: + x (np.ndarray): Input array + kernel_size (int): Size of the pooling kernel + stride (int): Stride of the pooling operation + padding (int, optional): Amount of padding to add on both sides. Defaults to 0. + axis (int, optional): Axis to perform max pooling over. Defaults to -1. + + Returns: + np.ndarray: Max pooled array with reduced size along the specified axis + + Note: + - For floating point arrays, padding is done with np.nan values + - For integer arrays, padding is done with the minimum value of the dtype + - np.nanmax is used to handle NaN values in the computation + + Example: + >>> x = np.array([1, 3, 2, 4, 5, 1, 2]) + >>> max_pool_1d(x, kernel_size=3, stride=2) + array([3, 5, 2]) + """ + axis = axis % x.ndim + if padding > 0: + fill_value = np.nan if x.dtype.kind == "f" else np.iinfo(x.dtype).min + padding_arr = np.full( + (*x.shape[:axis], padding, *x.shape[axis + 1 :]), + fill_value=fill_value, + dtype=x.dtype, + ) + x = np.concatenate([padding_arr, x, padding_arr], axis=axis) + a_sliding = sliding_window_1d(x, kernel_size, stride, axis) + max_pool = np.nanmax(a_sliding, axis=-1) + return max_pool + + +def max_pool_nd( + x: np.ndarray, + kernel_size: Tuple[int, ...], + stride: Tuple[int, ...], + padding: Tuple[int, ...], + axis: Tuple[int, ...], +) -> np.ndarray: + """ + Perform N-dimensional max pooling on the input array. + + This function applies max_pool_1d sequentially along multiple axes to perform + multi-dimensional max pooling. This is useful for downsampling multi-dimensional + data while preserving the most important features. + + Args: + x (np.ndarray): Input array + kernel_size (Tuple[int, ...]): Size of the pooling kernel for each axis + stride (Tuple[int, ...]): Stride of the pooling operation for each axis + padding (Tuple[int, ...]): Amount of padding for each axis + axis (Tuple[int, ...]): Axes to perform max pooling over + + Returns: + np.ndarray: Max pooled array with reduced size along the specified axes + + Note: + The length of kernel_size, stride, padding, and axis tuples must be equal. + Max pooling is applied sequentially along each axis in the order specified. + + Example: + >>> x = np.random.rand(10, 10, 10) + >>> pooled = max_pool_nd(x, kernel_size=(2, 2, 2), stride=(2, 2, 2), + ... padding=(0, 0, 0), axis=(-3, -2, -1)) + >>> # Reduces each dimension by half with 2x2x2 max pooling + """ + for i in range(len(axis)): + x = max_pool_1d(x, kernel_size[i], stride[i], padding[i], axis[i]) + return x + + +def max_pool_2d( + x: np.ndarray, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]], + padding: Union[int, Tuple[int, int]], + axis: Tuple[int, int] = (-2, -1), +): + """ + Perform 2D max pooling on the input array. + + Convenience function for 2D max pooling, commonly used in computer vision + and image processing for downsampling images while preserving important features. + + Args: + x (np.ndarray): Input array + kernel_size (Union[int, Tuple[int, int]]): Size of the 2D pooling kernel. + If int, same size is used for both dimensions. + stride (Union[int, Tuple[int, int]]): Stride of the 2D pooling operation. + If int, same stride is used for both dimensions. + padding (Union[int, Tuple[int, int]]): Amount of padding for both dimensions. + If int, same padding is used for both dimensions. + axis (Tuple[int, int], optional): Two axes to perform max pooling over. + Defaults to (-2, -1) (last two dimensions). + + Returns: + np.ndarray: 2D max pooled array with reduced size along the specified axes + + Example: + >>> image = np.random.rand(64, 64) + >>> pooled = max_pool_2d(image, kernel_size=2, stride=2, padding=0) + >>> # Reduces image size from 64x64 to 32x32 with 2x2 max pooling + """ + if isinstance(kernel_size, Number): + kernel_size = (kernel_size, kernel_size) + if isinstance(stride, Number): + stride = (stride, stride) + if isinstance(padding, Number): + padding = (padding, padding) + axis = tuple(axis) + return max_pool_nd(x, kernel_size, stride, padding, axis) + + +@no_warnings(category=RuntimeWarning) +def depth_edge( + depth: np.ndarray, + atol: float = None, + rtol: float = None, + kernel_size: int = 3, + mask: np.ndarray = None, +) -> np.ndarray: + """ + Compute the edge mask from depth map. The edge is defined as the pixels whose neighbors have large difference in depth. + + Args: + depth (np.ndarray): shape (..., height, width), linear depth map + atol (float): absolute tolerance + rtol (float): relative tolerance + + Returns: + edge (np.ndarray): shape (..., height, width) of dtype torch.bool + """ + if mask is None: + diff = max_pool_2d( + depth, kernel_size, stride=1, padding=kernel_size // 2 + ) + max_pool_2d(-depth, kernel_size, stride=1, padding=kernel_size // 2) + else: + diff = max_pool_2d( + np.where(mask, depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + max_pool_2d( + np.where(mask, -depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + + edge = np.zeros_like(depth, dtype=bool) + if atol is not None: + edge |= diff > atol + + if rtol is not None: + edge |= diff / depth > rtol + return edge + + +def depth_aliasing( + depth: np.ndarray, + atol: float = None, + rtol: float = None, + kernel_size: int = 3, + mask: np.ndarray = None, +) -> np.ndarray: + """ + Compute the map that indicates the aliasing of x depth map. The aliasing is defined as the pixels which neither close to the maximum nor the minimum of its neighbors. + Args: + depth (np.ndarray): shape (..., height, width), linear depth map + atol (float): absolute tolerance + rtol (float): relative tolerance + + Returns: + edge (np.ndarray): shape (..., height, width) of dtype torch.bool + """ + if mask is None: + diff_max = ( + max_pool_2d(depth, kernel_size, stride=1, padding=kernel_size // 2) - depth + ) + diff_min = ( + max_pool_2d(-depth, kernel_size, stride=1, padding=kernel_size // 2) + depth + ) + else: + diff_max = ( + max_pool_2d( + np.where(mask, depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + - depth + ) + diff_min = ( + max_pool_2d( + np.where(mask, -depth, -np.inf), + kernel_size, + stride=1, + padding=kernel_size // 2, + ) + + depth + ) + diff = np.minimum(diff_max, diff_min) + + edge = np.zeros_like(depth, dtype=bool) + if atol is not None: + edge |= diff > atol + if rtol is not None: + edge |= diff / depth > rtol + return edge + + +@no_warnings(category=RuntimeWarning) +def normals_edge( + normals: np.ndarray, tol: float, kernel_size: int = 3, mask: np.ndarray = None +) -> np.ndarray: + """ + Compute the edge mask from normal map. + + Args: + normal (np.ndarray): shape (..., height, width, 3), normal map + tol (float): tolerance in degrees + + Returns: + edge (np.ndarray): shape (..., height, width) of dtype torch.bool + """ + assert normals.ndim >= 3 and normals.shape[-1] == 3, ( + "normal should be of shape (..., height, width, 3)" + ) + normals = normals / (np.linalg.norm(normals, axis=-1, keepdims=True) + 1e-12) + + padding = kernel_size // 2 + normals_window = sliding_window_2d( + np.pad( + normals, + ( + *([(0, 0)] * (normals.ndim - 3)), + (padding, padding), + (padding, padding), + (0, 0), + ), + mode="edge", + ), + window_size=kernel_size, + stride=1, + axis=(-3, -2), + ) + if mask is None: + angle_diff = np.arccos( + (normals[..., None, None] * normals_window).sum(axis=-3) + ).max(axis=(-2, -1)) + else: + mask_window = sliding_window_2d( + np.pad( + mask, + (*([(0, 0)] * (mask.ndim - 3)), (padding, padding), (padding, padding)), + mode="edge", + ), + window_size=kernel_size, + stride=1, + axis=(-3, -2), + ) + angle_diff = np.where( + mask_window, + np.arccos((normals[..., None, None] * normals_window).sum(axis=-3)), + 0, + ).max(axis=(-2, -1)) + + angle_diff = max_pool_2d( + angle_diff, kernel_size, stride=1, padding=kernel_size // 2 + ) + edge = angle_diff > np.deg2rad(tol) + return edge diff --git a/mapanything/utils/hf_utils/__init__.py b/mapanything/utils/hf_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/mapanything/utils/hf_utils/css_and_html.py b/mapanything/utils/hf_utils/css_and_html.py new file mode 100644 index 0000000000000000000000000000000000000000..e5ae13c7d40419ce7b5410fd85d6c1f6719dbc61 --- /dev/null +++ b/mapanything/utils/hf_utils/css_and_html.py @@ -0,0 +1,211 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +""" +CSS and HTML content for the MapAnything Gradio application. +This module contains all the CSS styles and HTML content blocks +used in the Gradio interface. +""" + +# CSS Styles for the Gradio interface +GRADIO_CSS = """ +.custom-log * { + font-style: italic; + font-size: 22px !important; + background-image: linear-gradient(120deg, #ffb366 0%, #ffa366 60%, #ff9966 100%); + -webkit-background-clip: text; + background-clip: text; + font-weight: bold !important; + color: transparent !important; + text-align: center !important; +} + +.example-log * { + font-style: italic; + font-size: 16px !important; + background-image: linear-gradient(120deg, #ffb366 0%, #ffa366 60%, #ff9966 100%); + -webkit-background-clip: text; + background-clip: text; + color: transparent !important; +} + +#my_radio .wrap { + display: flex; + flex-wrap: nowrap; + justify-content: center; + align-items: center; +} + +#my_radio .wrap label { + display: flex; + width: 50%; + justify-content: center; + align-items: center; + margin: 0; + padding: 10px 0; + box-sizing: border-box; +} + +/* Align navigation buttons with dropdown bottom */ +.navigation-row { + display: flex !important; + align-items: flex-end !important; + gap: 8px !important; +} + +.navigation-row > div:nth-child(1), +.navigation-row > div:nth-child(3) { + align-self: flex-end !important; +} + +.navigation-row > div:nth-child(2) { + flex: 1 !important; +} + +/* Make thumbnails clickable with pointer cursor */ +.clickable-thumbnail img { + cursor: pointer !important; +} + +.clickable-thumbnail:hover img { + cursor: pointer !important; + opacity: 0.8; + transition: opacity 0.3s ease; +} + +/* Make thumbnail containers narrower horizontally */ +.clickable-thumbnail { + padding: 5px 2px !important; + margin: 0 2px !important; +} + +.clickable-thumbnail .image-container { + margin: 0 !important; + padding: 0 !important; +} + +.scene-info { + text-align: center !important; + padding: 5px 2px !important; + margin: 0 !important; +} +""" + + +def get_header_html(logo_base64=None): + """ + Generate the main header HTML with logo and title. + + Args: + logo_base64 (str, optional): Base64 encoded logo image + + Returns: + str: HTML string for the header + """ + logo_style = "display: none;" if not logo_base64 else "" + logo_src = logo_base64 or "" + + return f""" +
+ WAI Logo +

MapAnything: Metric 3D Scene Reconstruction

+
+

+ 🌟 GitHub Repository | + 🚀 Project Page +

+ """ + + +def get_description_html(): + """ + Generate the main description and getting started HTML. + + Returns: + str: HTML string for the description + """ + return """ +
+

Upload a video or a set of images to create a 3D reconstruction of a scene or object. MapAnything takes these images and generates 3D point clouds directly from multi-view images.

+

This demo demonstrates the image input configuration only. For trying out additional inputs like cameras and depth, please check out the code in our Github repo.

+ +

Getting Started:

+
    +
  1. Upload Your Data: Use the "Upload Video" or "Upload Images" buttons on the left to provide your input. Videos will be automatically split into individual frames (one frame per second).
    2. +
  2. Preview: Your uploaded images will appear in the gallery on the left.
    3. +
  3. Reconstruct: Click the "Reconstruct" button to start the 3D reconstruction process.
    4. +
  4. Visualize: The 3D reconstruction will appear in the viewer on the right. You can rotate, pan, and zoom to explore the model, and download the GLB file. Note the visualization of 3D points may be slow for a large number of input images.
    5. +
  5. + Adjust Reconstruction & Visualization (Optional): + You can fine-tune the visualization using the options below the viewer +
    + (click to expand): +
      +
    • Show Camera: Toggle the display of estimated camera positions.
      • +
    • Show Mesh: Use meshes for the prediction visualization.
      • +
    • Show Points from Frame: Select specific frames to display in the viewer.
      • +
    • Filter Black Background: Remove black background pixels.
      • +
    • Filter White Background: Remove white background pixels.
      • +
    +
    +
    6. +
+

Please note: Depending on the amount of images, our model usually only needs about 1 second to reconstruct a scene. However, downloading model weights and visualizing 3D points may take tens of seconds. Please be patient or, for faster visualization, use a local machine to run our demo from our GitHub repository.

+
+ """ + + +def get_acknowledgements_html(): + """ + Generate the acknowledgements section HTML. + + Returns: + str: HTML string for the acknowledgements + """ + return """ +
+
+

Acknowledgements

+

This site builds upon code from:

+ +

We extend our gratitude to these projects for their valuable contributions to the research community.

+
+ """ + + +def get_gradio_theme(): + """ + Get the configured Gradio theme. + + Returns: + gr.themes.Base: Configured Gradio theme + """ + import gradio as gr + + return gr.themes.Base( + primary_hue=gr.themes.Color( + c100="#ffedd5", + c200="#ffddb3", + c300="rgba(242.78125, 182.89427563548466, 120.32579495614034, 1)", + c400="#fb923c", + c50="#fff7ed", + c500="#f97316", + c600="#ea580c", + c700="#c2410c", + c800="#9a3412", + c900="#7c2d12", + c950="#6c2e12", + ), + secondary_hue="amber", + ) + + +# Measure tab instructions HTML +MEASURE_INSTRUCTIONS_HTML = """ +### Click on the image to measure the distance between two points. +""" diff --git a/mapanything/utils/hf_utils/hf_helpers.py b/mapanything/utils/hf_utils/hf_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..8a7dbe668c99810f0d534217df79a3c3167013e5 --- /dev/null +++ b/mapanything/utils/hf_utils/hf_helpers.py @@ -0,0 +1,186 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +""" +Helper functions for HuggingFace integration and model initialization. +""" + +import json +import os + + +def load_hf_token(): + """Load HuggingFace access token from local file""" + token_file_paths = [ + "/home/aknapitsch/hf_token.txt", + ] + + for token_path in token_file_paths: + if os.path.exists(token_path): + try: + with open(token_path, "r") as f: + token = f.read().strip() + print(f"Loaded HuggingFace token from: {token_path}") + return token + except Exception as e: + print(f"Error reading token from {token_path}: {e}") + continue + else: + print(token_path, "token_path doesnt exist") + + # Also try environment variable + # see https://huggingface.co/docs/hub/spaces-overview#managing-secrets on options + token = ( + os.getenv("HF_TOKEN") + or os.getenv("HUGGING_FACE_HUB_TOKEN") + or os.getenv("HUGGING_FACE_MODEL_TOKEN") + ) + if token: + print("Loaded HuggingFace token from environment variable") + return token + + print( + "Warning: No HuggingFace token found. Model loading may fail for private repositories." + ) + return None + + +def init_hydra_config(config_path, overrides=None): + """Initialize Hydra config""" + import hydra + + config_dir = os.path.dirname(config_path) + config_name = os.path.basename(config_path).split(".")[0] + relative_path = os.path.relpath(config_dir, os.path.dirname(__file__)) + hydra.core.global_hydra.GlobalHydra.instance().clear() + hydra.initialize(version_base=None, config_path=relative_path) + if overrides is not None: + cfg = hydra.compose(config_name=config_name, overrides=overrides) + else: + cfg = hydra.compose(config_name=config_name) + return cfg + + +def initialize_mapanything_model(high_level_config, device): + """ + Initialize MapAnything model with three-tier fallback approach: + 1. Try HuggingFace from_pretrained() + 2. Download HF config + use local model factory + load HF weights + 3. Pure local configuration fallback + + Args: + high_level_config (dict): Configuration dictionary containing model settings + device (torch.device): Device to load the model on + + Returns: + torch.nn.Module: Initialized MapAnything model + """ + import torch + from huggingface_hub import hf_hub_download + + from mapanything.models import init_model, MapAnything + + print("Initializing MapAnything model...") + + # Initialize Hydra config and create model from configuration + cfg = init_hydra_config( + high_level_config["path"], overrides=high_level_config["config_overrides"] + ) + + # Try using from_pretrained first + try: + print("Loading MapAnything model from_pretrained...") + model = MapAnything.from_pretrained(high_level_config["hf_model_name"]).to( + device + ) + print("Loading MapAnything model from_pretrained succeeded...") + return model + except Exception as e: + print(f"from_pretrained failed: {e}") + print("Falling back to local configuration approach using hf_hub_download...") + + # Create model from local configuration instead of using from_pretrained + # Try to download and use the config from HuggingFace Hub + try: + print("Downloading model configuration from HuggingFace Hub...") + config_path = hf_hub_download( + repo_id=high_level_config["hf_model_name"], + filename=high_level_config["config_name"], + token=load_hf_token(), + ) + + # Load the config from the downloaded file + with open(config_path, "r") as f: + downloaded_config = json.load(f) + + print("Using downloaded configuration for model initialization") + model = init_model( + model_str=downloaded_config.get( + "model_str", high_level_config["model_str"] + ), + model_config=downloaded_config.get( + "model_config", cfg.model.model_config + ), + torch_hub_force_reload=high_level_config.get( + "torch_hub_force_reload", False + ), + ) + except Exception as config_e: + print(f"Failed to download/use HuggingFace config: {config_e}") + print("Falling back to local configuration...") + # Fall back to local configuration as before + model = init_model( + model_str=cfg.model.model_str, + model_config=cfg.model.model_config, + torch_hub_force_reload=high_level_config.get( + "torch_hub_force_reload", False + ), + ) + + # Load the pretrained weights from HuggingFace Hub + try: + # First, let's see what files are available in the repository + try: + checkpoint_filename = high_level_config["checkpoint_name"] + # Download the model weights + checkpoint_path = hf_hub_download( + repo_id=high_level_config["hf_model_name"], + filename=checkpoint_filename, + token=load_hf_token(), + ) + + # Load the weights + print("start loading checkpoint") + if checkpoint_filename.endswith(".safetensors"): + from safetensors.torch import load_file + + checkpoint = load_file(checkpoint_path) + else: + checkpoint = torch.load( + checkpoint_path, map_location="cpu", weights_only=False + ) + + print("start loading state_dict") + if "model" in checkpoint: + model.load_state_dict(checkpoint["model"], strict=False) + elif "state_dict" in checkpoint: + model.load_state_dict(checkpoint["state_dict"], strict=False) + else: + model.load_state_dict(checkpoint, strict=False) + + print( + f"Successfully loaded pretrained weights from HuggingFace Hub ({checkpoint_filename})" + ) + + except Exception as inner_e: + print(f"Error listing repository files or loading weights: {inner_e}") + raise inner_e + + except Exception as e: + print(f"Warning: Could not load pretrained weights: {e}") + print("Proceeding with randomly initialized model...") + + model = model.to(device) + return model diff --git a/mapanything/utils/hf_utils/viz.py b/mapanything/utils/hf_utils/viz.py new file mode 100644 index 0000000000000000000000000000000000000000..f6e4c08b442758800718ff2494c5379ee7911bc9 --- /dev/null +++ b/mapanything/utils/hf_utils/viz.py @@ -0,0 +1,681 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# This source code is licensed under the Apache License, Version 2.0 +# found in the LICENSE file in the root directory of this source tree. + +""" +Utility functions for Gradio demo visualizations +""" + +import copy +import os +from typing import Tuple + +import cv2 +import matplotlib +import numpy as np +import requests +import trimesh +from scipy.spatial.transform import Rotation + + +def remove_unreferenced_vertices( + faces: np.ndarray, *vertice_attrs, return_indices: bool = False +) -> Tuple[np.ndarray, ...]: + """ + Remove unreferenced vertices of a mesh. + Unreferenced vertices are removed, and the face indices are updated accordingly. + + Args: + faces (np.ndarray): [T, P] face indices + *vertice_attrs: vertex attributes + + Returns: + faces (np.ndarray): [T, P] face indices + *vertice_attrs: vertex attributes + indices (np.ndarray, optional): [N] indices of vertices that are kept. Defaults to None. + """ + P = faces.shape[-1] + fewer_indices, inv_map = np.unique(faces, return_inverse=True) + faces = inv_map.astype(np.int32).reshape(-1, P) + ret = [faces] + for attr in vertice_attrs: + ret.append(attr[fewer_indices]) + if return_indices: + ret.append(fewer_indices) + return tuple(ret) + + +def triangulate( + faces: np.ndarray, vertices: np.ndarray = None, backslash: np.ndarray = None +) -> np.ndarray: + """ + Triangulate a polygonal mesh. + + Args: + faces (np.ndarray): [L, P] polygonal faces + vertices (np.ndarray, optional): [N, 3] 3-dimensional vertices. + If given, the triangulation is performed according to the distance + between vertices. Defaults to None. + backslash (np.ndarray, optional): [L] boolean array indicating + how to triangulate the quad faces. Defaults to None. + + Returns: + (np.ndarray): [L * (P - 2), 3] triangular faces + """ + if faces.shape[-1] == 3: + return faces + P = faces.shape[-1] + if vertices is not None: + assert faces.shape[-1] == 4, "now only support quad mesh" + if backslash is None: + backslash = np.linalg.norm( + vertices[faces[:, 0]] - vertices[faces[:, 2]], axis=-1 + ) < np.linalg.norm(vertices[faces[:, 1]] - vertices[faces[:, 3]], axis=-1) + if backslash is None: + loop_indice = np.stack( + [ + np.zeros(P - 2, dtype=int), + np.arange(1, P - 1, 1, dtype=int), + np.arange(2, P, 1, dtype=int), + ], + axis=1, + ) + return faces[:, loop_indice].reshape((-1, 3)) + else: + assert faces.shape[-1] == 4, "now only support quad mesh" + faces = np.where( + backslash[:, None], + faces[:, [0, 1, 2, 0, 2, 3]], + faces[:, [0, 1, 3, 3, 1, 2]], + ).reshape((-1, 3)) + return faces + + +def image_mesh( + *image_attrs: np.ndarray, + mask: np.ndarray = None, + tri: bool = False, + return_indices: bool = False, +) -> Tuple[np.ndarray, ...]: + """ + Get a mesh regarding image pixel uv coordinates as vertices and image grid as faces. + + Args: + *image_attrs (np.ndarray): image attributes in shape (height, width, [channels]) + mask (np.ndarray, optional): binary mask of shape (height, width), dtype=bool. Defaults to None. + + Returns: + faces (np.ndarray): faces connecting neighboring pixels. shape (T, 4) if tri is False, else (T, 3) + *vertex_attrs (np.ndarray): vertex attributes in corresponding order with input image_attrs + indices (np.ndarray, optional): indices of vertices in the original mesh + """ + assert (len(image_attrs) > 0) or (mask is not None), ( + "At least one of image_attrs or mask should be provided" + ) + height, width = next(image_attrs).shape[:2] if mask is None else mask.shape + assert all(img.shape[:2] == (height, width) for img in image_attrs), ( + "All image_attrs should have the same shape" + ) + + row_faces = np.stack( + [ + np.arange(0, width - 1, dtype=np.int32), + np.arange(width, 2 * width - 1, dtype=np.int32), + np.arange(1 + width, 2 * width, dtype=np.int32), + np.arange(1, width, dtype=np.int32), + ], + axis=1, + ) + faces = ( + np.arange(0, (height - 1) * width, width, dtype=np.int32)[:, None, None] + + row_faces[None, :, :] + ).reshape((-1, 4)) + if mask is None: + if tri: + faces = triangulate(faces) + ret = [faces, *(img.reshape(-1, *img.shape[2:]) for img in image_attrs)] + if return_indices: + ret.append(np.arange(height * width, dtype=np.int32)) + return tuple(ret) + else: + quad_mask = ( + mask[:-1, :-1] & mask[1:, :-1] & mask[1:, 1:] & mask[:-1, 1:] + ).ravel() + faces = faces[quad_mask] + if tri: + faces = triangulate(faces) + return remove_unreferenced_vertices( + faces, + *(x.reshape(-1, *x.shape[2:]) for x in image_attrs), + return_indices=return_indices, + ) + + +def predictions_to_glb( + predictions, + filter_by_frames="all", + mask_black_bg=False, + mask_white_bg=False, + show_cam=True, + mask_ambiguous=False, + as_mesh=True, +) -> trimesh.Scene: + """ + Converts VGGT predictions to a 3D scene represented as a GLB file. + + Args: + predictions (dict): Dictionary containing model predictions with keys: + - world_points: 3D point coordinates (S, H, W, 3) + - images: Input images (S, H, W, 3) + - extrinsic: Camera extrinsic matrices (S, 3, 4) + filter_by_frames (str): Frame filter specification (default: "all") + mask_black_bg (bool): Mask out black background pixels (default: False) + mask_white_bg (bool): Mask out white background pixels (default: False) + show_cam (bool): Include camera visualization (default: True) + mask_ambiguous (bool): Apply final mask to filter ambiguous predictions (default: False) + as_mesh (bool): Represent the data as a mesh instead of point cloud (default: False) + + Returns: + trimesh.Scene: Processed 3D scene containing point cloud/mesh and cameras + + Raises: + ValueError: If input predictions structure is invalid + """ + if not isinstance(predictions, dict): + raise ValueError("predictions must be a dictionary") + + print("Building GLB scene") + selected_frame_idx = None + if filter_by_frames != "all" and filter_by_frames != "All": + try: + # Extract the index part before the colon + selected_frame_idx = int(filter_by_frames.split(":")[0]) + except (ValueError, IndexError): + pass + + # Always use Pointmap Branch + print("Using Pointmap Branch") + if "world_points" not in predictions: + raise ValueError( + "world_points not found in predictions. Pointmap Branch requires 'world_points' key. " + "Depthmap and Camera branches have been removed." + ) + + pred_world_points = predictions["world_points"] + + # Get images from predictions + images = predictions["images"] + # Use extrinsic matrices instead of pred_extrinsic_list + camera_matrices = predictions["extrinsic"] + + if selected_frame_idx is not None: + pred_world_points = pred_world_points[selected_frame_idx][None] + images = images[selected_frame_idx][None] + camera_matrices = camera_matrices[selected_frame_idx][None] + + vertices_3d = pred_world_points.reshape(-1, 3) + # Handle different image formats - check if images need transposing + if images.ndim == 4 and images.shape[1] == 3: # NCHW format + colors_rgb = np.transpose(images, (0, 2, 3, 1)) + else: # Assume already in NHWC format + colors_rgb = images + colors_rgb = (colors_rgb.reshape(-1, 3) * 255).astype(np.uint8) + + # Create mask for filtering + mask = np.ones(len(vertices_3d), dtype=bool) + final_mask = predictions["final_mask"].reshape(-1) + + if mask_black_bg: + black_bg_mask = colors_rgb.sum(axis=1) >= 16 + mask = mask & black_bg_mask + + if mask_white_bg: + # Filter out white background pixels (RGB values close to white) + # Consider pixels white if all RGB values are above 240 + white_bg_mask = ( + (colors_rgb[:, 0] > 240) + & (colors_rgb[:, 1] > 240) + & (colors_rgb[:, 2] > 240) + ) + mask = mask & ~white_bg_mask + + # Use final_mask when mask_ambiguous is checked + if mask_ambiguous: + mask = mask & final_mask + + vertices_3d = vertices_3d[mask].copy() + colors_rgb = colors_rgb[mask].copy() + + if vertices_3d is None or np.asarray(vertices_3d).size == 0: + vertices_3d = np.array([[1, 0, 0]]) + colors_rgb = np.array([[255, 255, 255]]) + scene_scale = 1 + else: + # Calculate the 5th and 95th percentiles along each axis + lower_percentile = np.percentile(vertices_3d, 5, axis=0) + upper_percentile = np.percentile(vertices_3d, 95, axis=0) + + # Calculate the diagonal length of the percentile bounding box + scene_scale = np.linalg.norm(upper_percentile - lower_percentile) + + colormap = matplotlib.colormaps.get_cmap("gist_rainbow") + + # Initialize a 3D scene + scene_3d = trimesh.Scene() + + # Add point cloud data to the scene + if as_mesh: + # Create mesh from pointcloud + # try: + if selected_frame_idx is not None: + # Single frame case - we can create a proper mesh + H, W = pred_world_points.shape[1:3] + + # Get original unfiltered data for mesh creation + original_points = pred_world_points.reshape(H, W, 3) + + # Reshape original image data properly + if images.ndim == 4 and images.shape[1] == 3: # NCHW format + original_image_colors = np.transpose(images[0], (1, 2, 0)) + else: # Assume already in HWC format + original_image_colors = images[0] + original_image_colors *= 255 + # Get original final mask + original_final_mask = predictions["final_mask"][selected_frame_idx].reshape( + H, W + ) + + # Create mask based on final mask + mask = original_final_mask + + # Additional background masks if needed + if mask_black_bg: + black_bg_mask = original_image_colors.sum(axis=2) >= 16 + mask = mask & black_bg_mask + + if mask_white_bg: + white_bg_mask = ~( + (original_image_colors[:, :, 0] > 240) + & (original_image_colors[:, :, 1] > 240) + & (original_image_colors[:, :, 2] > 240) + ) + mask = mask & white_bg_mask + + # Check if normals are available in predictions + vertex_normals = None + if "normal" in predictions and predictions["normal"] is not None: + # Get normals for the selected frame + frame_normals = ( + predictions["normal"][selected_frame_idx] + if selected_frame_idx is not None + else predictions["normal"][0] + ) + + # Create faces and vertices using image_mesh with normals support + faces, vertices, vertex_colors, vertex_normals = image_mesh( + original_points * np.array([1, -1, 1], dtype=np.float32), + original_image_colors / 255.0, + frame_normals * np.array([1, -1, 1], dtype=np.float32), + mask=mask, + tri=True, + return_indices=False, + ) + + # Apply coordinate transformations to normals + vertex_normals = vertex_normals * np.array([1, -1, 1], dtype=np.float32) + else: + # Create faces and vertices using image_mesh without normals + faces, vertices, vertex_colors = image_mesh( + original_points * np.array([1, -1, 1], dtype=np.float32), + original_image_colors / 255.0, + mask=mask, + tri=True, + return_indices=False, + ) + + # vertices = vertices * np.array([1, -1, 1], dtype=np.float32) + + # Create trimesh object with optional normals + mesh_data = trimesh.Trimesh( + vertices=vertices * np.array([1, -1, 1], dtype=np.float32), + faces=faces, + vertex_colors=(vertex_colors * 255).astype(np.uint8), + vertex_normals=(vertex_normals if vertex_normals is not None else None), + process=False, + ) + scene_3d.add_geometry(mesh_data) + + else: + # Multi-frame case - create separate meshes for each frame + print("Creating mesh for multi-frame data...") + + for frame_idx in range(pred_world_points.shape[0]): + H, W = pred_world_points.shape[1:3] + + # Get data for this frame + frame_points = pred_world_points[frame_idx] + frame_final_mask = predictions["final_mask"][frame_idx] + + # Get frame image + if images.ndim == 4 and images.shape[1] == 3: # NCHW format + frame_image = np.transpose(images[frame_idx], (1, 2, 0)) + else: # Assume already in HWC format + frame_image = images[frame_idx] + frame_image *= 255 + # Create mask for this frame using final_mask + mask = frame_final_mask + + # Additional background masks if needed + if mask_black_bg: + black_bg_mask = frame_image.sum(axis=2) >= 16 + mask = mask & black_bg_mask + + if mask_white_bg: + white_bg_mask = ~( + (frame_image[:, :, 0] > 240) + & (frame_image[:, :, 1] > 240) + & (frame_image[:, :, 2] > 240) + ) + mask = mask & white_bg_mask + + # Create mesh for this frame + faces, vertices, vertex_colors = image_mesh( + frame_points * np.array([1, -1, 1], dtype=np.float32), + frame_image / 255.0, + mask=mask, + tri=True, + return_indices=False, + ) + + vertices = vertices * np.array([1, -1, 1], dtype=np.float32) + # Create trimesh object for this frame + frame_mesh = trimesh.Trimesh( + vertices=vertices, + faces=faces, + vertex_colors=(vertex_colors * 255).astype(np.uint8), + process=False, + ) + scene_3d.add_geometry(frame_mesh) + else: + point_cloud_data = trimesh.PointCloud(vertices=vertices_3d, colors=colors_rgb) + scene_3d.add_geometry(point_cloud_data) + + # Prepare 4x4 matrices for camera extrinsics + num_cameras = len(camera_matrices) + + if show_cam: + # Add camera models to the scene + for i in range(num_cameras): + world_to_camera = camera_matrices[i] + rgba_color = colormap(i / num_cameras) + current_color = tuple(int(255 * x) for x in rgba_color[:3]) + + integrate_camera_into_scene( + scene_3d, world_to_camera, current_color, scene_scale + ) + + # Align scene to the observation of the first camera + scene_3d = apply_scene_alignment(scene_3d, camera_matrices) + + print("GLB Scene built") + return scene_3d + + +def integrate_camera_into_scene( + scene: trimesh.Scene, + transform: np.ndarray, + face_colors: tuple, + scene_scale: float, +): + """ + Integrates a fake camera mesh into the 3D scene. + + Args: + scene (trimesh.Scene): The 3D scene to add the camera model. + transform (np.ndarray): Transformation matrix for camera positioning. + face_colors (tuple): Color of the camera face. + scene_scale (float): Scale of the scene. + """ + scene_scale = 12 + cam_width = scene_scale * 0.05 + cam_height = scene_scale * 0.1 + # cam_width = scene_scale * 0.05 + # cam_height = scene_scale * 0.1 + + # Create cone shape for camera + rot_45_degree = np.eye(4) + rot_45_degree[:3, :3] = Rotation.from_euler("z", 45, degrees=True).as_matrix() + rot_45_degree[2, 3] = -cam_height + + opengl_transform = get_opengl_conversion_matrix() + # Combine transformations + complete_transform = transform @ opengl_transform @ rot_45_degree + camera_cone_shape = trimesh.creation.cone(cam_width, cam_height, sections=4) + + # Generate mesh for the camera + slight_rotation = np.eye(4) + slight_rotation[:3, :3] = Rotation.from_euler("z", 2, degrees=True).as_matrix() + + vertices_combined = np.concatenate( + [ + camera_cone_shape.vertices, + 0.95 * camera_cone_shape.vertices, + transform_points(slight_rotation, camera_cone_shape.vertices), + ] + ) + vertices_transformed = transform_points(complete_transform, vertices_combined) + + mesh_faces = compute_camera_faces(camera_cone_shape) + + # Add the camera mesh to the scene + camera_mesh = trimesh.Trimesh(vertices=vertices_transformed, faces=mesh_faces) + camera_mesh.visual.face_colors[:, :3] = face_colors + scene.add_geometry(camera_mesh) + + +def apply_scene_alignment( + scene_3d: trimesh.Scene, extrinsics_matrices: np.ndarray +) -> trimesh.Scene: + """ + Aligns the 3D scene based on the extrinsics of the first camera. + + Args: + scene_3d (trimesh.Scene): The 3D scene to be aligned. + extrinsics_matrices (np.ndarray): Camera extrinsic matrices. + + Returns: + trimesh.Scene: Aligned 3D scene. + """ + # Set transformations for scene alignment + opengl_conversion_matrix = get_opengl_conversion_matrix() + + # Rotation matrix for alignment (180 degrees around the y-axis) + align_rotation = np.eye(4) + align_rotation[:3, :3] = Rotation.from_euler("y", 0, degrees=True).as_matrix() + + # Apply transformation + initial_transformation = ( + np.linalg.inv(extrinsics_matrices[0]) + @ opengl_conversion_matrix + @ align_rotation + ) + scene_3d.apply_transform(initial_transformation) + return scene_3d + + +def get_opengl_conversion_matrix() -> np.ndarray: + """ + Constructs and returns the OpenGL conversion matrix. + + Returns: + numpy.ndarray: A 4x4 OpenGL conversion matrix. + """ + # Create an identity matrix + matrix = np.identity(4) + + # Flip the y and z axes + matrix[1, 1] = -1 + matrix[2, 2] = -1 + + return matrix + + +def transform_points( + transformation: np.ndarray, points: np.ndarray, dim: int = None +) -> np.ndarray: + """ + Applies a 4x4 transformation to a set of points. + + Args: + transformation (np.ndarray): Transformation matrix. + points (np.ndarray): Points to be transformed. + dim (int, optional): Dimension for reshaping the result. + + Returns: + np.ndarray: Transformed points. + """ + points = np.asarray(points) + initial_shape = points.shape[:-1] + dim = dim or points.shape[-1] + + # Apply transformation + transformation = transformation.swapaxes( + -1, -2 + ) # Transpose the transformation matrix + points = points @ transformation[..., :-1, :] + transformation[..., -1:, :] + + # Reshape the result + result = points[..., :dim].reshape(*initial_shape, dim) + return result + + +def compute_camera_faces(cone_shape: trimesh.Trimesh) -> np.ndarray: + """ + Computes the faces for the camera mesh. + + Args: + cone_shape (trimesh.Trimesh): The shape of the camera cone. + + Returns: + np.ndarray: Array of faces for the camera mesh. + """ + # Create pseudo cameras + faces_list = [] + num_vertices_cone = len(cone_shape.vertices) + + for face in cone_shape.faces: + if 0 in face: + continue + v1, v2, v3 = face + v1_offset, v2_offset, v3_offset = face + num_vertices_cone + v1_offset_2, v2_offset_2, v3_offset_2 = face + 2 * num_vertices_cone + + faces_list.extend( + [ + (v1, v2, v2_offset), + (v1, v1_offset, v3), + (v3_offset, v2, v3), + (v1, v2, v2_offset_2), + (v1, v1_offset_2, v3), + (v3_offset_2, v2, v3), + ] + ) + + faces_list += [(v3, v2, v1) for v1, v2, v3 in faces_list] + return np.array(faces_list) + + +def segment_sky(image_path, onnx_session, mask_filename=None): + """ + Segments sky from an image using an ONNX model. + Thanks for the great model provided by https://github.com/xiongzhu666/Sky-Segmentation-and-Post-processing + + Args: + image_path: Path to input image + onnx_session: ONNX runtime session with loaded model + mask_filename: Path to save the output mask + + Returns: + np.ndarray: Binary mask where 255 indicates non-sky regions + """ + + assert mask_filename is not None + image = cv2.imread(image_path) + + result_map = run_skyseg(onnx_session, [320, 320], image) + # resize the result_map to the original image size + result_map_original = cv2.resize(result_map, (image.shape[1], image.shape[0])) + + # Fix: Invert the mask so that 255 = non-sky, 0 = sky + # The model outputs low values for sky, high values for non-sky + output_mask = np.zeros_like(result_map_original) + output_mask[result_map_original < 32] = 255 # Use threshold of 32 + + os.makedirs(os.path.dirname(mask_filename), exist_ok=True) + cv2.imwrite(mask_filename, output_mask) + return output_mask + + +def run_skyseg(onnx_session, input_size, image): + """ + Runs sky segmentation inference using ONNX model. + + Args: + onnx_session: ONNX runtime session + input_size: Target size for model input (width, height) + image: Input image in BGR format + + Returns: + np.ndarray: Segmentation mask + """ + + # Pre process:Resize, BGR->RGB, Transpose, PyTorch standardization, float32 cast + temp_image = copy.deepcopy(image) + resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1])) + x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB) + x = np.array(x, dtype=np.float32) + mean = [0.485, 0.456, 0.406] + std = [0.229, 0.224, 0.225] + x = (x / 255 - mean) / std + x = x.transpose(2, 0, 1) + x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32") + + # Inference + input_name = onnx_session.get_inputs()[0].name + output_name = onnx_session.get_outputs()[0].name + onnx_result = onnx_session.run([output_name], {input_name: x}) + + # Post process + onnx_result = np.array(onnx_result).squeeze() + min_value = np.min(onnx_result) + max_value = np.max(onnx_result) + onnx_result = (onnx_result - min_value) / (max_value - min_value) + onnx_result *= 255 + onnx_result = onnx_result.astype("uint8") + + return onnx_result + + +def download_file_from_url(url, filename): + """Downloads a file from a Hugging Face model repo, handling redirects.""" + try: + # Get the redirect URL + response = requests.get(url, allow_redirects=False) + response.raise_for_status() # Raise HTTPError for bad requests (4xx or 5xx) + + if response.status_code == 302: # Expecting a redirect + redirect_url = response.headers["Location"] + response = requests.get(redirect_url, stream=True) + response.raise_for_status() + else: + print(f"Unexpected status code: {response.status_code}") + return + + with open(filename, "wb") as f: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + print(f"Downloaded {filename} successfully.") + + except requests.exceptions.RequestException as e: + print(f"Error downloading file: {e}") diff --git a/mapanything/utils/image.py b/mapanything/utils/image.py new file mode 100644 index 0000000000000000000000000000000000000000..fee62cff92227065b93c94fa887e2934552b3f92 --- /dev/null +++ b/mapanything/utils/image.py @@ -0,0 +1,380 @@ +""" +Utility functions for loading, converting, and manipulating images. + +This module provides functions for: +- Converting between different image formats and representations +- Resizing and cropping images to specific resolutions +- Loading and normalizing images for model input +- Handling various image file formats including HEIF/HEIC when available +""" + +import os + +import numpy as np +import PIL.Image +import torch +import torchvision.transforms as tvf +from PIL.ImageOps import exif_transpose + +os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1" +import cv2 + +from mapanything.utils.moge_inference import run_moge_inference, load_moge_model + +try: + from pillow_heif import register_heif_opener + + register_heif_opener() + heif_support_enabled = True +except ImportError: + heif_support_enabled = False + +from mapanything.utils.cropping import crop_resize_if_necessary +from uniception.models.encoders.image_normalizations import IMAGE_NORMALIZATION_DICT + +# Fixed resolution mappings with precomputed aspect ratios as keys +RESOLUTION_MAPPINGS = { + 518: { + 1.000: (518, 518), # 1:1 + 1.321: (518, 392), # 4:3 + 1.542: (518, 336), # 3:2 + 1.762: (518, 294), # 16:9 + 2.056: (518, 252), # 2:1 + 3.083: (518, 168), # 3.2:1 + 0.757: (392, 518), # 3:4 + 0.649: (336, 518), # 2:3 + 0.567: (294, 518), # 9:16 + 0.486: (252, 518), # 1:2 + }, + 512: { + 1.000: (512, 512), # 1:1 + 1.333: (512, 384), # 4:3 + 1.524: (512, 336), # 3:2 + 1.778: (512, 288), # 16:9 + 2.000: (512, 256), # 2:1 + 3.200: (512, 160), # 3.2:1 + 0.750: (384, 512), # 3:4 + 0.656: (336, 512), # 2:3 + 0.562: (288, 512), # 9:16 + 0.500: (256, 512), # 1:2 + }, +} + +# Precomputed sorted aspect ratio keys for efficient lookup +ASPECT_RATIO_KEYS = { + 518: sorted(RESOLUTION_MAPPINGS[518].keys()), + 512: sorted(RESOLUTION_MAPPINGS[512].keys()), +} + +def img_to_arr(img): + if isinstance(img, str): + img = imread_cv2(img) + return img + + +def imread_cv2(path, options=cv2.IMREAD_COLOR): + """Open an image or a depthmap with opencv-python.""" + if path.endswith((".exr", "EXR")): + options = cv2.IMREAD_ANYDEPTH + img = cv2.imread(path, options) + if img is None: + raise IOError(f"Could not load image={path} with {options=}") + if img.ndim == 3: + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return img + +def find_closest_aspect_ratio(aspect_ratio, resolution_set): + """ + Find the closest aspect ratio from the resolution mappings using efficient key lookup. + + Args: + aspect_ratio (float): Target aspect ratio + resolution_set (int): Resolution set to use (518 or 512) + + Returns: + tuple: (target_width, target_height) from the resolution mapping + """ + aspect_keys = ASPECT_RATIO_KEYS[resolution_set] + + # Find the closest aspect ratio key using binary search approach + closest_key = min(aspect_keys, key=lambda x: abs(x - aspect_ratio)) + + return RESOLUTION_MAPPINGS[resolution_set][closest_key] + + +def rgb(ftensor, norm_type, true_shape=None): + """ + Convert normalized image tensor to RGB image for visualization. + + Args: + ftensor (torch.Tensor or numpy.ndarray or list): Image tensor or list of image tensors + norm_type (str): Normalization type, see UniCeption IMAGE_NORMALIZATION_DICT keys or use "identity" + true_shape (tuple, optional): If provided, the image will be cropped to this shape (H, W) + + Returns: + numpy.ndarray: RGB image with values in range [0, 1] + """ + if isinstance(ftensor, list): + return [rgb(x, norm_type, true_shape=true_shape) for x in ftensor] + if isinstance(ftensor, torch.Tensor): + ftensor = ftensor.detach().cpu().numpy() # H,W,3 + if ftensor.ndim == 3 and ftensor.shape[0] == 3: + ftensor = ftensor.transpose(1, 2, 0) + elif ftensor.ndim == 4 and ftensor.shape[1] == 3: + ftensor = ftensor.transpose(0, 2, 3, 1) + if true_shape is not None: + H, W = true_shape + ftensor = ftensor[:H, :W] + if ftensor.dtype == np.uint8: + img = np.float32(ftensor) / 255 + else: + if norm_type in IMAGE_NORMALIZATION_DICT.keys(): + img_norm = IMAGE_NORMALIZATION_DICT[norm_type] + mean = img_norm.mean.numpy() + std = img_norm.std.numpy() + elif norm_type == "identity": + mean = 0.0 + std = 1.0 + else: + raise ValueError( + f"Unknown image normalization type: {norm_type}. Available types: identity or {IMAGE_NORMALIZATION_DICT.keys()}" + ) + img = ftensor * std + mean + return img.clip(min=0, max=1) + + +def load_images( + folder_or_list, + resize_mode="fixed_mapping", + size=None, + norm_type="dinov2", + patch_size=14, + verbose=False, + bayer_format=False, + resolution_set=518, + stride=1, + compute_moge_mask=False, + moge_model=None, + binary_mask_path=None +): + """ + Open and convert all images in a list or folder to proper input format for model + + Args: + folder_or_list (str or list): Path to folder or list of image paths. + resize_mode (str): Resize mode - "fixed_mapping", "longest_side", "square", or "fixed_size". Defaults to "fixed_mapping". + size (int or tuple, optional): Required for "longest_side", "square", and "fixed_size" modes. + - For "longest_side" and "square": int value for resize dimension + - For "fixed_size": tuple of (width, height) + norm_type (str, optional): Image normalization type. See UniCeption IMAGE_NORMALIZATION_DICT keys. Defaults to "dinov2". + patch_size (int, optional): Patch size for image processing. Defaults to 14. + verbose (bool, optional): If True, print progress messages. Defaults to False. + bayer_format (bool, optional): If True, read images in Bayer format. Defaults to False. + resolution_set (int, optional): Resolution set to use for "fixed_mapping" mode (518 or 512). Defaults to 518. + stride (int, optional): Load every nth image from the input. stride=1 loads all images, stride=2 loads every 2nd image, etc. Defaults to 1. + + Returns: + list: List of dictionaries containing image data and metadata + """ + # Validate resize_mode and size parameter requirements + valid_resize_modes = ["fixed_mapping", "longest_side", "square", "fixed_size"] + if resize_mode not in valid_resize_modes: + raise ValueError( + f"Resize_mode must be one of {valid_resize_modes}, got '{resize_mode}'" + ) + + if resize_mode in ["longest_side", "square", "fixed_size"] and size is None: + raise ValueError(f"Size parameter is required for resize_mode='{resize_mode}'") + + # Validate size type based on resize mode + if resize_mode in ["longest_side", "square"]: + if not isinstance(size, int): + raise ValueError( + f"Size must be an int for resize_mode='{resize_mode}', got {type(size)}" + ) + elif resize_mode == "fixed_size": + if not isinstance(size, (tuple, list)) or len(size) != 2: + raise ValueError( + f"Size must be a tuple/list of (width, height) for resize_mode='fixed_size', got {size}" + ) + if not all(isinstance(x, int) for x in size): + raise ValueError( + f"Size values must be integers for resize_mode='fixed_size', got {size}" + ) + + # Get list of image paths + if isinstance(folder_or_list, str): + # If folder_or_list is a string, assume it's a path to a folder + if verbose: + print(f"Loading images from {folder_or_list}") + root, folder_content = folder_or_list, sorted(os.listdir(folder_or_list)) + elif isinstance(folder_or_list, list): + # If folder_or_list is a list, assume it's a list of image paths + if verbose: + print(f"Loading a list of {len(folder_or_list)} images") + root, folder_content = "", folder_or_list + else: + # If folder_or_list is neither a string nor a list, raise an error + raise ValueError(f"Bad {folder_or_list=} ({type(folder_or_list)})") + + # Define supported image extensions + supported_images_extensions = [".jpg", ".jpeg", ".png"] + if heif_support_enabled: + supported_images_extensions += [".heic", ".heif"] + supported_images_extensions = tuple(supported_images_extensions) + + # First pass: Load all images and collect aspect ratios + loaded_images = [] + aspect_ratios = [] + for i, path in enumerate(folder_content): + # Skip images based on stride + if i % stride != 0: + continue + + # Check if the file has a supported image extension + if not path.lower().endswith(supported_images_extensions): + continue + + try: + if bayer_format: + # If bayer_format is True, read the image in Bayer format + color_bayer = cv2.imread(os.path.join(root, path), cv2.IMREAD_UNCHANGED) + color = cv2.cvtColor(color_bayer, cv2.COLOR_BAYER_RG2BGR) + img = PIL.Image.fromarray(color) + img = exif_transpose(img).convert("RGB") + else: + # Otherwise, read the image normally + img = exif_transpose(PIL.Image.open(os.path.join(root, path))).convert( + "RGB" + ) + + W1, H1 = img.size + aspect_ratios.append(W1 / H1) + loaded_images.append((path, img, W1, H1)) + + except Exception as e: + if verbose: + print(f"Warning: Could not load {path}: {e}") + continue + + # Check if any images were loaded + if not loaded_images: + raise ValueError("No valid images found") + + # Calculate average aspect ratio and determine target size + average_aspect_ratio = sum(aspect_ratios) / len(aspect_ratios) + if verbose: + print( + f"Calculated average aspect ratio: {average_aspect_ratio:.3f} from {len(aspect_ratios)} images" + ) + + # Determine target size for all images based on resize mode + if resize_mode == "fixed_mapping": + # Resolution mappings are already compatible with their respective patch sizes + # 518 mappings are divisible by 14, 512 mappings are divisible by 16 + target_width, target_height = find_closest_aspect_ratio( + average_aspect_ratio, resolution_set + ) + target_size = (target_width, target_height) + elif resize_mode == "square": + target_size = ( + round((size // patch_size)) * patch_size, + round((size // patch_size)) * patch_size, + ) + elif resize_mode == "longest_side": + # Use average aspect ratio to determine size for all images + # Longest side should be the input size + if average_aspect_ratio >= 1: # Landscape or square + # Width is the longest side + target_size = ( + size, + round((size // patch_size) / average_aspect_ratio) * patch_size, + ) + else: # Portrait + # Height is the longest side + target_size = ( + round((size // patch_size) * average_aspect_ratio) * patch_size, + size, + ) + elif resize_mode == "fixed_size": + # Use exact size provided, aligned to patch_size + target_size = ( + (size[0] // patch_size) * patch_size, + (size[1] // patch_size) * patch_size, + ) + + if verbose: + print( + f"Using target resolution {target_size[0]}x{target_size[1]} (W x H) for all images" + ) + + # Get the image normalization function based on the norm_type + if norm_type in IMAGE_NORMALIZATION_DICT.keys(): + img_norm = IMAGE_NORMALIZATION_DICT[norm_type] + ImgNorm = tvf.Compose( + [tvf.ToTensor(), tvf.Normalize(mean=img_norm.mean, std=img_norm.std)] + ) + else: + raise ValueError( + f"Unknown image normalization type: {norm_type}. Available options: {list(IMAGE_NORMALIZATION_DICT.keys())}" + ) + + # Second pass: Resize all images to the same target size + imgs = [] + for path, img, W1, H1 in loaded_images: + + # Load binary mask if path is provided + if binary_mask_path is not None: + mask_img = cv2.imread(binary_mask_path) + mask_img = cv2.cvtColor(mask_img, cv2.COLOR_BGR2GRAY) + binary_mask = (mask_img > 0).astype(np.float32) + + # # Dilate binary_mask further + kernel_size = 3 + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + binary_mask = cv2.erode(binary_mask, kernel, iterations=1) + + else: + binary_mask = np.ones((H1, W1), dtype=np.float32) # Default mask, all pixels are valid + + # check if we need to compute moge mask + if compute_moge_mask: + moge_img = np.array(PIL.Image.open(os.path.join(root, path))) # (H, W, 3) + transform = tvf.Compose([tvf.ToTensor()]) + input_moge_img = transform(moge_img).unsqueeze(0) # (1, 3, H, W) + moge_output = run_moge_inference(moge_model, input_moge_img, device="cuda") + non_ambiguous_mask = moge_output["mask"].squeeze(0).cpu().numpy().astype(np.float32) # (H, W) + additional_quantities = [non_ambiguous_mask, binary_mask] + # img, additional_quantities = crop_resize_if_necessary(img, resolution=size, additional_quantities=additional_quantities) + img, additional_quantities = crop_resize_if_necessary(img, resolution=target_size, additional_quantities=additional_quantities) + non_ambiguous_mask = torch.tensor(additional_quantities[0]).bool() + binary_mask = torch.tensor(additional_quantities[1]).bool() + else: + additional_quantities = None + # img = crop_resize_if_necessary(img, resolution=size)[0] + img = crop_resize_if_necessary(img, resolution=target_size)[0] + non_ambiguous_mask = torch.tensor(np.ones_like(img)).bool() # Default mask, all pixels are valid + binary_mask = torch.tensor(np.ones_like(img)).bool() + + # Normalize image and add it to the list + W2, H2 = img.size + if verbose: + print(f" - Adding {path} with resolution {W1}x{H1} --> {W2}x{H2}") + + imgs.append( + dict( + img=ImgNorm(img)[None], + true_shape=np.int32([img.size[::-1]]), + idx=len(imgs), + instance=str(len(imgs)), + data_norm_type=[norm_type], + non_ambiguous_mask=non_ambiguous_mask, + binary_mask=binary_mask + ) + ) + + assert imgs, "No images foud at " + root + if verbose: + print(f" (Found {len(imgs)} images)") + + return imgs diff --git a/mapanything/utils/inference.py b/mapanything/utils/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..d8b7b10a3925ae7b9349c44aa97f90c110da79d4 --- /dev/null +++ b/mapanything/utils/inference.py @@ -0,0 +1,476 @@ +""" +Inference utilities. +""" + +import warnings +from typing import Any, Dict, List + +import numpy as np +import torch + +from mapanything.utils.geometry import ( + depth_edge, + get_rays_in_camera_frame, + normals_edge, + points_to_normals, + quaternion_to_rotation_matrix, + recover_pinhole_intrinsics_from_ray_directions, + rotation_matrix_to_quaternion, +) +from mapanything.utils.image import rgb + +# Hard constraints - exactly what users can provide +ALLOWED_VIEW_KEYS = { + "img", # Required - input images + "data_norm_type", # Required - normalization type of the input images + "depth_z", # Optional - Z depth maps + "ray_directions", # Optional - ray directions in camera frame + "intrinsics", # Optional - pinhole camera intrinsics (conflicts with ray_directions) + "camera_poses", # Optional - camera poses + "is_metric_scale", # Optional - whether inputs are metric scale + "true_shape", # Optional - original image shape + "idx", # Optional - index of the view + "instance", # Optional - instance info of the view +} + +REQUIRED_KEYS = {"img", "data_norm_type"} + +# Define conflicting keys that cannot be used together +CONFLICTING_KEYS = [ + ("intrinsics", "ray_directions") # Both represent camera projection +] + + +def loss_of_one_batch_multi_view( + batch, + model, + criterion, + device, + use_amp=False, + amp_dtype="bf16", + ret=None, + ignore_keys=None, +): + """ + Calculate loss for a batch with multiple views. + + Args: + batch (list): List of view dictionaries containing input data. + model (torch.nn.Module): Model to run inference with. + criterion (callable, optional): Loss function to compute the loss. + device (torch.device): Device to run the computation on. + use_amp (bool, optional): Whether to use automatic mixed precision. Defaults to False. + amp_dtype (str, optional): Floating point type to use for automatic mixed precision. Options: ["fp32", "fp16", "bf16"]. Defaults to "bf16". + ret (str, optional): If provided, return only the specified key from the result dictionary. + ignore_keys (set, optional): Set of keys to ignore when moving tensors to device. + Defaults to {"dataset", "label", "instance", + "idx", "true_shape", "rng", "data_norm_type"}. + + Returns: + dict or Any: If ret is None, returns a dictionary containing views, predictions, and loss. + Otherwise, returns the value associated with the ret key. + """ + # Move necessary tensors to device + if ignore_keys is None: + ignore_keys = set( + [ + "depthmap", + "dataset", + "label", + "instance", + "idx", + "true_shape", + "rng", + "data_norm_type", + "scene_flow_compute_type", + ] + ) + for view in batch: + for name in view.keys(): + if name in ignore_keys: + continue + view[name] = view[name].to(device, non_blocking=True) + + # Determine the mixed precision floating point type + if use_amp: + if amp_dtype == "fp16": + amp_dtype = torch.float16 + elif amp_dtype == "bf16": + if torch.cuda.is_bf16_supported(): + amp_dtype = torch.bfloat16 + else: + warnings.warn( + "bf16 is not supported on this device. Using fp16 instead." + ) + amp_dtype = torch.float16 + elif amp_dtype == "fp32": + amp_dtype = torch.float32 + else: + amp_dtype = torch.float32 + + # Run model and compute loss + with torch.autocast("cuda", enabled=bool(use_amp), dtype=amp_dtype): + preds = model(batch) + with torch.autocast("cuda", enabled=False): + loss = criterion(batch, preds) if criterion is not None else None + + result = {f"view{i + 1}": view for i, view in enumerate(batch)} + result.update({f"pred{i + 1}": pred for i, pred in enumerate(preds)}) + result["loss"] = loss + + return result[ret] if ret else result + + +def validate_input_views_for_inference( + views: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """ + Strict validation and preprocessing of input views. + + Args: + views: List of view dictionaries + + Returns: + Validated and preprocessed views + + Raises: + ValueError: For invalid keys, missing required keys, conflicting inputs, or invalid camera pose constraints + """ + # Ensure input is not empty + if not views: + raise ValueError("At least one view must be provided") + + # Track which views have camera poses + views_with_poses = [] + + # Validate each view + for view_idx, view in enumerate(views): + # Check for invalid keys + provided_keys = set(view.keys()) + invalid_keys = provided_keys - ALLOWED_VIEW_KEYS + if invalid_keys: + raise ValueError( + f"View {view_idx} contains invalid keys: {invalid_keys}. " + f"Allowed keys are: {sorted(ALLOWED_VIEW_KEYS)}" + ) + + # Check for missing required keys + missing_keys = REQUIRED_KEYS - provided_keys + if missing_keys: + raise ValueError(f"View {view_idx} missing required keys: {missing_keys}") + + # Check for conflicting keys + for conflict_set in CONFLICTING_KEYS: + present_conflicts = [key for key in conflict_set if key in provided_keys] + if len(present_conflicts) > 1: + raise ValueError( + f"View {view_idx} contains conflicting keys: {present_conflicts}. " + f"Only one of {conflict_set} can be provided at a time." + ) + + # Check depth constraint: If depth is provided, intrinsics or ray_directions must also be provided + if "depth_z" in provided_keys: + if ( + "intrinsics" not in provided_keys + and "ray_directions" not in provided_keys + ): + raise ValueError( + f"View {view_idx} depth constraint violation: If 'depth_z' is provided, " + f"then 'intrinsics' or 'ray_directions' must also be provided. " + f"Z Depth values require camera calibration information to be meaningful for an image." + ) + + # Track views with camera poses + if "camera_poses" in provided_keys: + views_with_poses.append(view_idx) + + # Cross-view constraint: If any view has camera_poses, view 0 must have them too + if views_with_poses and 0 not in views_with_poses: + raise ValueError( + f"Camera pose constraint violation: Views {views_with_poses} have camera_poses, " + f"but view 0 (reference view) does not. When using camera_poses, the first view " + f"must also provide camera_poses to serve as the reference frame." + ) + + return views + + +def preprocess_input_views_for_inference( + views: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """ + Pre-process input views to match the expected internal input format. + + The following steps are performed: + 1. Convert intrinsics to ray directions when required. If ray directions are already provided, unit normalize them. + 2. Convert depth_z to depth_along_ray + 3. Convert camera_poses to the expected input keys (camera_pose_quats and camera_pose_trans) + 4. Default is_metric_scale to True when not provided + + Args: + views: List of view dictionaries + + Returns: + Preprocessed views with consistent internal format + """ + processed_views = [] + + for view_idx, view in enumerate(views): + # Copy the view dictionary to avoid modifying the original input + processed_view = dict(view) + + # Step 1: Convert intrinsics to ray_directions when required. If ray_directions are provided, unit normalize them. + if "intrinsics" in view: + images = view["img"] + height, width = images.shape[-2:] + intrinsics = view["intrinsics"] + _, ray_directions = get_rays_in_camera_frame( + intrinsics=intrinsics, + height=height, + width=width, + normalize_to_unit_sphere=True, + ) + processed_view["ray_directions"] = ray_directions + del processed_view["intrinsics"] + elif "ray_directions" in view: + ray_directions = view["ray_directions"] + ray_norm = torch.norm(ray_directions, dim=-1, keepdim=True) + processed_view["ray_directions"] = ray_directions / (ray_norm + 1e-8) + + # Step 2: Convert depth_z to depth_along_ray + if "depth_z" in view: + depth_z = view["depth_z"] + ray_directions = processed_view["ray_directions"] + ray_directions_unit_plane = ray_directions / ray_directions[..., 2:3] + pts3d_cam = depth_z * ray_directions_unit_plane + depth_along_ray = torch.norm(pts3d_cam, dim=-1, keepdim=True) + processed_view["depth_along_ray"] = depth_along_ray + del processed_view["depth_z"] + + # Step 3: Convert camera_poses to expected input keys + if "camera_poses" in view: + camera_poses = view["camera_poses"] + if isinstance(camera_poses, tuple) and len(camera_poses) == 2: + quats, trans = camera_poses + processed_view["camera_pose_quats"] = quats + processed_view["camera_pose_trans"] = trans + elif torch.is_tensor(camera_poses) and camera_poses.shape[-2:] == (4, 4): + rotation_matrices = camera_poses[:, :3, :3] + translation_vectors = camera_poses[:, :3, 3] + quats = rotation_matrix_to_quaternion(rotation_matrices) + processed_view["camera_pose_quats"] = quats + processed_view["camera_pose_trans"] = translation_vectors + else: + raise ValueError( + f"View {view_idx}: camera_poses must be either a tuple of (quats, trans) " + f"or a tensor of (B, 4, 4) transformation matrices." + ) + del processed_view["camera_poses"] + + # Step 4: Default is_metric_scale to True when not provided + if "is_metric_scale" not in processed_view: + # Get batch size from the image tensor + batch_size = view["img"].shape[0] + # Default to True for all samples in the batch + processed_view["is_metric_scale"] = torch.ones( + batch_size, dtype=torch.bool, device=view["img"].device + ) + + # Rename keys to match expected model input format + if "ray_directions" in processed_view: + processed_view["ray_directions_cam"] = processed_view["ray_directions"] + del processed_view["ray_directions"] + + # Append the processed view to the list + processed_views.append(processed_view) + + return processed_views + + +def postprocess_model_outputs_for_inference( + raw_outputs: List[Dict[str, torch.Tensor]], + input_views: List[Dict[str, Any]], + apply_mask: bool = True, + mask_edges: bool = True, + edge_normal_threshold: float = 5.0, + edge_depth_threshold: float = 0.03, + apply_confidence_mask: bool = False, + confidence_percentile: float = 10, +) -> List[Dict[str, torch.Tensor]]: + """ + Post-process raw model outputs by copying raw outputs and adding essential derived fields. + + This function simplifies the raw model outputs by: + 1. Copying all raw outputs as-is + 2. Adding denormalized images (img_no_norm) + 3. Adding Z depth (depth_z) from camera frame points + 4. Recovering pinhole camera intrinsics from ray directions + 5. Adding camera pose matrices (camera_poses) if pose data is available + 6. Applying mask to dense geometry outputs if requested (supports edge masking and confidence masking) + + Args: + raw_outputs: List of raw model output dictionaries, one per view + input_views: List of original input view dictionaries, one per view + apply_mask: Whether to apply non-ambiguous mask to dense outputs. Defaults to True. + mask_edges: Whether to compute an edge mask based on normals and depth and apply it to the output. Defaults to True. + apply_confidence_mask: Whether to apply the confidence mask to the output. Defaults to False. + confidence_percentile: The percentile to use for the confidence threshold. Defaults to 10. + + Returns: + List of processed output dictionaries containing: + - All original raw outputs (after masking dense geometry outputs if requested) + - 'img_no_norm': Denormalized RGB images (B, H, W, 3) + - 'depth_z': Z depth from camera frame (B, H, W, 1) if points in camera frame available + - 'intrinsics': Recovered pinhole camera intrinsics (B, 3, 3) if ray directions available + - 'camera_poses': 4x4 pose matrices (B, 4, 4) if pose data available + - 'mask': comprehensive mask for dense geometry outputs (B, H, W, 1) if requested + + """ + processed_outputs = [] + + for view_idx, (raw_output, original_view) in enumerate( + zip(raw_outputs, input_views) + ): + # Start by copying all raw outputs + processed_output = dict(raw_output) + + # 1. Add denormalized images + img = original_view["img"] # Shape: (B, 3, H, W) + data_norm_type = original_view["data_norm_type"][0] + img_hwc = rgb(img, data_norm_type) + + # Convert numpy back to torch if needed (rgb returns numpy) + if isinstance(img_hwc, np.ndarray): + img_hwc = torch.from_numpy(img_hwc).to(img.device) + + processed_output["img_no_norm"] = img_hwc + + # 2. Add Z depth if we have camera frame points + if "pts3d_cam" in processed_output: + processed_output["depth_z"] = processed_output["pts3d_cam"][..., 2:3] + + # 3. Recover pinhole camera intrinsics from ray directions if available + if "ray_directions" in processed_output: + intrinsics = recover_pinhole_intrinsics_from_ray_directions( + processed_output["ray_directions"] + ) + processed_output["intrinsics"] = intrinsics + + # 4. Add camera pose matrices if both translation and quaternions are available + if "cam_trans" in processed_output and "cam_quats" in processed_output: + cam_trans = processed_output["cam_trans"] # (B, 3) + cam_quats = processed_output["cam_quats"] # (B, 4) + batch_size = cam_trans.shape[0] + + # Convert quaternions to rotation matrices + rotation_matrices = quaternion_to_rotation_matrix(cam_quats) # (B, 3, 3) + + # Create 4x4 pose matrices + pose_matrices = ( + torch.eye(4, device=img.device).unsqueeze(0).repeat(batch_size, 1, 1) + ) + pose_matrices[:, :3, :3] = rotation_matrices + pose_matrices[:, :3, 3] = cam_trans + + processed_output["camera_poses"] = pose_matrices # (B, 4, 4) + + # 5. Apply comprehensive mask to dense geometry outputs if requested + if apply_mask: + final_mask = None + + # Start with non-ambiguous mask if available + if "non_ambiguous_mask" in processed_output: + non_ambiguous_mask = ( + processed_output["non_ambiguous_mask"].cpu().numpy() + ) # (B, H, W) + final_mask = non_ambiguous_mask + + # Apply confidence mask if requested and available + if apply_confidence_mask and "conf" in processed_output: + confidences = processed_output["conf"].cpu() # (B, H, W) + # Compute percentile threshold for each batch element + batch_size = confidences.shape[0] + conf_mask = torch.zeros_like(confidences, dtype=torch.bool) + percentile_threshold = ( + torch.quantile( + confidences.reshape(batch_size, -1), + confidence_percentile / 100.0, + dim=1, + ) + .unsqueeze(-1) + .unsqueeze(-1) + ) # Shape: (B, 1, 1) + + # Compute mask for each batch element + conf_mask = confidences > percentile_threshold + conf_mask = conf_mask.numpy() + + if final_mask is not None: + final_mask = final_mask & conf_mask + else: + final_mask = conf_mask + + # Apply edge mask if requested and we have the required data + if mask_edges and final_mask is not None and "pts3d" in processed_output: + # Get 3D points for edge computation + pred_pts3d = processed_output["pts3d"].cpu().numpy() # (B, H, W, 3) + batch_size, height, width = final_mask.shape + + edge_masks = [] + for b in range(batch_size): + batch_final_mask = final_mask[b] # (H, W) + batch_pts3d = pred_pts3d[b] # (H, W, 3) + + if batch_final_mask.any(): # Only compute if we have valid points + # Compute normals and normal-based edge mask + normals, normals_mask = points_to_normals( + batch_pts3d, mask=batch_final_mask + ) + normal_edges = normals_edge( + normals, tol=edge_normal_threshold, mask=normals_mask + ) + + # Compute depth-based edge mask + depth_z = ( + processed_output["depth_z"][b].squeeze(-1).cpu().numpy() + ) + depth_edges = depth_edge( + depth_z, rtol=edge_depth_threshold, mask=batch_final_mask + ) + + # Combine both edge types + edge_mask = ~(depth_edges & normal_edges) + edge_masks.append(edge_mask) + else: + # No valid points, keep all as invalid + edge_masks.append(np.zeros_like(batch_final_mask, dtype=bool)) + + # Stack batch edge masks and combine with final mask + edge_mask = np.stack(edge_masks, axis=0) # (B, H, W) + final_mask = final_mask & edge_mask + + # Apply final mask to dense geometry outputs if we have a mask + if final_mask is not None: + # Convert mask to torch tensor + final_mask_torch = torch.from_numpy(final_mask).to( + processed_output["pts3d"].device + ) + final_mask_torch = final_mask_torch.unsqueeze(-1) # (B, H, W, 1) + + # Apply mask to dense geometry outputs (zero out invalid regions) + dense_geometry_keys = [ + "pts3d", + "pts3d_cam", + "depth_along_ray", + "depth_z", + ] + for key in dense_geometry_keys: + if key in processed_output: + processed_output[key] = processed_output[key] * final_mask_torch + + # Add mask to processed output + processed_output["mask"] = final_mask_torch + + processed_outputs.append(processed_output) + + return processed_outputs diff --git a/mapanything/utils/metrics.py b/mapanything/utils/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..29d28b5a6856daface74ddebe50d51ca9548ee51 --- /dev/null +++ b/mapanything/utils/metrics.py @@ -0,0 +1,504 @@ +""" +Utils for Metrics +Source for Pose AUC Metrics: VGGT +""" + +import math + +import numpy as np +import torch +import torch.nn.functional as F + + +def l2_distance_of_unit_quats_to_angular_error(l2_distance): + """ + Converts a given L2 distance (for unit quaternions) to the angular error in degrees. + For two quaternions differing by an angle θ the relationship is: + L2 distance = 2 * sin(θ/4) + Hence, the angular error in degrees is computed as: + 4 * asin(l2_distance / 2) * (180/π) + + Args: + l2_distance: L2 distance between two unit quaternions (torch.Tensor, shape: (N,)) + Returns: + angular_error_degrees: Angular error in degrees (torch.Tensor, shape: (N,)) + """ + angular_error_radians = 4 * torch.asin(l2_distance / 2) + angular_error_degrees = angular_error_radians * 180.0 / math.pi + + return angular_error_degrees + + +def l2_distance_of_unit_ray_directions_to_angular_error(l2_distance): + """ + Converts a given L2 distance (for unit ray directions) to the angular error in degrees. + For two unit ray directions differing by an angle θ the relationship is: + L2 distance = 2 * sin(θ/2) + Hence, the angular error in degrees is computed as: + 2 * asin(l2_distance / 2) * (180/π) + + Args: + l2_distance: L2 distance between two unit ray directions (torch.Tensor, shape: (N,)) + Returns: + angular_error_degrees: Angular error in degrees (torch.Tensor, shape: (N,)) + """ + angular_error_radians = 2 * torch.asin(l2_distance / 2) + angular_error_degrees = angular_error_radians * 180.0 / math.pi + + return angular_error_degrees + + +def valid_mean(arr, mask, axis=None, keepdims=np._NoValue): + """Compute mean of elements across given dimensions of an array, considering only valid elements. + + Args: + arr: The array to compute the mean. + mask: Array with numerical or boolean values for element weights or validity. For bool, False means invalid. + axis: Dimensions to reduce. + keepdims: If true, retains reduced dimensions with length 1. + + Returns: + Mean array/scalar and a valid array/scalar that indicates where the mean could be computed successfully. + """ + + mask = mask.astype(arr.dtype) if mask.dtype == bool else mask + num_valid = np.sum(mask, axis=axis, keepdims=keepdims) + masked_arr = arr * mask + masked_arr_sum = np.sum(masked_arr, axis=axis, keepdims=keepdims) + + with np.errstate(divide="ignore", invalid="ignore"): + valid_mean = masked_arr_sum / num_valid + is_valid = np.isfinite(valid_mean) + valid_mean = np.nan_to_num(valid_mean, nan=0, posinf=0, neginf=0) + + return valid_mean, is_valid + + +def thresh_inliers(gt, pred, thresh=1.03, mask=None, output_scaling_factor=1.0): + """Computes the inlier (=error within a threshold) ratio for a predicted and ground truth dense map of size H x W x C. + + Args: + gt: Ground truth depth map as numpy array of shape HxW. Negative or 0 values are invalid and ignored. + pred: Predicted depth map as numpy array of shape HxW. + thresh: Threshold for the relative difference between the prediction and ground truth. Default: 1.03 + mask: Array of shape HxW with boolean values to indicate validity. For bool, False means invalid. Default: None + output_scaling_factor: Scaling factor that is applied after computing the metrics (e.g. to get [%]). Default: 1 + + Returns: + Scalar that indicates the inlier ratio. Scalar is np.nan if the result is invalid. + """ + # Compute the norms + gt_norm = np.linalg.norm(gt, axis=-1) + pred_norm = np.linalg.norm(pred, axis=-1) + + gt_norm_valid = (gt_norm) > 0 + if mask is not None: + combined_mask = mask & gt_norm_valid + else: + combined_mask = gt_norm_valid + + with np.errstate(divide="ignore", invalid="ignore"): + rel_1 = np.nan_to_num( + gt_norm / pred_norm, nan=thresh + 1, posinf=thresh + 1, neginf=thresh + 1 + ) # pred=0 should be an outlier + rel_2 = np.nan_to_num( + pred_norm / gt_norm, nan=0, posinf=0, neginf=0 + ) # gt=0 is masked out anyways + + max_rel = np.maximum(rel_1, rel_2) + inliers = ((0 < max_rel) & (max_rel < thresh)).astype( + np.float32 + ) # 1 for inliers, 0 for outliers + + inlier_ratio, valid = valid_mean(inliers, combined_mask) + + inlier_ratio = inlier_ratio * output_scaling_factor + inlier_ratio = inlier_ratio if valid else np.nan + + return inlier_ratio + + +def m_rel_ae(gt, pred, mask=None, output_scaling_factor=1.0): + """Computes the mean-relative-absolute-error for a predicted and ground truth dense map of size HxWxC. + + Args: + gt: Ground truth map as numpy array of shape H x W x C. + pred: Predicted map as numpy array of shape H x W x C. + mask: Array of shape HxW with boolean values to indicate validity. For bool, False means invalid. Default: None + output_scaling_factor: Scaling factor that is applied after computing the metrics (e.g. to get [%]). Default: 1 + + Returns: + Scalar that indicates the mean-relative-absolute-error. Scalar is np.nan if the result is invalid. + """ + error_norm = np.linalg.norm(pred - gt, axis=-1) + gt_norm = np.linalg.norm(gt, axis=-1) + + gt_norm_valid = (gt_norm) > 0 + if mask is not None: + combined_mask = mask & gt_norm_valid + else: + combined_mask = gt_norm_valid + + with np.errstate(divide="ignore", invalid="ignore"): + rel_ae = np.nan_to_num(error_norm / gt_norm, nan=0, posinf=0, neginf=0) + + m_rel_ae, valid = valid_mean(rel_ae, combined_mask) + + m_rel_ae = m_rel_ae * output_scaling_factor + m_rel_ae = m_rel_ae if valid else np.nan + + return m_rel_ae + + +def align(model, data): + """Align two trajectories using the method of Horn (closed-form). + + Args: + model -- first trajectory (3xn) + data -- second trajectory (3xn) + + Returns: + rot -- rotation matrix (3x3) + trans -- translation vector (3x1) + trans_error -- translational error per point (1xn) + + """ + np.set_printoptions(precision=3, suppress=True) + model_zerocentered = model - model.mean(1).reshape((3, -1)) + data_zerocentered = data - data.mean(1).reshape((3, -1)) + + W = np.zeros((3, 3)) + for column in range(model.shape[1]): + W += np.outer(model_zerocentered[:, column], data_zerocentered[:, column]) + U, d, Vh = np.linalg.linalg.svd(W.transpose()) + S = np.matrix(np.identity(3)) + if np.linalg.det(U) * np.linalg.det(Vh) < 0: + S[2, 2] = -1 + rot = U * S * Vh + trans = data.mean(1).reshape((3, -1)) - rot * model.mean(1).reshape((3, -1)) + + model_aligned = rot * model + trans + alignment_error = model_aligned - data + + trans_error = np.sqrt(np.sum(np.multiply(alignment_error, alignment_error), 0)).A[0] + + return rot, trans, trans_error + + +def evaluate_ate(gt_traj, est_traj): + """ + Input : + gt_traj: list of 4x4 matrices + est_traj: list of 4x4 matrices + len(gt_traj) == len(est_traj) + """ + gt_traj_pts = [gt_traj[idx][:3, 3] for idx in range(len(gt_traj))] + est_traj_pts = [est_traj[idx][:3, 3] for idx in range(len(est_traj))] + + gt_traj_pts = torch.stack(gt_traj_pts).detach().cpu().numpy().T + est_traj_pts = torch.stack(est_traj_pts).detach().cpu().numpy().T + + _, _, trans_error = align(gt_traj_pts, est_traj_pts) + + avg_trans_error = trans_error.mean() + + return avg_trans_error + + +def build_pair_index(N, B=1): + """ + Build indices for all possible pairs of frames. + + Args: + N: Number of frames + B: Batch size + + Returns: + i1, i2: Indices for all possible pairs + """ + i1_, i2_ = torch.combinations(torch.arange(N), 2, with_replacement=False).unbind(-1) + i1, i2 = [(i[None] + torch.arange(B)[:, None] * N).reshape(-1) for i in [i1_, i2_]] + return i1, i2 + + +def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: + """ + Returns torch.sqrt(torch.max(0, x)) + but with a zero subgradient where x is 0. + """ + ret = torch.zeros_like(x) + positive_mask = x > 0 + if torch.is_grad_enabled(): + ret[positive_mask] = torch.sqrt(x[positive_mask]) + else: + ret = torch.where(positive_mask, torch.sqrt(x), ret) + return ret + + +def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor: + """ + Convert rotations given as rotation matrices to quaternions. + + Args: + matrix: Rotation matrices as tensor of shape (..., 3, 3). + + Returns: + quaternions with real part last, as tensor of shape (..., 4). + Quaternion Order: XYZW or say ijkr, scalar-last + """ + if matrix.size(-1) != 3 or matrix.size(-2) != 3: + raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.") + + batch_dim = matrix.shape[:-2] + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + matrix.reshape(batch_dim + (9,)), dim=-1 + ) + + q_abs = _sqrt_positive_part( + torch.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + dim=-1, + ) + ) + + # we produce the desired quaternion multiplied by each of r, i, j, k + quat_by_rijk = torch.stack( + [ + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), + # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and + # `int`. + torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), + ], + dim=-2, + ) + + # We floor here at 0.1 but the exact level is not important; if q_abs is small, + # the candidate won't be picked. + flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr)) + + # if not for numerical problems, quat_candidates[i] should be same (up to a sign), + # forall i; we pick the best-conditioned one (with the largest denominator) + out = quat_candidates[ + F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : + ].reshape(batch_dim + (4,)) # pylint: disable=not-callable + + # Convert from rijk to ijkr + out = out[..., [1, 2, 3, 0]] + + out = standardize_quaternion(out) + + return out + + +def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor: + """ + Convert a unit quaternion to a standard form: one in which the real + part is non negative. + + Args: + quaternions: Quaternions with real part last, + as tensor of shape (..., 4). + + Returns: + Standardized quaternions as tensor of shape (..., 4). + """ + return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) + + +def rotation_angle(rot_gt, rot_pred, batch_size=None, eps=1e-15): + """ + Calculate rotation angle error between ground truth and predicted rotations. + + Args: + rot_gt: Ground truth rotation matrices + rot_pred: Predicted rotation matrices + batch_size: Batch size for reshaping the result + eps: Small value to avoid numerical issues + + Returns: + Rotation angle error in degrees + """ + q_pred = mat_to_quat(rot_pred) + q_gt = mat_to_quat(rot_gt) + + loss_q = (1 - (q_pred * q_gt).sum(dim=1) ** 2).clamp(min=eps) + err_q = torch.arccos(1 - 2 * loss_q) + + rel_rangle_deg = err_q * 180 / np.pi + + if batch_size is not None: + rel_rangle_deg = rel_rangle_deg.reshape(batch_size, -1) + + return rel_rangle_deg + + +def translation_angle(tvec_gt, tvec_pred, batch_size=None, ambiguity=True): + """ + Calculate translation angle error between ground truth and predicted translations. + + Args: + tvec_gt: Ground truth translation vectors + tvec_pred: Predicted translation vectors + batch_size: Batch size for reshaping the result + ambiguity: Whether to handle direction ambiguity + + Returns: + Translation angle error in degrees + """ + rel_tangle_deg = compare_translation_by_angle(tvec_gt, tvec_pred) + rel_tangle_deg = rel_tangle_deg * 180.0 / np.pi + + if ambiguity: + rel_tangle_deg = torch.min(rel_tangle_deg, (180 - rel_tangle_deg).abs()) + + if batch_size is not None: + rel_tangle_deg = rel_tangle_deg.reshape(batch_size, -1) + + return rel_tangle_deg + + +def compare_translation_by_angle(t_gt, t, eps=1e-15, default_err=1e6): + """ + Normalize the translation vectors and compute the angle between them. + + Args: + t_gt: Ground truth translation vectors + t: Predicted translation vectors + eps: Small value to avoid division by zero + default_err: Default error value for invalid cases + + Returns: + Angular error between translation vectors in radians + """ + t_norm = torch.norm(t, dim=1, keepdim=True) + t = t / (t_norm + eps) + + t_gt_norm = torch.norm(t_gt, dim=1, keepdim=True) + t_gt = t_gt / (t_gt_norm + eps) + + loss_t = torch.clamp_min(1.0 - torch.sum(t * t_gt, dim=1) ** 2, eps) + err_t = torch.acos(torch.sqrt(1 - loss_t)) + + err_t[torch.isnan(err_t) | torch.isinf(err_t)] = default_err + return err_t + + +def calculate_auc_np(r_error, t_error, max_threshold=30): + """ + Calculate the Area Under the Curve (AUC) for the given error arrays using NumPy. + + Args: + r_error: numpy array representing R error values (Degree) + t_error: numpy array representing T error values (Degree) + max_threshold: Maximum threshold value for binning the histogram + + Returns: + AUC value and the normalized histogram + """ + error_matrix = np.concatenate((r_error[:, None], t_error[:, None]), axis=1) + max_errors = np.max(error_matrix, axis=1) + bins = np.arange(max_threshold + 1) + histogram, _ = np.histogram(max_errors, bins=bins) + num_pairs = float(len(max_errors)) + normalized_histogram = histogram.astype(float) / num_pairs + return np.mean(np.cumsum(normalized_histogram)), normalized_histogram + + +def closed_form_inverse_se3(se3, R=None, T=None): + """ + Compute the inverse of each 4x4 (or 3x4) SE3 matrix in a batch. + + If `R` and `T` are provided, they must correspond to the rotation and translation + components of `se3`. Otherwise, they will be extracted from `se3`. + + Args: + se3: Nx4x4 or Nx3x4 array or tensor of SE3 matrices. + R (optional): Nx3x3 array or tensor of rotation matrices. + T (optional): Nx3x1 array or tensor of translation vectors. + + Returns: + Inverted SE3 matrices with the same type and device as `se3`. + + Shapes: + se3: (N, 4, 4) + R: (N, 3, 3) + T: (N, 3, 1) + """ + # Check if se3 is a numpy array or a torch tensor + is_numpy = isinstance(se3, np.ndarray) + + # Validate shapes + if se3.shape[-2:] != (4, 4) and se3.shape[-2:] != (3, 4): + raise ValueError(f"se3 must be of shape (N,4,4), got {se3.shape}.") + + # Extract R and T if not provided + if R is None: + R = se3[:, :3, :3] # (N,3,3) + if T is None: + T = se3[:, :3, 3:] # (N,3,1) + + # Transpose R + if is_numpy: + # Compute the transpose of the rotation for NumPy + R_transposed = np.transpose(R, (0, 2, 1)) + # -R^T t for NumPy + top_right = -np.matmul(R_transposed, T) + inverted_matrix = np.tile(np.eye(4), (len(R), 1, 1)) + else: + R_transposed = R.transpose(1, 2) # (N,3,3) + top_right = -torch.bmm(R_transposed, T) # (N,3,1) + inverted_matrix = torch.eye(4, 4)[None].repeat(len(R), 1, 1) + inverted_matrix = inverted_matrix.to(R.dtype).to(R.device) + + inverted_matrix[:, :3, :3] = R_transposed + inverted_matrix[:, :3, 3:] = top_right + + return inverted_matrix + + +def se3_to_relative_pose_error(pred_se3, gt_se3, num_frames): + """ + Compute rotation and translation errors between predicted and ground truth poses. + + Args: + pred_se3: Predicted SE(3) transformations + gt_se3: Ground truth SE(3) transformations + num_frames: Number of frames + + Returns: + Rotation and translation angle errors in degrees + """ + pair_idx_i1, pair_idx_i2 = build_pair_index(num_frames) + + # Compute relative camera poses between pairs + # We use closed_form_inverse to avoid potential numerical loss by torch.inverse() + relative_pose_gt = closed_form_inverse_se3(gt_se3[pair_idx_i1]).bmm( + gt_se3[pair_idx_i2] + ) + relative_pose_pred = closed_form_inverse_se3(pred_se3[pair_idx_i1]).bmm( + pred_se3[pair_idx_i2] + ) + + # Compute the difference in rotation and translation + rel_rangle_deg = rotation_angle( + relative_pose_gt[:, :3, :3], relative_pose_pred[:, :3, :3] + ) + rel_tangle_deg = translation_angle( + relative_pose_gt[:, :3, 3], relative_pose_pred[:, :3, 3] + ) + + return rel_rangle_deg, rel_tangle_deg diff --git a/mapanything/utils/misc.py b/mapanything/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..c6b976fa1df7430465961bc946d1744ae5665faa --- /dev/null +++ b/mapanything/utils/misc.py @@ -0,0 +1,109 @@ +""" +Miscellaneous utility functions. +""" + +import logging +import os +import random + +import numpy as np +import torch + + +class StreamToLogger: + """ + A class that redirects stream writes to a logger. + + This class can be used to redirect stdout or stderr to a logger + by implementing a file-like interface with write and flush methods. + + Parameters: + - logger: A logger instance that will receive the log messages + - log_level: The logging level to use (default: logging.INFO) + """ + + def __init__(self, logger, log_level=logging.INFO): + self.logger = logger + self.log_level = log_level + self.linebuf = "" + + def write(self, buf): + """ + Write the buffer content to the logger. + + Parameters: + - buf: The string buffer to write + """ + for line in buf.rstrip().splitlines(): + self.logger.log(self.log_level, line.rstrip()) + + def flush(self): + """ + Flush method to comply with file-like object interface. + This method is required but does nothing in this implementation. + """ + pass + + +def seed_everything(seed: int = 42): + """ + Set the `seed` value for torch and numpy seeds. Also turns on + deterministic execution for cudnn. + + Parameters: + - seed: A hashable seed value + """ + random.seed(seed) + os.environ["PYTHONHASHSEED"] = str(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + print(f"Seed set to: {seed}") + + +def invalid_to_nans(arr, valid_mask, ndim=999): + """ + Replace invalid values in an array with NaN values based on a validity mask. + + Parameters: + - arr: Input array (typically a PyTorch tensor) + - valid_mask: Boolean mask indicating valid elements (True) and invalid elements (False) + - ndim: Maximum number of dimensions to keep; flattens dimensions if arr.ndim > ndim + + Returns: + - Modified array with invalid values replaced by NaN + """ + if valid_mask is not None: + arr = arr.clone() + arr[~valid_mask] = float("nan") + if arr.ndim > ndim: + arr = arr.flatten(-2 - (arr.ndim - ndim), -2) + return arr + + +def invalid_to_zeros(arr, valid_mask, ndim=999): + """ + Replace invalid values in an array with zeros based on a validity mask. + + Parameters: + - arr: Input array (typically a PyTorch tensor) + - valid_mask: Boolean mask indicating valid elements (True) and invalid elements (False) + - ndim: Maximum number of dimensions to keep; flattens dimensions if arr.ndim > ndim + + Returns: + - Tuple containing: + - Modified array with invalid values replaced by zeros + - nnz: Number of non-zero (valid) elements per sample in the batch + """ + if valid_mask is not None: + arr = arr.clone() + arr[~valid_mask] = 0 + nnz = valid_mask.view(len(valid_mask), -1).sum(1) + else: + nnz = ( + arr[..., 0].numel() // len(arr) if len(arr) else 0 + ) # Number of pixels per image + if arr.ndim > ndim: + arr = arr.flatten(-2 - (arr.ndim - ndim), -2) + return arr, nnz diff --git a/mapanything/utils/moge_inference.py b/mapanything/utils/moge_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..a34c54ad0ba5cf73f5c15d40e9430dbd1cd0d72a --- /dev/null +++ b/mapanything/utils/moge_inference.py @@ -0,0 +1,111 @@ +""" +Util functions to run inference with MoGe +""" + +import argparse +import os +import warnings +from pathlib import Path + +warnings.filterwarnings("ignore", category=FutureWarning) # Suppress XFormers warnings + +import numpy as np +import rerun as rr +import torch +import torchvision +import torchvision.transforms as tvf +from PIL import Image + +from mapanything.utils.viz import log_data_to_rerun, script_add_rerun_args + + +def load_moge_model( + model_code_path: str = "MoGe", + ckpt_path: str = "Ruicheng/moge-vitl", #"/mnt/xri_mapsresearch/data/nkeetha/cache/huggingface/hub/models--Ruicheng--moge-vitl/snapshots/979e84da9415762c30e6c0cf8dc0962896c793df/model.pt", + device="cuda", +): + """ + Load the MoGe (ViT-L) model from huggingface hub (or load from local). + """ + if not Path(model_code_path).exists(): + raise FileNotFoundError(f"MoGe code not found at {model_code_path}") + import sys + + # Add the MoGe code to the system path + sys.path.append(str(model_code_path)) + + # Init the MoGe model + from moge.model.v1 import MoGeModel + + model = MoGeModel.from_pretrained(ckpt_path).to(device).eval() + + return model + + +@torch.no_grad() +def run_moge_inference(model: torch.nn.Module, image: torch.tensor, device="cuda"): + """ + Run MoGe inference on a batch of images or single image. + Output is a dictionary with the following keys: + - points: (B, H, W, 3) # scale-invariant point map in OpenCV camera coordinate system (x right, y down, z forward) + - depth: (B, H, W) # scale-invariant depth map + - mask: (B, H, W) # a binary mask for valid pixels + - intrinsics: (B, 3, 3) # normalized camera intrinsics + + Args: + model: MoGe model + image: (B, 3, H, W) or (3, H, W) # RGB image in range [0, 1] + """ + image = image.to(device) + return model.infer(image) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-ip", "--image_path", default='/ocean/projects/cis220039p/mdt2/jkarhade/Any4D/benchmarking/monst3r/demo_data/lady-running/00000.jpg', type=str) + parser.add_argument("--viz", action="store_true") + script_add_rerun_args(parser) # Options: --headless, --connect, --serve, --addr, --save, --stdout + args = parser.parse_args() + + # Setup Rerun if needed + if args.viz: + rr.script_setup(args, f"MoGe_Pred_Viz") + rr.set_time_seconds("stable_time", 0) + rr.log("moge", rr.ViewCoordinates.RDF, static=True) + + # Load the input data + img = np.array(Image.open(args.image_path)) # (H, W, 3) + transform = tvf.Compose([tvf.ToTensor()]) + input_img = transform(img).unsqueeze(0) # (B, 3, H, W) + + # Load the model + model = load_moge_model() + + # Run the model inference + output = run_moge_inference(model, input_img) + + # Get the different outputs + pts3d = output["points"].cpu().squeeze(0).numpy() # (H, W, 3) + depth = output["depth"].cpu().squeeze(0).numpy() # (H, W) + mask = output["mask"].cpu().squeeze(0).numpy() # (H, W) + intrinsics = output["intrinsics"].cpu().squeeze(0).numpy() # (3, 3), normalized + intrinsics[0, :] = intrinsics[0, :] * depth.shape[1] + intrinsics[1, :] = intrinsics[1, :] * depth.shape[0] + + # Log prediction to Rerun + if args.viz: + base_name = "moge" + log_data_to_rerun( + image=img, depthmap=depth, pose=np.eye(4), intrinsics=intrinsics, base_name=base_name, mask=np.float32(mask) + ) + # Log the predicted 3D points + filtered_pts = pts3d[mask] + filtered_pts_col = img[mask] + pts_name = f"{base_name}/points" + rr.log( + pts_name, + rr.Points3D( + positions=filtered_pts.reshape(-1, 3), + colors=filtered_pts_col.reshape(-1, 3), + ), + ) diff --git a/mapanything/utils/parallel.py b/mapanything/utils/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..e5c6f826291646a6c4059d3b6ae4c3d71d5e0104 --- /dev/null +++ b/mapanything/utils/parallel.py @@ -0,0 +1,158 @@ +""" +Utility functions for multiprocessing +""" + +import os +from multiprocessing.dummy import Pool as ThreadPool + +import torch +from torch.multiprocessing import Pool as TorchPool, set_start_method +from tqdm import tqdm + + +def cpu_count(): + """ + Returns the number of available CPUs for the python process + """ + return len(os.sched_getaffinity(0)) + + +def parallel_threads( + function, + args, + workers=0, + star_args=False, + kw_args=False, + front_num=1, + Pool=ThreadPool, + ordered_res=True, + **tqdm_kw, +): + """tqdm but with parallel execution. + + Will essentially return + res = [ function(arg) # default + function(*arg) # if star_args is True + function(**arg) # if kw_args is True + for arg in args] + + Note: + the first elements of args will not be parallelized. + This can be useful for debugging. + """ + # Determine the number of workers + while workers <= 0: + workers += cpu_count() + + # Convert args to an iterable + try: + n_args_parallel = len(args) - front_num + except TypeError: + n_args_parallel = None + args = iter(args) + + # Sequential execution for the first few elements (useful for debugging) + front = [] + while len(front) < front_num: + try: + a = next(args) + except StopIteration: + return front # end of the iterable + front.append( + function(*a) if star_args else function(**a) if kw_args else function(a) + ) + + # Parallel execution using multiprocessing.dummy + out = [] + with Pool(workers) as pool: + if star_args: + map_func = pool.imap if ordered_res else pool.imap_unordered + futures = map_func(starcall, [(function, a) for a in args]) + elif kw_args: + map_func = pool.imap if ordered_res else pool.imap_unordered + futures = map_func(starstarcall, [(function, a) for a in args]) + else: + map_func = pool.imap if ordered_res else pool.imap_unordered + futures = map_func(function, args) + # Track progress with tqdm + for f in tqdm(futures, total=n_args_parallel, **tqdm_kw): + out.append(f) + return front + out + + +def cuda_parallel_threads( + function, + args, + workers=0, + star_args=False, + kw_args=False, + front_num=1, + Pool=TorchPool, + ordered_res=True, + **tqdm_kw, +): + """ + Parallel execution of a function using torch.multiprocessing with CUDA support. + This is the CUDA variant of the parallel_threads function. + """ + # Set the start method for multiprocessing + set_start_method("spawn", force=True) + + # Determine the number of workers + while workers <= 0: + workers += torch.multiprocessing.cpu_count() + + # Convert args to an iterable + try: + n_args_parallel = len(args) - front_num + except TypeError: + n_args_parallel = None + args = iter(args) + + # Sequential execution for the first few elements (useful for debugging) + front = [] + while len(front) < front_num: + try: + a = next(args) + except StopIteration: + return front # End of the iterable + front.append( + function(*a) if star_args else function(**a) if kw_args else function(a) + ) + + # Parallel execution using torch.multiprocessing + out = [] + with Pool(workers) as pool: + if star_args: + map_func = pool.imap if ordered_res else pool.imap_unordered + futures = map_func(starcall, [(function, a) for a in args]) + elif kw_args: + map_func = pool.imap if ordered_res else pool.imap_unordered + futures = map_func(starstarcall, [(function, a) for a in args]) + else: + map_func = pool.imap if ordered_res else pool.imap_unordered + futures = map_func(function, args) + # Track progress with tqdm + for f in tqdm(futures, total=n_args_parallel, **tqdm_kw): + out.append(f) + return front + out + + +def parallel_processes(*args, **kwargs): + """Same as parallel_threads, with processes""" + import multiprocessing as mp + + kwargs["Pool"] = mp.Pool + return parallel_threads(*args, **kwargs) + + +def starcall(args): + """convenient wrapper for Process.Pool""" + function, args = args + return function(*args) + + +def starstarcall(args): + """convenient wrapper for Process.Pool""" + function, args = args + return function(**args) diff --git a/mapanything/utils/timing.py b/mapanything/utils/timing.py new file mode 100644 index 0000000000000000000000000000000000000000..ba9ee891d33420c51644f11f38bf169a1a610c27 --- /dev/null +++ b/mapanything/utils/timing.py @@ -0,0 +1,304 @@ +""" +Utility functions for timing code blocks +""" + +import time +from contextlib import ContextDecorator + +import numpy as np + + +class BlockTimeManager: + """ + Manages a collection of timers and their formatting options. + + This class serves as a central registry for Timer objects, allowing them to be + accessed by name and maintaining their formatting preferences. + + Attributes: + timers (dict): Dictionary mapping timer names to Timer objects + timer_fmts (dict): Dictionary mapping timer names to their display formats + window_size (int): Default window size for calculating windowed averages + buf_size (int): Default buffer size for storing timing measurements + """ + + def __init__(self, window_size=10, buf_size=100000): + self.timers = dict() + self.timer_fmts = dict() + self.window_size = window_size + self.buf_size = buf_size + + +btm = BlockTimeManager(window_size=100000) + + +class Timer: + """ + Core timing class that tracks execution times. + + This class provides the fundamental timing functionality, storing timing measurements + and calculating various statistics. + + Attributes: + name (str): Identifier for this timer + buf_size (int): Maximum number of timing measurements to store + window_size (int): Number of most recent measurements to use for windowed statistics + measures_arr (numpy.ndarray): Array storing start and end times of measurements + current_start (float or None): Start time of current measurement + current_end (float or None): End time of current measurement + """ + + def __init__(self, name, window_size, buf_size=100000): + self.name = name + self.buf_size = buf_size + self.window_size = window_size + self.init() + + def init(self): + """Initialize or reset the timer's state.""" + self.measures_arr = np.empty((0, 2)) # LIFO + self.current_start = None + self.current_end = None + + def reset(self): + """Reset the timer to its initial state.""" + self.init() + + def tic(self): + """Start a new timing measurement.""" + if self.current_start is not None: + # another tic executed before a toc + self.toc() + self.current_start = time.perf_counter() + + def toc(self): + """End the current timing measurement.""" + self.current_end = time.perf_counter() + self._add_current_measure() + + def _add_current_measure(self): + """Add the current timing measurement to the measurements array.""" + self.measures_arr = np.concatenate( + [ + np.array([[self.current_start, self.current_end]]), + self.measures_arr[: self.buf_size], + ] + ) + self.current_start = None + self.current_end = None + + @property + def avg(self) -> float: + """Calculate the average execution time across all measurements.""" + return np.mean(self.measures_arr[:, 1] - self.measures_arr[:, 0]) + + @property + def wavg(self) -> float: + """Calculate the windowed average execution time using the most recent measurements.""" + return np.mean( + self.measures_arr[: self.window_size, 1] + - self.measures_arr[: self.window_size, 0] + ) + + @property + def max(self) -> float: + """Return the maximum execution time.""" + return np.max(self.measures_arr[:, 1] - self.measures_arr[:, 0]) + + @property + def min(self) -> float: + """Return the minimum execution time.""" + return np.min(self.measures_arr[:, 1] - self.measures_arr[:, 0]) + + @property + def total(self) -> float: + """Return the total execution time across all measurements.""" + return np.sum(self.measures_arr[:, 1] - self.measures_arr[:, 0]) + + @property + def latest(self) -> float: + """Return the most recent execution time.""" + return self.measures_arr[0, 1] - self.measures_arr[0, 0] + + @property + def median(self) -> float: + """Return the median execution time.""" + return np.median(self.measures_arr[:, 1] - self.measures_arr[:, 0]) + + @property + def var(self) -> float: + """Return the variance of execution times.""" + return np.var(self.measures_arr[:, 1] - self.measures_arr[:, 0]) + + +class BlockTimer(ContextDecorator): + """ + A context manager and decorator for timing code blocks. + + This class provides a convenient interface for timing code execution, either as a + context manager (with statement) or as a decorator. It uses the Timer class for + the actual timing functionality. + + Attributes: + name (str): Identifier for this timer + fmt (str or None): Format string for displaying timing information + timer (Timer): The underlying Timer object + num_calls (int): Number of times this timer has been called + """ + + @staticmethod + def timers(): + """Return a list of all registered timer names.""" + return list(btm.timers.keys()) + + def __init__(self, name, fmt=None, window_size=100): + self.name = name + if name in btm.timers: + self.timer = btm.timers[name] + # restore format + self.fmt = fmt if fmt is not None else btm.timer_fmts[name] + else: + self.timer = Timer(name, btm.window_size, btm.buf_size) + btm.timers[name] = self.timer + btm.timer_fmts[name] = fmt + self.timer.window_size = window_size + self._default_fmt = "[{name}] num: {num} latest: {latest:.4f} --wind_avg: {wavg:.4f} -- avg: {avg:.4f} --var: {var:.4f} -- total: {total:.4f}" + if fmt == "default": + self.fmt = self._default_fmt + # extend here for new formats + else: + self.fmt = None + + self.num_calls = 0 + + def __enter__(self) -> "Timer": + """Start timing when entering a context.""" + self.tic() + return self + + def __exit__(self, *args): + """End timing when exiting a context and optionally display results.""" + self.toc() + if self.fmt is not None: + print(str(self)) + + def __str__(self) -> str: + """Return a string representation of the timer.""" + return self.display() + + def reset(self): + """Reset the timer and call counter.""" + self.timer.reset() + self.num_calls = 0 + + def display(self, fmt=None): + """ + Format and return timing information. + + Args: + fmt (str, optional): Format string to use. If None, uses the timer's format. + + Returns: + str: Formatted timing information + """ + if fmt is None: + if self.fmt is not None: + fmt = self.fmt + else: + fmt = self._default_fmt + return fmt.format( + name=self.name, + num=self.num_calls, + latest=self.latest, + wavg=self.wavg, + avg=self.avg, + var=self.var, + total=self.total, + ) + + def tic(self): + """Start a new timing measurement and increment the call counter.""" + self.timer.tic() + self.num_calls += 1 + + def toc(self, display=False): + """ + End the current timing measurement. + + Args: + display (bool): Whether to return a formatted display string + + Returns: + str or None: Formatted timing information if display is True + """ + self.timer.toc() + if display: + return self.display() + + @property + def latest(self) -> float: + """Return the most recent execution time.""" + return self.timer.latest + + @property + def avg(self) -> float: + """Return the average execution time.""" + return self.timer.avg + + @property + def wavg(self) -> float: + """Return the windowed average execution time.""" + return self.timer.wavg + + @property + def max(self) -> float: + """Return the maximum execution time.""" + return self.timer.max + + @property + def min(self) -> float: + """Return the minimum execution time.""" + return self.timer.min + + @property + def total(self) -> float: + """Return the total execution time.""" + return self.timer.total + + @property + def median(self) -> float: + """Return the median execution time.""" + return self.timer.median + + @property + def var(self) -> float: + """Return the variance of execution times.""" + return self.timer.var + + +if __name__ == "__main__": + + @BlockTimer("fct", "default") + def fct(bobo): + time.sleep(0.5) + + fct(2) + + for i in range(10): + with BlockTimer("affe", "default"): + time.sleep(0.1) + for i in range(1000): + with BlockTimer("test", None): + time.sleep(0.001) + + # BlockTimer("test").display = f"""avg: {BlockTimer("test").avg} total: {BlockTimer("test").total}""" + # print(str(BlockTimer("test"))) + + print(BlockTimer("test")) + BlockTimer("test").tic() + BlockTimer("t2", "default").tic() + time.sleep(0.4) + print(BlockTimer("t2").toc(True)) + + time.sleep(0.4) + print(BlockTimer("test").toc(True)) diff --git a/mapanything/utils/train_tools.py b/mapanything/utils/train_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..0f95d207f541a4d4de61f9111014d0927f91cd08 --- /dev/null +++ b/mapanything/utils/train_tools.py @@ -0,0 +1,978 @@ +""" +Utility functions for training deep learning models, particularly focused on distributed training, +metric logging, and gradient handling. + +This module provides tools for: +- Tracking and logging metrics during training +- Setting up distributed training environments +- Handling gradient scaling and normalization +- Managing learning rates and parameter groups +- Saving and loading model checkpoints + +References: CroCo (https://github.com/naver/croco) +""" + +import builtins +import datetime +import json +import math +import os +import time +from collections import defaultdict, deque +from pathlib import Path + +import torch +import torch.distributed as dist +from torch import inf + + +class SmoothedValue(object): + """ + Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda") + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + return self.total / self.count + + @property + def max(self): + return max(self.deque) + + @property + def value(self): + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +class MetricLogger(object): + """ + Logger for tracking and displaying training metrics. + + This class maintains a collection of metrics during training, provides + methods to update them, and formats them for display. It also handles + synchronization of metrics across processes in distributed training. + """ + + def __init__(self, delimiter="\t", print_per_view_stats=False): + """ + Initialize the MetricLogger. + + Args: + delimiter (str, optional): Delimiter for formatting output. Defaults to "\t". + print_per_view_stats (bool, optional): Whether to print per-view statistics. Defaults to False. + """ + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + self.print_per_view_stats = print_per_view_stats + + def update(self, **kwargs): + """ + Update metrics with new values. + + Args: + **kwargs: Key-value pairs where keys are metric names and values are metric values + Values can be tensors or numbers + + Raises: + AssertionError: If a value is not a float or int after conversion from tensor + """ + for k, v in kwargs.items(): + if v is None: + continue + if isinstance(v, torch.Tensor): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + """ + Get a meter by attribute name. + + This allows accessing meters as attributes of the logger. + + Args: + attr (str): Name of the attribute to get + + Returns: + SmoothedValue: The meter corresponding to the attribute name + + Raises: + AttributeError: If the attribute doesn't exist as a meter or regular attribute + """ + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError( + "'{}' object has no attribute '{}'".format(type(self).__name__, attr) + ) + + def __str__(self): + """ + Format all metrics as a string. + + Returns: + str: Formatted string containing all metrics + """ + loss_str = [] + for name, meter in self.meters.items(): + # Skip printing per-view stats if not enabled + if not self.print_per_view_stats and "view" in name: + continue + loss_str.append("{}: {}".format(name, str(meter))) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + """ + Synchronize metrics across processes in distributed training. + + This method calls synchronize_between_processes on each meter to + ensure consistent values across all processes. + """ + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + """ + Add a custom meter to the logger. + + Args: + name (str): Name of the meter + meter (SmoothedValue): The meter to add + """ + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None, max_iter=None): + """ + Log metrics at regular intervals while iterating. + + This method wraps an iterable and logs metrics every print_freq iterations. + It also tracks iteration time, data loading time, and memory usage. + + Args: + iterable: Iterable to iterate over (typically a data loader) + print_freq (int): How often to log metrics (in iterations) + header (str, optional): Header string to print before metrics. Defaults to None. + max_iter (int, optional): Maximum number of iterations. Defaults to None. + + Yields: + object: Items from the original iterable + """ + i = 0 + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + len_iterable = min(len(iterable), max_iter) if max_iter else len(iterable) + space_fmt = ":" + str(len(str(len_iterable))) + "d" + log_msg = [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + if torch.cuda.is_available(): + log_msg.append("max mem: {memory:.0f}") + log_msg = self.delimiter.join(log_msg) + MB = 1024.0 * 1024.0 + for it, obj in enumerate(iterable): + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == len_iterable - 1: + eta_seconds = iter_time.global_avg * (len_iterable - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print( + log_msg.format( + i, + len_iterable, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB, + ) + ) + else: + print( + log_msg.format( + i, + len_iterable, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + ) + ) + i += 1 + end = time.time() + if max_iter and it >= max_iter: + break + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print( + "{} Total time: {} ({:.4f} s / it)".format( + header, total_time_str, total_time / len_iterable + ) + ) + + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process. + + It replaces the built-in print function with a custom version that only prints + when the current process is the master process or when explicitly forced. + + Args: + is_master (bool): Whether the current process is the master process + """ + builtin_print = builtins.print + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + # force = force or (get_world_size() > 8) + if is_master or force: + now = datetime.datetime.now().time() + builtin_print("[{}] ".format(now), end="") # print with time stamp + builtin_print(*args, **kwargs) + + builtins.print = print + + +def is_dist_avail_and_initialized(): + """ + Check if distributed training is available and initialized. + + Returns: + bool: True if distributed training is available and initialized, False otherwise + """ + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + """ + Get the number of processes in the distributed training group. + + Returns: + int: Number of processes in the distributed group, or 1 if not using distributed training + """ + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + """ + Get the rank of the current process in the distributed training group. + + Returns: + int: Rank of the current process, or 0 if not using distributed training + """ + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + """ + Check if the current process is the main process (rank 0). + + Returns: + bool: True if the current process is the main process, False otherwise + """ + return get_rank() == 0 + + +def save_on_master(*args, **kwargs): + """ + Save a PyTorch object only on the master process. + + This function is useful in distributed training to avoid multiple processes + trying to save the same file simultaneously. + + Args: + *args: Positional arguments to pass to torch.save() + **kwargs: Keyword arguments to pass to torch.save() + """ + if is_main_process(): + torch.save(*args, **kwargs) + + +def init_distributed_mode(args): + """ + Initialize distributed training mode. + + This function sets up the distributed training environment based on environment + variables and command-line arguments. It initializes the process group, + sets the appropriate device, and configures printing for the distributed setup. + + Args: + args: Arguments object containing distributed training configuration. + Expected to have attributes like dist_url, and will be modified + to include rank, world_size, gpu, and distributed flag. + """ + nodist = args.nodist if hasattr(args, "nodist") else False + if "RANK" in os.environ and "WORLD_SIZE" in os.environ and not nodist: + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ["WORLD_SIZE"]) + args.gpu = int(os.environ["LOCAL_RANK"]) + else: + print("Not using distributed mode") + setup_for_distributed(is_master=True) # hack + args.distributed = False + return + + args.distributed = True + + torch.cuda.set_device(args.gpu) + args.dist_backend = "nccl" + print( + "| distributed init (rank {}): {}, gpu {}".format( + args.rank, args.dist_url, args.gpu + ), + flush=True, + ) + torch.distributed.init_process_group( + backend=args.dist_backend, + init_method=args.dist_url, + world_size=args.world_size, + rank=args.rank, + ) + torch.distributed.barrier() + setup_for_distributed(args.rank == 0) + + +class NativeScalerWithGradNormCount: + """ + A gradient scaler that handles gradient scaling and norm computation for mixed precision training. + + This class wraps PyTorch's GradScaler to provide additional functionality for gradient norm tracking + and clipping during mixed precision training. + """ + + state_dict_key = "amp_scaler" + + def __init__(self, enabled=True): + """Initialize the scaler. + + Args: + enabled (bool): Whether to enable gradient scaling. Default: True + """ + self._scaler = torch.GradScaler("cuda", enabled=enabled) + + def __call__( + self, + loss, + optimizer, + clip_grad=None, + parameters=None, + create_graph=False, + update_grad=True, + ): + """Scales loss and performs backward pass with optional gradient clipping. + + Args: + loss: The loss to backpropagate + optimizer: The optimizer being used + clip_grad: Max norm for gradient clipping. None means no clipping + parameters: Model parameters or list of parameters for gradient norm computation + create_graph: Whether to create graph during backward pass + update_grad: Whether to update gradients + + Returns: + norm: The gradient norm if computed, else None. Returns list of norms if parameters is a list. + """ + self._scaler.scale(loss).backward(create_graph=create_graph) + if update_grad: + if clip_grad is not None: + assert parameters is not None + self._scaler.unscale_( + optimizer + ) # unscale the gradients of optimizer's assigned params in-place + if isinstance(parameters, (list, tuple)): + norm = [ + torch.nn.utils.clip_grad_norm_(p, clip_grad) for p in parameters + ] + else: + norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad) + else: + self._scaler.unscale_(optimizer) + norm = get_grad_norm_(parameters) + self._scaler.step(optimizer) + self._scaler.update() + else: + norm = None + return norm + + def state_dict(self): + """Returns the state dict of the underlying scaler. + + Returns: + dict: The state dict of the gradient scaler + """ + return self._scaler.state_dict() + + def load_state_dict(self, state_dict): + """Loads the state dict into the underlying scaler. + + Args: + state_dict: The state dict to load + """ + self._scaler.load_state_dict(state_dict) + + +def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor: + """ + Calculate the gradient norm of parameters. + + This function computes the norm of gradients for a set of parameters. It can handle + both single parameter groups and multiple parameter groups (list/tuple of parameters). + + Args: + parameters: A tensor or iterable of tensors or iterable of iterables of tensors + containing model parameters for which to compute gradient norms + norm_type (float): Type of norm to use (e.g., 2.0 for L2 norm, inf for infinity norm) + + Returns: + torch.Tensor: The computed gradient norm. If parameters is a list/tuple of parameter + groups, returns a list of norms, one for each group. + """ + if isinstance(parameters, (list, tuple)): + # If parameters is already a list/tuple, process each parameter group + all_norms = [] + for params in parameters: + if isinstance(params, torch.Tensor): + params = [params] + params = [p for p in params if p.grad is not None] + if len(params) > 0: + device = params[0].grad.device + if norm_type == inf: + group_norm = max( + p.grad.detach().abs().max().to(device) for p in params + ) + else: + group_norm = torch.norm( + torch.stack( + [ + torch.norm(p.grad.detach(), norm_type).to(device) + for p in params + ] + ), + norm_type, + ) + else: + group_norm = torch.tensor(0.0) + all_norms.append(group_norm) + return all_norms + + # Original logic for single parameter group + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + parameters = [p for p in parameters if p.grad is not None] + norm_type = float(norm_type) + if len(parameters) == 0: + return torch.tensor(0.0) + device = parameters[0].grad.device + if norm_type == inf: + total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters) + else: + total_norm = torch.norm( + torch.stack( + [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters] + ), + norm_type, + ) + return total_norm + + +def save_model( + args, epoch, model_without_ddp, optimizer, loss_scaler, fname=None, best_so_far=None +): + """ + Save model checkpoint to disk. + + This function saves the model state, optimizer state, loss scaler state, + training arguments, current epoch, and optionally the best metric value so far. + The checkpoint is only saved on the master process in distributed training. + + Args: + args: Arguments containing output directory information + epoch (int): Current training epoch + model_without_ddp (torch.nn.Module): Model without DistributedDataParallel wrapper + optimizer (torch.optim.Optimizer): Optimizer instance + loss_scaler: Gradient scaler for mixed precision training + fname (str, optional): Custom filename suffix. If None, uses the epoch number. Defaults to None. + best_so_far (float, optional): Best metric value achieved so far. Defaults to None. + """ + output_dir = Path(args.output_dir) + if fname is None: + fname = str(epoch) + checkpoint_path = output_dir / ("checkpoint-%s.pth" % fname) + to_save = { + "model": model_without_ddp.state_dict(), + "optimizer": optimizer.state_dict(), + "scaler": loss_scaler.state_dict(), + "args": args, + "epoch": epoch, + } + if best_so_far is not None: + to_save["best_so_far"] = best_so_far + print(f">> Saving model to {checkpoint_path} ...") + save_on_master(to_save, checkpoint_path) + + +def load_model(train_args, model_without_ddp, optimizer, loss_scaler): + """ + Load model checkpoint from disk or URL. + + This function loads a saved checkpoint, restoring the model state, optimizer state, + loss scaler state, and training epoch. It can load from a local file or a URL. + + Args: + train_args: Training arguments containing resume information + model_without_ddp (torch.nn.Module): Model without DistributedDataParallel wrapper + optimizer (torch.optim.Optimizer): Optimizer instance + loss_scaler: Gradient scaler for mixed precision training + + Returns: + float or None: Best metric value from the checkpoint if available, otherwise None + """ + train_args.start_epoch = 0 + best_so_far = None + if train_args.resume and train_args.resume_ckpt is not None: + if train_args.resume_ckpt.startswith("https"): + checkpoint = torch.hub.load_state_dict_from_url( + train_args.resume_ckpt, map_location="cpu", check_hash=True + ) + else: + checkpoint = torch.load( + train_args.resume_ckpt, map_location="cpu", weights_only=False + ) + print("Resume checkpoint %s" % train_args.resume_ckpt) + model_without_ddp.load_state_dict(checkpoint["model"], strict=False) + train_args.start_epoch = checkpoint["epoch"] + 1 + optimizer.load_state_dict(checkpoint["optimizer"]) + if "scaler" in checkpoint: + loss_scaler.load_state_dict(checkpoint["scaler"]) + if "best_so_far" in checkpoint: + best_so_far = checkpoint["best_so_far"] + print(" & best_so_far={:g}".format(best_so_far)) + else: + print("") + print( + "With optim & sched! start_epoch={:d}".format(train_args.start_epoch), + end="", + ) + return best_so_far + + +def all_reduce_mean(x): + """ + Compute the mean of a value across all processes in distributed training. + + This function takes a value, reduces it across all processes using all_reduce, + and returns the mean value. + + Args: + x: The value to reduce (typically a scalar) + + Returns: + float: The mean value across all processes + """ + world_size = get_world_size() + if world_size > 1: + x_reduce = torch.tensor(x).cuda() + dist.all_reduce(x_reduce) + x_reduce /= world_size + return x_reduce.item() + else: + return x + + +def _replace(text, src, tgt, rm=""): + """ + Advanced string replacement utility. + + Given a text: + - replace all elements in src by the corresponding element in tgt + - remove all elements in rm + + Args: + text (str): The input text to modify + src (str): String of characters to replace + tgt (str): String of replacement characters (must be same length as src or length 1) + rm (str, optional): String of characters to remove. Defaults to "". + + Returns: + str: The modified text after replacements and removals + + Raises: + AssertionError: If src and tgt have different lengths (unless tgt has length 1) + """ + if len(tgt) == 1: + tgt = tgt * len(src) + assert len(src) == len(tgt), f"'{src}' and '{tgt}' should have the same len" + for s, t in zip(src, tgt): + text = text.replace(s, t) + for c in rm: + text = text.replace(c, "") + return text + + +def filename(obj): + """ + Transform a Python object or command into a proper filename. + + This function converts a Python object or command string into a valid filename + by replacing special characters and ensuring the filename is not too long. + + Special replacements: + - \1 gets replaced by slash '/' + - \2 gets replaced by comma ',' + + Args: + obj: The Python object or string to convert to a filename + + Returns: + str: A valid filename derived from the input object + + Raises: + AssertionError: If any part of the resulting path is longer than 256 characters + """ + if not isinstance(obj, str): + obj = repr(obj) + obj = str(obj).replace("()", "") + obj = _replace(obj, "_,(*/\1\2", "-__x%/,", rm=" )'\"") + assert all(len(s) < 256 for s in obj.split(os.sep)), ( + "filename too long (>256 characters):\n" + obj + ) + return obj + + +def compute_effective_lrs(train_args): + """ + Compute the effective learning rates based on batch size scaling. + + This function calculates the effective learning rates for the main model and + any submodules based on the effective batch size (accounting for gradient accumulation + and distributed training) and the base learning rates. + + Args: + train_args: Training arguments containing batch size, accumulation iterations, + learning rates, and submodule configurations + + Returns: + train_args: Updated training arguments with computed effective learning rates + """ + + # Compute the effective batch size + eff_batch_size = train_args.batch_size * train_args.accum_iter * get_world_size() + print("Accumulate grad iterations: %d" % train_args.accum_iter) + print("Effective batch size: %d" % eff_batch_size) + # Compute the effective default learning rate + if train_args.lr is None: # only base_lr is specified + train_args.lr = train_args.blr * math.sqrt( + eff_batch_size / train_args.base_eff_batch_size + ) + print( + f"Base default lr for effective batch size {eff_batch_size}: %.2e" + % (train_args.lr * math.sqrt(train_args.base_eff_batch_size / eff_batch_size)) + ) + print("Actual default lr: %.2e" % train_args.lr) + for submodule, config in train_args.submodule_configs.items(): + if config.get("lr") is None: # only base_lr is specified + config["lr"] = config["blr"] * math.sqrt( + eff_batch_size / train_args.base_eff_batch_size + ) + print( + f"Submodule {submodule} base lr for effective batch size {eff_batch_size}: %.2e" + % ( + config["lr"] + * math.sqrt(train_args.base_eff_batch_size / eff_batch_size) + ) + ) + print(f"Submodule {submodule} actual lr: %.2e" % config["lr"]) + + return train_args + + +def get_parameter_groups( + model, + lr, + weight_decay, + skip_list=[], + submodule_configs=None, + warn_not_in_submodule=False, +): + """ + Get parameter groups for optimizer with customized learning rates and weight decay. + + This function organizes model parameters into groups for the optimizer, allowing + different learning rates and weight decay values for different parts of the model. + Parameters are grouped by: + 1. Whether they should have weight decay applied (bias terms and 1D tensors typically don't) + 2. Which submodule they belong to (if submodule_configs is provided) + + Args: + model (torch.nn.Module): Model to get parameter groups for + lr (float): Default learning rate for parameters not in submodule_configs + weight_decay (float): Default weight decay for parameters not in submodule_configs + skip_list (list): List of parameter names to skip weight decay for + submodule_configs (dict, optional): Dictionary mapping submodule prefixes to configs + with 'lr' and 'weight_decay' keys + warn_not_in_submodule (bool, optional): Whether to warn if a parameter does not + belong to any submodule. Defaults to False. + + Returns: + tuple: A tuple containing: + - parameter_group_vars (list): List of parameter groups for optimizer + - parameter_group_name_to_idx_map (dict): Mapping from submodule name to parameter group indices + - parameter_group_idx_to_name_map (dict): Mapping from parameter group index to submodule name + """ + + if submodule_configs is None: + submodule_configs = {} + + parameter_group_names = {} + parameter_group_vars = {} + parameter_group_name_to_idx_map = {} + parameter_group_idx_to_name_map = {} + mapping_index = 0 + + for name, param in model.named_parameters(): + # Skip frozen parameters + if not param.requires_grad: + continue + + # Determine the submodule this parameter belongs to + submodule_name = None + for submodule, config in submodule_configs.items(): + if name.startswith(submodule): + submodule_name = submodule + break + + if submodule_name: + config = submodule_configs[submodule_name] + this_weight_decay = config.get("weight_decay", weight_decay) + this_lr = config.get("lr", lr) + # Freeze the parameters if lr is 0 + if this_lr == 0: + param.requires_grad = False + continue + else: + this_weight_decay = weight_decay + this_lr = lr + if warn_not_in_submodule and submodule_configs is not None: + print( + f"Warning: Parameter {name} does not belong to any submodule in {submodule_configs.keys()}." + ) + + # Assign weight decay values + if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list: + group_name = f"{submodule_name}_no_decay" if submodule_name else "no_decay" + this_weight_decay = 0.0 + else: + group_name = f"{submodule_name}_decay" if submodule_name else "decay" + + if group_name not in parameter_group_names: + parameter_group_names[group_name] = { + "weight_decay": this_weight_decay, + "lr": this_lr, + "params": [], + } + parameter_group_vars[group_name] = { + "weight_decay": this_weight_decay, + "lr": this_lr, + "params": [], + } + submodule_name_mapping = submodule_name if submodule_name else "default" + if submodule_name_mapping not in parameter_group_name_to_idx_map: + parameter_group_name_to_idx_map[submodule_name_mapping] = [ + mapping_index + ] + else: + parameter_group_name_to_idx_map[submodule_name_mapping].append( + mapping_index + ) + parameter_group_idx_to_name_map[mapping_index] = submodule_name_mapping + mapping_index += 1 + + parameter_group_vars[group_name]["params"].append(param) + parameter_group_names[group_name]["params"].append(name) + + # Print the parameter groups + print("Param groups = %s" % json.dumps(parameter_group_names, indent=2)) + + return ( + list(parameter_group_vars.values()), + parameter_group_name_to_idx_map, + parameter_group_idx_to_name_map, + ) + + +def adjust_learning_rate( + optimizer, + epoch, + train_args, + parameter_group_idx_to_name_map, + submodule_configs=None, +): + """ + Adjust the learning rate based on the schedule type and current epoch. + + This function updates the learning rates for all parameter groups in the optimizer + according to the specified learning rate schedule. Different submodules can have + different learning rate schedules. + + Currently supported schedule types: + - linear_warmup_half_cycle_cosine_decay: Linear warmup followed by cosine decay + + Args: + optimizer (torch.optim.Optimizer): The optimizer to update + epoch (int): Current training epoch + train_args: Training arguments containing schedule type, warmup epochs, etc. + parameter_group_idx_to_name_map (dict): Mapping from parameter group index to submodule name + submodule_configs (dict, optional): Dictionary of submodule-specific configurations + for learning rate schedules + + Raises: + ValueError: If an unsupported schedule type is specified + """ + + if submodule_configs is None: + submodule_configs = {} + + for group_num, param_group in enumerate(optimizer.param_groups): + submodule_name = parameter_group_idx_to_name_map.get(group_num) + + if submodule_name in submodule_configs: + config = submodule_configs[submodule_name] + lr = config.get("lr", train_args.lr) + warmup_epochs = config.get("warmup_epochs", train_args.warmup_epochs) + min_lr = config.get("min_lr", train_args.min_lr) + schedule_type = config.get("schedule_type", train_args.schedule_type) + else: + lr = train_args.lr + warmup_epochs = train_args.warmup_epochs + min_lr = train_args.min_lr + schedule_type = train_args.schedule_type + + if schedule_type == "linear_warmup_half_cycle_cosine_decay": + if epoch < warmup_epochs: + lr = lr * epoch / warmup_epochs + else: + lr = min_lr + (lr - min_lr) * 0.5 * ( + 1.0 + + math.cos( + math.pi + * (epoch - warmup_epochs) + / (train_args.epochs - warmup_epochs) + ) + ) + else: + raise ValueError(f"Schedule type {schedule_type} not implemented") + + param_group["lr"] = lr + + +def debug_after_backward( + model, + check_missing_gradients=True, + check_gradient_mismatch=False, + target_size=(256, 256, 1, 1), + target_stride=(256, 1, 256, 256), +): + """ + Debugging function to check for gradient issues after backward pass. + + This function performs two types of gradient debugging: + 1. Gradient mismatch: Checks for parameters with specific gradient shapes and strides + that might indicate incorrect gradient computation. + 2. Missing gradients: Identifies parameters that require gradients but didn't receive any. + + Args: + model (torch.nn.Module): The model to check gradients for + check_missing_gradients (bool, optional): Whether to check for missing gradients. Defaults to True. + check_gradient_mismatch (bool, optional): Whether to check for gradient mismatches. Defaults to False. + target_size (tuple, optional): Target tensor size to check for gradient mismatch. Defaults to (256, 256, 1, 1). + target_stride (tuple, optional): Target tensor stride to check for gradient mismatch. Defaults to (256, 1, 256, 256). + """ + # Debug for missing gradients + if check_missing_gradients: + missing_grad_params = [] + for name, param in model.named_parameters(): + if param.requires_grad and param.grad is None: + missing_grad_params.append(name) + + if missing_grad_params: + print("Parameters requiring gradients but missing gradients:") + for name in missing_grad_params: + print(f" - {name}") + else: + print("All parameters requiring gradients received gradients!") + + # Debug for gradient mismatch + if check_gradient_mismatch: + for name, param in model.named_parameters(): + grad = param.grad + if grad is None: + continue + if grad.size() == target_size and grad.stride() == target_stride: + print(f"Found parameter with incorrect gradient: '{name}'") + print(f"Gradient shape: {grad.size()}, strides: {grad.stride()}") diff --git a/mapanything/utils/viz.py b/mapanything/utils/viz.py new file mode 100644 index 0000000000000000000000000000000000000000..e738f08cd3ad700ac6cf01fd8f8a1d3361fce46c --- /dev/null +++ b/mapanything/utils/viz.py @@ -0,0 +1,167 @@ +""" +Utility functions for visualization +""" + +from argparse import ArgumentParser, Namespace +from distutils.util import strtobool + +import rerun as rr + + +def log_data_to_rerun(image, depthmap, pose, intrinsics, base_name, mask=None): + """ + Log camera and image data to Rerun visualization tool. + + Parameters + ---------- + image : numpy.ndarray + RGB image to be logged + depthmap : numpy.ndarray + Depth map corresponding to the image + pose : numpy.ndarray + 4x4 camera pose matrix with rotation (3x3) and translation (3x1) + intrinsics : numpy.ndarray + Camera intrinsic matrix + base_name : str + Base name for the logged entities in Rerun + mask : numpy.ndarray, optional + Optional segmentation mask for the depth image + """ + # Log camera info and loaded data + height, width = image.shape[0], image.shape[1] + rr.log( + base_name, + rr.Transform3D( + translation=pose[:3, 3], + mat3x3=pose[:3, :3], + ), + ) + rr.log( + f"{base_name}/pinhole", + rr.Pinhole( + image_from_camera=intrinsics, + height=height, + width=width, + camera_xyz=rr.ViewCoordinates.RDF, + ), + ) + rr.log( + f"{base_name}/pinhole/rgb", + rr.Image(image), + ) + rr.log( + f"{base_name}/pinhole/depth", + rr.DepthImage(depthmap), + ) + if mask is not None: + rr.log( + f"{base_name}/pinhole/depth_mask", + rr.SegmentationImage(mask), + ) + + +def str2bool(v): + return bool(strtobool(v)) + + +def script_add_rerun_args(parser: ArgumentParser) -> None: + """ + Add common Rerun script arguments to `parser`. + + Change Log from https://github.com/rerun-io/rerun/blob/29eb8954b08e59ff96943dc0677f46f7ea4ea734/rerun_py/rerun_sdk/rerun/script_helpers.py#L65: + - Added default portforwarding url for ease of use + - Update parser types + + Parameters + ---------- + parser : ArgumentParser + The parser to add arguments to. + + Returns + ------- + None + """ + parser.add_argument( + "--headless", + type=str2bool, + nargs="?", + const=True, + default=True, + help="Don't show GUI", + ) + parser.add_argument( + "--connect", + dest="connect", + type=str2bool, + nargs="?", + const=True, + default=True, + help="Connect to an external viewer", + ) + parser.add_argument( + "--serve", + dest="serve", + type=str2bool, + nargs="?", + const=True, + default=False, + help="Serve a web viewer (WARNING: experimental feature)", + ) + parser.add_argument( + "--url", + type=str, + default="rerun+http://127.0.0.1:2004/proxy", + help="Connect to this HTTP(S) URL", + ) + parser.add_argument( + "--save", type=str, default=None, help="Save data to a .rrd file at this path" + ) + parser.add_argument( + "-o", + "--stdout", + dest="stdout", + action="store_true", + help="Log data to standard output, to be piped into a Rerun Viewer", + ) + + +def init_rerun_args( + headless=True, + connect=True, + serve=False, + url="rerun+http://127.0.0.1:2004/proxy", + save=None, + stdout=False, +) -> Namespace: + """ + Initialize common Rerun script arguments. + + Parameters + ---------- + headless : bool, optional + Don't show GUI, by default True + connect : bool, optional + Connect to an external viewer, by default True + serve : bool, optional + Serve a web viewer (WARNING: experimental feature), by default False + url : str, optional + Connect to this HTTP(S) URL, by default rerun+http://127.0.0.1:2004/proxy + save : str, optional + Save data to a .rrd file at this path, by default None + stdout : bool, optional + Log data to standard output, to be piped into a Rerun Viewer, by default False + + Returns + ------- + Namespace + The parsed arguments. + """ + rerun_args = Namespace() + rerun_args.headless = headless + rerun_args.connect = connect + rerun_args.serve = serve + rerun_args.url = url + rerun_args.save = save + rerun_args.stdout = stdout + + return rerun_args diff --git a/mapanything/utils/warnings.py b/mapanything/utils/warnings.py new file mode 100644 index 0000000000000000000000000000000000000000..8422416bac8ba5893f6a50f2b32125f4f9ab65bb --- /dev/null +++ b/mapanything/utils/warnings.py @@ -0,0 +1,41 @@ +""" +Wrapper utilities for warnings. +""" + +import warnings +from functools import wraps + + +def suppress_traceback(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except Exception as e: + e.__traceback__ = e.__traceback__.tb_next.tb_next + raise + + return wrapper + + +class no_warnings: + def __init__(self, action: str = "ignore", **kwargs): + self.action = action + self.filter_kwargs = kwargs + + def __call__(self, fn): + @wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + warnings.simplefilter(self.action, **self.filter_kwargs) + return fn(*args, **kwargs) + + return wrapper + + def __enter__(self): + self.warnings_manager = warnings.catch_warnings() + self.warnings_manager.__enter__() + warnings.simplefilter(self.action, **self.filter_kwargs) + + def __exit__(self, exc_type, exc_val, exc_tb): + self.warnings_manager.__exit__(exc_type, exc_val, exc_tb) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..601c0c0c8977475334ba6dcd15d6677663107906 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,28 @@ +# Core dependencies +huggingface_hub +hydra-core +natsort +opencv-python-headless==4.10.0.84 +orjson +pillow-heif +plyfile +python-box +safetensors +tensorboard +tqdm +rerun-sdk>=0.23,<0.24 +torchaudio +torch~=2.6.0 +torchvision +trimesh +matplotlib +einops + +# Git dependencies +git+https://github.com/JayKarhade/UniCeption.git@dev/any4d +git+https://github.com/EasternJournalist/utils3d.git@3913c65d81e05e47b9f367250cf8c0f7462a0900 + +# Gradio dependencies +gradio==5.45.0 +gradio_rerun +spaces