Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/fig_teaser.png filter=lfs diff=lfs merge=lfs -text
+triton-2.0.0-cp310-cp310-win_amd64.whl filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,192 @@
+# Initially taken from Github's Python gitignore file
+
+ckpts
+sam_pt
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# tests and logs
+tests/fixtures/cached_*_text.txt
+logs/
+lightning_logs/
+lang_code_data/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# vscode
+.vs
+.vscode
+
+# Pycharm
+.idea
+
+# TF code
+tensorflow_code
+
+# Models
+proc_data
+
+# examples
+runs
+/runs_old
+/wandb
+/examples/runs
+/examples/**/*.args
+/examples/rag/sweep
+
+# data
+/data
+serialization_dir
+
+# emacs
+*.*~
+debug.env
+
+# vim
+.*.swp
+
+#ctags
+tags
+
+# pre-commit
+.pre-commit*
+
+# .lock
+*.lock
+
+# DS_Store (MacOS)
+.DS_Store
+# RL pipelines may produce mp4 outputs
+*.mp4
+
+# dependencies
+/transformers
+
+# ruff
+.ruff_cache
+
+# ckpts
+*.ckpt
+
+outputs/*
+
+NeuS/exp/*
+NeuS/test_scenes/*
+NeuS/mesh2tex/*
+neus_configs
+vast/*
+render_results
+experiments/*
+ckpts/*
+neus/*
+instant-nsr-pl/exp/*

1gpu.yaml

@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: 'NO'
+downcast_bf16: 'no'
+gpu_ids: '0'
+machine_rank: 0
+main_training_function: main
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

LICENSE

@@ -0,0 +1,661 @@
+                    GNU AFFERO GENERAL PUBLIC LICENSE
+                       Version 3, 19 November 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+                            Preamble
+
+  The GNU Affero General Public License is a free, copyleft license for
+software and other kinds of works, specifically designed to ensure
+cooperation with the community in the case of network server software.
+
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+our General Public Licenses are intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  Developers that use our General Public Licenses protect your rights
+with two steps: (1) assert copyright on the software, and (2) offer
+you this License which gives you legal permission to copy, distribute
+and/or modify the software.
+
+  A secondary benefit of defending all users' freedom is that
+improvements made in alternate versions of the program, if they
+receive widespread use, become available for other developers to
+incorporate.  Many developers of free software are heartened and
+encouraged by the resulting cooperation.  However, in the case of
+software used on network servers, this result may fail to come about.
+The GNU General Public License permits making a modified version and
+letting the public access it on a server without ever releasing its
+source code to the public.
+
+  The GNU Affero General Public License is designed specifically to
+ensure that, in such cases, the modified source code becomes available
+to the community.  It requires the operator of a network server to
+provide the source code of the modified version running there to the
+users of that server.  Therefore, public use of a modified version, on
+a publicly accessible server, gives the public access to the source
+code of the modified version.
+
+  An older license, called the Affero General Public License and
+published by Affero, was designed to accomplish similar goals.  This is
+a different license, not a version of the Affero GPL, but Affero has
+released a new version of the Affero GPL which permits relicensing under
+this license.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU Affero General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Remote Network Interaction; Use with the GNU General Public License.
+
+  Notwithstanding any other provision of this License, if you modify the
+Program, your modified version must prominently offer all users
+interacting with it remotely through a computer network (if your version
+supports such interaction) an opportunity to receive the Corresponding
+Source of your version by providing access to the Corresponding Source
+from a network server at no charge, through some standard or customary
+means of facilitating copying of software.  This Corresponding Source
+shall include the Corresponding Source for any work covered by version 3
+of the GNU General Public License that is incorporated pursuant to the
+following paragraph.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the work with which it is combined will remain governed by version
+3 of the GNU General Public License.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU Affero General Public License from time to time.  Such new versions
+will be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU Affero General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU Affero General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU Affero General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    <one line to give the program's name and a brief idea of what it does.>
+    Copyright (C) <year>  <name of author>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU Affero General Public License as published
+    by the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU Affero General Public License for more details.
+
+    You should have received a copy of the GNU Affero General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If your software can interact with users remotely through a computer
+network, you should also make sure that it provides a way for users to
+get its source.  For example, if your program is a web application, its
+interface could display a "Source" link that leads users to an archive
+of the code.  There are many ways you could offer source, and different
+solutions will be better for different programs; see section 13 for the
+specific requirements.
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU AGPL, see
+<https://www.gnu.org/licenses/>.

README.md

@@ -1,12 +1,211 @@
 ---
 title: Wonder3D
-emoji: 🏢
-colorFrom: blue
-colorTo: pink
+app_file: gradio_app_mv.py
 sdk: gradio
-sdk_version: 4.15.0
-app_file: app.py
-pinned: false
+sdk_version: 3.50.2
 ---
+**中文版本 [中文](README_zh.md)**
+# Wonder3D
+Single Image to 3D using Cross-Domain Diffusion
+## [Paper](https://arxiv.org/abs/2310.15008) | [Project page](https://www.xxlong.site/Wonder3D/) | [Hugging Face Demo](https://huggingface.co/spaces/flamehaze1115/Wonder3D-demo) | [Colab from @camenduru](https://github.com/camenduru/Wonder3D-colab)
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+![](assets/fig_teaser.png)
+
+Wonder3D reconstructs highly-detailed textured meshes from a single-view image in only 2 ∼ 3 minutes. Wonder3D first generates consistent multi-view normal maps with corresponding color images via a cross-domain diffusion model, and then leverages a novel normal fusion method to achieve fast and high-quality reconstruction.
+
+## Usage
+```bash
+
+# First clone the repo, and use the commands in the repo
+
+import torch
+import requests
+from PIL import Image
+import numpy as np
+from torchvision.utils import make_grid, save_image
+from diffusers import DiffusionPipeline  # only tested on diffusers[torch]==0.19.3, may have conflicts with newer versions of diffusers
+
+def load_wonder3d_pipeline():
+
+    pipeline = DiffusionPipeline.from_pretrained(
+    'flamehaze1115/wonder3d-v1.0', # or use local checkpoint './ckpts'
+    custom_pipeline='flamehaze1115/wonder3d-pipeline',
+    torch_dtype=torch.float16
+    )
+
+    # enable xformers
+    pipeline.unet.enable_xformers_memory_efficient_attention()
+
+    if torch.cuda.is_available():
+        pipeline.to('cuda:0')
+    return pipeline
+
+pipeline = load_wonder3d_pipeline()
+
+# Download an example image.
+cond = Image.open(requests.get("https://d.skis.ltd/nrp/sample-data/lysol.png", stream=True).raw)
+
+# The object should be located in the center and resized to 80% of image height.
+cond = Image.fromarray(np.array(cond)[:, :, :3])
+
+# Run the pipeline!
+images = pipeline(cond, num_inference_steps=20, output_type='pt', guidance_scale=1.0).images
+
+result = make_grid(images, nrow=6, ncol=2, padding=0, value_range=(0, 1))
+
+save_image(result, 'result.png')
+```
+
+## Collaborations
+Our overarching mission is to enhance the speed, affordability, and quality of 3D AIGC, making the creation of 3D content accessible to all. While significant progress has been achieved in the recent years, we acknowledge there is still a substantial journey ahead. We enthusiastically invite you to engage in discussions and explore potential collaborations in any capacity. <span style="color:red">**If you're interested in connecting or partnering with us, please don't hesitate to reach out via email (xxlong@connect.hku.hk)**</span> .
+
+## More features
+
+The repo is still being under construction, thanks for your patience. 
+- [x] Local gradio demo.
+- [x] Detailed tutorial.
+- [x] GUI demo for mesh reconstruction
+- [x] Windows support
+- [x] Docker support
+
+## Schedule
+- [x] Inference code and pretrained models.
+- [x] Huggingface demo.
+- [ ] New model with higher resolution.
+
+
+### Preparation for inference
+
+#### Linux System Setup.
+```angular2html
+conda create -n wonder3d
+conda activate wonder3d
+pip install -r requirements.txt
+pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
+```
+#### Windows System Setup.
+
+Please switch to branch `main-windows` to see details of windows setup.
+
+#### Docker Setup
+see [docker/README.MD](docker/README.md)
+
+### Inference
+1. Optional. If you have troubles to connect to huggingface. Make sure you have downloaded the following models.
+Download the [checkpoints](https://connecthkuhk-my.sharepoint.com/:f:/g/personal/xxlong_connect_hku_hk/Ej7fMT1PwXtKvsELTvDuzuMBebQXEkmf2IwhSjBWtKAJiA) and into the root folder.
+
+If you are in mainland China, you may download via [aliyun](https://www.alipan.com/s/T4rLUNAVq6V).
+
+```bash
+Wonder3D
+|-- ckpts
+    |-- unet
+    |-- scheduler
+    |-- vae
+    ...
+```
+Then modify the file ./configs/mvdiffusion-joint-ortho-6views.yaml, set `pretrained_model_name_or_path="./ckpts"`
+
+2. Download the [SAM](https://huggingface.co/spaces/abhishek/StableSAM/blob/main/sam_vit_h_4b8939.pth) model. Put it to the ``sam_pt`` folder.
+```
+Wonder3D
+|-- sam_pt
+    |-- sam_vit_h_4b8939.pth
+```
+3. Predict foreground mask as the alpha channel. We use [Clipdrop](https://clipdrop.co/remove-background) to segment the foreground object interactively. 
+You may also use `rembg` to remove the backgrounds.
+```bash
+# !pip install rembg
+import rembg
+result = rembg.remove(result)
+result.show()
+```
+4. Run Wonder3d to produce multiview-consistent normal maps and color images. Then you can check the results in the folder `./outputs`. (we use `rembg` to remove backgrounds of the results, but the segmentations are not always perfect. May consider using [Clipdrop](https://clipdrop.co/remove-background) to get masks for the generated normal maps and color images, since the quality of masks will significantly influence the reconstructed mesh quality.) 
+```bash
+accelerate launch --config_file 1gpu.yaml test_mvdiffusion_seq.py \
+            --config configs/mvdiffusion-joint-ortho-6views.yaml validation_dataset.root_dir={your_data_path} \
+            validation_dataset.filepaths=['your_img_file'] save_dir={your_save_path}
+```
+
+see example:
+
+```bash
+accelerate launch --config_file 1gpu.yaml test_mvdiffusion_seq.py \
+            --config configs/mvdiffusion-joint-ortho-6views.yaml validation_dataset.root_dir=./example_images \
+            validation_dataset.filepaths=['owl.png'] save_dir=./outputs
+```
+
+#### Interactive inference: run your local gradio demo. (Only generate normals and colors without reconstruction)
+```bash
+python gradio_app_mv.py   # generate multi-view normals and colors
+```
+
+5. Mesh Extraction
+
+#### Instant-NSR Mesh Extraction
+
+```bash
+cd ./instant-nsr-pl
+python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../{your_save_path}/cropsize-{crop_size}-cfg{guidance_scale:.1f}/ dataset.scene={scene}
+```
+
+see example:
+
+```bash
+cd ./instant-nsr-pl
+python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../outputs/cropsize-192-cfg1.0/ dataset.scene=owl
+```
+
+Our generated normals and color images are defined in orthographic views, so the reconstructed mesh is also in orthographic camera space. If you use MeshLab to view the meshes, you can click `Toggle Orthographic Camera` in `View` tab.
+
+#### Interactive inference: run your local gradio demo. (First generate normals and colors, and then do reconstructions. No need to perform gradio_app_mv.py first.)
+```bash
+python gradio_app_recon.py   
+```
+
+#### NeuS-based Mesh Extraction
+
+Since there are many complaints about the Windows setup of instant-nsr-pl, we provide the NeuS-based reconstruction, which may get rid of the requirement problems. 
+
+NeuS consumes less GPU memory and favors smooth surfaces without parameters tuning. However, NeuS consumes more times and its texture may be less sharp. If you are not sensitive to time, we recommend NeuS for optimization due to its robustness.
+
+```bash
+cd ./NeuS
+bash run.sh output_folder_path scene_name 
+```
+
+## Common questions
+Q: Tips to get better results.
+1. Wonder3D is sensitive the facing direciton of input images. By experiments, front-facing images always lead to good reconstruction.
+2. Limited by resources, current implemetation only supports limited views (6 views) and low resolution (256x256). Any images will be first resized into 256x256 for generation, so images after such a downsample that still keep clear and sharp features will lead to good results.
+3. Images with occlusions will cause worse reconstructions, since 6 views cannot cover the complete object. Images with less occlsuions lead to better results.
+4. Increate optimization steps in instant-nsr-pl, modify `trainer.max_steps: 3000` in `instant-nsr-pl/configs/neuralangelo-ortho-wmask.yaml` to more steps like `trainer.max_steps: 10000`. Longer optimization leads to better texture.
+
+Q: The evelation and azimuth degrees of the generated views?
+
+A: Unlike that the prior works such as Zero123, SyncDreamer and One2345 adopt object world system, our views are defined in the camera system of the input image. The six views are in the plane with 0 elevation degree in the camera system of the input image. Therefore we don't need to estimate an elevation degree for input image. The azimuth degrees of the six views are 0, 45, 90, 180, -90, -45 respectively.
+
+Q: The focal length of the generated views?
+
+A: We assume the input images are captured by orthographic camera, so the generated views are also in orthographic space. This design enables our model to keep strong generlaization on unreal images, but sometimes it may suffer from focal lens distortions on real-captured images.
+## Acknowledgement
+We have intensively borrow codes from the following repositories. Many thanks to the authors for sharing their codes.
+- [stable diffusion](https://github.com/CompVis/stable-diffusion)
+- [zero123](https://github.com/cvlab-columbia/zero123)
+- [NeuS](https://github.com/Totoro97/NeuS)
+- [SyncDreamer](https://github.com/liuyuan-pal/SyncDreamer)
+- [instant-nsr-pl](https://github.com/bennyguo/instant-nsr-pl)
+
+## License
+Wonder3D is under [AGPL-3.0](https://www.gnu.org/licenses/agpl-3.0.en.html), so any downstream solution and products (including cloud services) that include wonder3d code or a trained model (both pretrained or custom trained) inside it should be open-sourced to comply with the AGPL conditions. If you have any questions about the usage of Wonder3D, please contact us first.
+
+## Citation
+If you find this repository useful in your project, please cite the following work. :)
+```
+@article{long2023wonder3d,
+  title={Wonder3D: Single Image to 3D using Cross-Domain Diffusion},
+  author={Long, Xiaoxiao and Guo, Yuan-Chen and Lin, Cheng and Liu, Yuan and Dou, Zhiyang and Liu, Lingjie and Ma, Yuexin and Zhang, Song-Hai and Habermann, Marc and Theobalt, Christian and others},
+  journal={arXiv preprint arXiv:2310.15008},
+  year={2023}
+}
+```

README_zh.md

@@ -0,0 +1,203 @@
+**其他语言版本 [English](README.md)**
+
+# Wonder3D
+Single Image to 3D using Cross-Domain Diffusion
+## [Paper](https://arxiv.org/abs/2310.15008) | [Project page](https://www.xxlong.site/Wonder3D/) | [Hugging Face Demo](https://huggingface.co/spaces/flamehaze1115/Wonder3D-demo) | [Colab from @camenduru](https://github.com/camenduru/Wonder3D-colab)
+
+![](assets/fig_teaser.png)
+
+Wonder3D仅需2至3分钟即可从单视图图像中重建出高度详细的纹理网格。Wonder3D首先通过跨域扩散模型生成一致的多视图法线图与相应的彩色图像，然后利用一种新颖的法线融合方法实现快速且高质量的重建。
+
+## Usage 使用
+```bash
+
+import torch
+import requests
+from PIL import Image
+import numpy as np
+from torchvision.utils import make_grid, save_image
+from diffusers import DiffusionPipeline  # only tested on diffusers[torch]==0.19.3, may have conflicts with newer versions of diffusers
+
+def load_wonder3d_pipeline():
+
+    pipeline = DiffusionPipeline.from_pretrained(
+    'flamehaze1115/wonder3d-v1.0', # or use local checkpoint './ckpts'
+    custom_pipeline='flamehaze1115/wonder3d-pipeline',
+    torch_dtype=torch.float16
+    )
+
+    # enable xformers
+    pipeline.unet.enable_xformers_memory_efficient_attention()
+
+    if torch.cuda.is_available():
+        pipeline.to('cuda:0')
+    return pipeline
+
+pipeline = load_wonder3d_pipeline()
+
+# Download an example image.
+cond = Image.open(requests.get("https://d.skis.ltd/nrp/sample-data/lysol.png", stream=True).raw)
+
+# The object should be located in the center and resized to 80% of image height.
+cond = Image.fromarray(np.array(cond)[:, :, :3])
+
+# Run the pipeline!
+images = pipeline(cond, num_inference_steps=20, output_type='pt', guidance_scale=1.0).images
+
+result = make_grid(images, nrow=6, ncol=2, padding=0, value_range=(0, 1))
+
+save_image(result, 'result.png')
+```
+
+## Collaborations 合作
+我们的总体使命是提高3D人工智能图形生成（3D AIGC）的速度、可负担性和质量，使所有人都能够轻松创建3D内容。尽管近年来取得了显著的进展，我们承认前方仍有很长的路要走。我们热切邀请您参与讨论并在任何方面探索潜在的合作机会。<span style="color:red">**如果您有兴趣与我们联系或合作，请随时通过电子邮件（xxlong@connect.hku.hk）联系我们**</span>。
+
+## More features 
+
+The repo is still being under construction, thanks for your patience. 
+- [x] Local gradio demo.
+- [x] Detailed tutorial.
+- [x] GUI demo for mesh reconstruction
+- [x] Windows support
+- [x] Docker support
+
+## Schedule
+- [x] Inference code and pretrained models.
+- [x] Huggingface demo.
+- [ ] New model with higher resolution.
+
+
+### Preparation for inference 测试准备
+
+#### Linux System Setup.
+```angular2html
+conda create -n wonder3d
+conda activate wonder3d
+pip install -r requirements.txt
+pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
+```
+#### Windows System Setup.
+
+请切换到`main-windows`分支以查看Windows设置的详细信息。
+
+#### Docker Setup
+详见 [docker/README.MD](docker/README.md)
+
+### Inference
+1. 可选。如果您在连接到Hugging Face时遇到问题，请确保已下载以下模型。
+下载[checkpoints](https://connecthkuhk-my.sharepoint.com/:f:/g/personal/xxlong_connect_hku_hk/Ej7fMT1PwXtKvsELTvDuzuMBebQXEkmf2IwhSjBWtKAJiA)并放入根文件夹中。
+
+国内用户可下载： [阿里云盘](https://www.alipan.com/s/T4rLUNAVq6V)
+
+```bash
+Wonder3D
+|-- ckpts
+    |-- unet
+    |-- scheduler
+    |-- vae
+    ...
+```
+然后更改文件 ./configs/mvdiffusion-joint-ortho-6views.yaml, 设置 `pretrained_model_name_or_path="./ckpts"`
+
+2. 下载模型 [SAM](https://huggingface.co/spaces/abhishek/StableSAM/blob/main/sam_vit_h_4b8939.pth) . 放置在 ``sam_pt`` 文件夹.
+```
+Wonder3D
+|-- sam_pt
+    |-- sam_vit_h_4b8939.pth
+```
+3. 预测前景蒙版作为阿尔法通道。我们使用[Clipdrop](https://clipdrop.co/remove-background)来交互地分割前景对象。
+您还可以使用`rembg`来去除背景。
+```bash
+# !pip install rembg
+import rembg
+result = rembg.remove(result)
+result.show()
+```
+4. 运行Wonder3D以生成多视角一致的法线图和彩色图像。然后，您可以在文件夹`./outputs`中检查结果（我们使用`rembg`去除结果的背景，但分割并不总是完美的。可以考虑使用[Clipdrop](https://clipdrop.co/remove-background)获取生成的法线图和彩色图像的蒙版，因为蒙版的质量将显著影响重建的网格质量）。
+```bash
+accelerate launch --config_file 1gpu.yaml test_mvdiffusion_seq.py \
+            --config configs/mvdiffusion-joint-ortho-6views.yaml validation_dataset.root_dir={your_data_path} \
+            validation_dataset.filepaths=['your_img_file'] save_dir={your_save_path}
+```
+
+示例:
+
+```bash
+accelerate launch --config_file 1gpu.yaml test_mvdiffusion_seq.py \
+            --config configs/mvdiffusion-joint-ortho-6views.yaml validation_dataset.root_dir=./example_images \
+            validation_dataset.filepaths=['owl.png'] save_dir=./outputs
+```
+
+#### 运行本地的Gradio演示。仅生成法线和颜色，无需进行重建。
+```bash
+python gradio_app_mv.py   # generate multi-view normals and colors
+```
+
+5. Mesh Extraction
+
+#### Instant-NSR Mesh Extraction
+
+```bash
+cd ./instant-nsr-pl
+python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../{your_save_path}/cropsize-{crop_size}-cfg{guidance_scale:.1f}/ dataset.scene={scene}
+```
+
+示例:
+
+```bash
+cd ./instant-nsr-pl
+python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../outputs/cropsize-192-cfg1.0/ dataset.scene=owl
+```
+
+我们生成的法线图和彩色图像是在正交视图中定义的，因此重建的网格也处于正交摄像机空间。如果您使用MeshLab查看网格，可以在“View”选项卡中单击“Toggle Orthographic Camera”切换到正交相机。
+
+#### 运行本地的Gradio演示。首先生成法线和颜色，然后进行重建。无需首先执行`gradio_app_mv.py`。
+```bash
+python gradio_app_recon.py   
+```
+
+#### NeuS-based Mesh Extraction
+
+由于许多用户对于instant-nsr-pl的Windows设置提出了抱怨，我们提供了基于NeuS的重建，这可能消除了一些要求方面的问题。
+
+NeuS消耗较少的GPU内存，对平滑表面有利，无需参数调整。然而，NeuS需要更多时间，其纹理可能不够清晰。如果您对时间不太敏感，我们建议由于其稳健性而使用NeuS进行优化。
+
+```bash
+cd ./NeuS
+bash run.sh output_folder_path scene_name 
+```
+
+## 常见问题
+**获取更好结果的提示:**
+1. **图片朝向方向敏感:** Wonder3D对输入图像的面向方向敏感。通过实验证明，面向前方的图像通常会导致良好的重建结果。
+2. **图像分辨率:** 受资源限制，当前实现仅支持有限的视图（6个视图）和低分辨率（256x256）。任何图像都将首先调整大小为256x256进行生成，因此在这样的降采样后仍然保持清晰而锐利特征的图像将导致良好的结果。
+3. **处理遮挡:** 具有遮挡的图像会导致更差的重建，因为6个视图无法完全覆盖整个对象。具有较少遮挡的图像通常会产生更好的结果。
+4. **增加instant-nsr-pl中的优化步骤:** 在instant-nsr-pl中增加优化步骤。在`instant-nsr-pl/configs/neuralangelo-ortho-wmask.yaml`中修改`trainer.max_steps: 3000`为更多步骤，例如`trainer.max_steps: 10000`。更长的优化步骤会导致更好的纹理。
+
+**生成视图信息:**
+- **仰角和方位角度:** 与Zero123、SyncDreamer和One2345等先前作品采用对象世界系统不同，我们的视图是在输入图像的相机系统中定义的。六个视图在输入图像的相机系统中的平面上，仰角为0度。因此，我们不需要为输入图像估算仰角。六个视图的方位角度分别为0、45、90、180、-90、-45。
+
+**生成视图的焦距:**
+- 我们假设输入图像是由正交相机捕获的，因此生成的视图也在正交空间中。这种设计使得我们的模型能够在虚构图像上保持强大的泛化能力，但有时可能在实际捕获的图像上受到焦距镜头畸变的影响。
+
+## 致谢
+We have intensively borrow codes from the following repositories. Many thanks to the authors for sharing their codes.
+- [stable diffusion](https://github.com/CompVis/stable-diffusion)
+- [zero123](https://github.com/cvlab-columbia/zero123)
+- [NeuS](https://github.com/Totoro97/NeuS)
+- [SyncDreamer](https://github.com/liuyuan-pal/SyncDreamer)
+- [instant-nsr-pl](https://github.com/bennyguo/instant-nsr-pl)
+
+## 协议
+Wonder3D采用[AGPL-3.0](https://www.gnu.org/licenses/agpl-3.0.en.html)许可，因此任何包含Wonder3D代码或其中训练的模型（无论是预训练还是定制训练）的下游解决方案和产品（包括云服务）都应该开源以符合AGPL条件。如果您对Wonder3D的使用有任何疑问，请首先与我们联系。
+
+## 引用
+如果您在项目中发现这个项目对您有用，请引用以下工作。 :)
+```
+@article{long2023wonder3d,
+  title={Wonder3D: Single Image to 3D using Cross-Domain Diffusion},
+  author={Long, Xiaoxiao and Guo, Yuan-Chen and Lin, Cheng and Liu, Yuan and Dou, Zhiyang and Liu, Lingjie and Ma, Yuexin and Zhang, Song-Hai and Habermann, Marc and Theobalt, Christian and others},
+  journal={arXiv preprint arXiv:2310.15008},
+  year={2023}
+}
+```

configs/mvdiffusion-joint-ortho-6views.yaml

@@ -0,0 +1,42 @@
+pretrained_model_name_or_path: 'flamehaze1115/wonder3d-v1.0' # or './ckpts'
+revision: null
+validation_dataset:
+  root_dir: "./example_images" # the folder path stores testing images
+  num_views: 6
+  bg_color: 'white'
+  img_wh: [256, 256]
+  num_validation_samples: 1000
+  crop_size: 192
+  filepaths: ['owl.png']  # the test image names. leave it empty, test all images in the folder
+
+save_dir: 'outputs/'
+
+pred_type: 'joint'
+seed: 42
+validation_batch_size: 1
+dataloader_num_workers: 64
+
+local_rank: -1
+
+pipe_kwargs:
+  camera_embedding_type: 'e_de_da_sincos'
+  num_views: 6
+
+validation_guidance_scales: [1.0]
+pipe_validation_kwargs:
+  eta: 1.0
+validation_grid_nrow: 6
+
+unet_from_pretrained_kwargs:
+  camera_embedding_type: 'e_de_da_sincos'
+  projection_class_embeddings_input_dim: 10  
+  num_views: 6
+  sample_size: 32
+  cd_attention_mid: true
+  zero_init_conv_in: false
+  zero_init_camera_projection: false  
+
+num_views: 6
+camera_embedding_type: 'e_de_da_sincos'
+
+enable_xformers_memory_efficient_attention: true

docker/Dockerfile

@@ -0,0 +1,56 @@
+# get the development image from nvidia cuda 11.7
+FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04
+
+LABEL name="Wonder3D" \
+    maintainer="Tiancheng <athinkingneal@gmail.com>" \
+    lastupdate="2024-01-05" 
+
+# create workspace folder and set it as working directory
+RUN mkdir -p /workspace
+WORKDIR /workspace
+
+# Set the timezone
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && \
+    apt-get install -y tzdata && \
+    ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && \
+    dpkg-reconfigure --frontend noninteractive tzdata
+
+# update package lists and install git, wget, vim, libgl1-mesa-glx, and libglib2.0-0
+RUN apt-get update && \
+    apt-get install -y git wget vim libgl1-mesa-glx libglib2.0-0 unzip
+
+# install conda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    chmod +x Miniconda3-latest-Linux-x86_64.sh && \
+    ./Miniconda3-latest-Linux-x86_64.sh -b -p /workspace/miniconda3 && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+# update PATH environment variable
+ENV PATH="/workspace/miniconda3/bin:${PATH}"
+
+# initialize conda
+RUN conda init bash
+
+# create and activate conda environment
+RUN conda create -n wonder3d python=3.8 && echo "source activate wonder3d" > ~/.bashrc
+ENV PATH /workspace/miniconda3/envs/wonder3d/bin:$PATH
+
+
+# clone the repository
+RUN git clone https://github.com/xxlong0/Wonder3D.git && \
+    cd /workspace/Wonder3D 
+
+# change the working directory to the repository
+WORKDIR /workspace/Wonder3D
+
+# install pytorch 1.13.1 and torchvision
+RUN pip install -r docker/requirements.txt
+
+# install the specific version of nerfacc corresponding to torch 1.13.0 and cuda 11.7, otherwise the nerfacc will freeze during cuda setup
+RUN pip install nerfacc==0.3.3 -f https://nerfacc-bucket.s3.us-west-2.amazonaws.com/whl/torch-1.13.0_cu117.html
+
+# install tiny cuda during docker setup will cause error, need to install it manually in the container
+# RUN pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
+
+

docker/README.md

@@ -0,0 +1,57 @@
+# Docker setup
+
+This docker setup is tested on Ubunu20.04.
+
+make sure you are under directory yourworkspace/Wonder3D/
+
+run
+
+`docker build --no-cache -t wonder3d/deploy:cuda11.7 -f docker/Dockerfile .`
+
+then run 
+
+`docker run --gpus all -it wonder3d/deploy:cuda11.7 bash`
+
+
+## Nvidia Container Toolkit setup
+
+You will have trouble enabling gpu for docker if you haven't installed **NVIDIA Container Toolkit** on you local machine before. You can skip this section if you have already installed it. Follow the instruction in this website to install it. 
+
+https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html
+
+or you can run the following command to install it with apt:
+
+1.Configure the production repository:
+   
+```bash
+curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
+  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
+    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
+    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
+```
+
+2.Update the packages list from the repository:
+
+`sed -i -e '/experimental/ s/^#//g' /etc/apt/sources.list.d/nvidia-container-toolkit.list`
+
+3.Install the NVIDIA Container Toolkit packages:
+
+`sudo apt-get install -y nvidia-container-toolkit`
+
+Remember to restart the docker:
+
+`sudo systemctl restart docker`
+
+now you can run the following command:
+
+`docker run --gpus all -it wonder3d/deploy:cuda11.7 bash`
+
+
+## Install Tiny Cudann
+
+After you start the container, run the following command to install tiny cudann. Somehow this pip installation can not be done during the docker build, so you have to do it manually after the docker is started.
+
+`pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch`
+
+
+Now you should be good to go, good luck and have fun :)

docker/requirements.txt

@@ -0,0 +1,36 @@
+--extra-index-url https://download.pytorch.org/whl/cu117
+
+# nerfacc==0.3.3, nefacc needs to be installed from the specific location
+# see installation part in this link: https://github.com/nerfstudio-project/nerfacc
+
+torch==1.13.1+cu117 
+torchvision==0.14.1+cu117
+diffusers[torch]==0.19.3
+xformers==0.0.16
+transformers>=4.25.1
+bitsandbytes==0.35.4
+decord==0.6.0
+pytorch-lightning<2
+omegaconf==2.2.3
+trimesh==3.9.8
+pyhocon==0.3.57
+icecream==2.1.0
+PyMCubes==0.1.2
+accelerate
+modelcards
+einops
+ftfy
+piq
+matplotlib
+opencv-python
+imageio
+imageio-ffmpeg
+scipy
+pyransac3d
+torch_efficient_distloss
+tensorboard
+rembg
+segment_anything
+gradio==3.50.2
+triton
+rich

gradio_app_mv.py

@@ -0,0 +1,439 @@
+import os
+import torch
+import fire
+import gradio as gr
+from PIL import Image
+from functools import partial
+
+import cv2
+import time
+import numpy as np
+from rembg import remove
+from segment_anything import sam_model_registry, SamPredictor
+
+import os
+import sys
+import numpy
+import torch
+import rembg
+import threading
+import urllib.request
+from PIL import Image
+from typing import Dict, Optional, Tuple, List
+from dataclasses import dataclass
+import streamlit as st
+import huggingface_hub
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+from mvdiffusion.models.unet_mv2d_condition import UNetMV2DConditionModel
+from mvdiffusion.data.single_image_dataset import SingleImageDataset as MVDiffusionDataset
+from mvdiffusion.pipelines.pipeline_mvdiffusion_image import MVDiffusionImagePipeline
+from diffusers import AutoencoderKL, DDPMScheduler, DDIMScheduler
+from einops import rearrange
+import numpy as np
+import subprocess
+from datetime import datetime
+
+def save_image(tensor):
+    ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    # pdb.set_trace()
+    im = Image.fromarray(ndarr)
+    return ndarr
+
+
+def save_image_to_disk(tensor, fp):
+    ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    # pdb.set_trace()
+    im = Image.fromarray(ndarr)
+    im.save(fp)
+    return ndarr
+
+
+def save_image_numpy(ndarr, fp):
+    im = Image.fromarray(ndarr)
+    im.save(fp)
+
+
+weight_dtype = torch.float16
+
+_TITLE = '''Wonder3D: Single Image to 3D using Cross-Domain Diffusion'''
+_DESCRIPTION = '''
+<div>
+Generate consistent multi-view normals maps and color images.
+<a style="display:inline-block; margin-left: .5em" href='https://github.com/xxlong0/Wonder3D/'><img src='https://img.shields.io/github/stars/xxlong0/Wonder3D?style=social' /></a>
+</div>
+<div>
+The demo does not include the mesh reconstruction part, please visit <a href="https://github.com/xxlong0/Wonder3D/">our github repo</a> to get a textured mesh.
+</div>
+'''
+_GPU_ID = 0
+
+
+if not hasattr(Image, 'Resampling'):
+    Image.Resampling = Image
+
+
+def sam_init():
+    sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
+    model_type = "vit_h"
+
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda:{_GPU_ID}")
+    predictor = SamPredictor(sam)
+    return predictor
+
+
+def sam_segment(predictor, input_image, *bbox_coords):
+    bbox = np.array(bbox_coords)
+    image = np.asarray(input_image)
+
+    start_time = time.time()
+    predictor.set_image(image)
+
+    masks_bbox, scores_bbox, logits_bbox = predictor.predict(box=bbox, multimask_output=True)
+
+    print(f"SAM Time: {time.time() - start_time:.3f}s")
+    out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
+    out_image[:, :, :3] = image
+    out_image_bbox = out_image.copy()
+    out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255
+    torch.cuda.empty_cache()
+    return Image.fromarray(out_image_bbox, mode='RGBA')
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def preprocess(predictor, input_image, chk_group=None, segment=True, rescale=False):
+    RES = 1024
+    input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
+    if chk_group is not None:
+        segment = "Background Removal" in chk_group
+        rescale = "Rescale" in chk_group
+    if segment:
+        image_rem = input_image.convert('RGBA')
+        image_nobg = remove(image_rem, alpha_matting=True)
+        arr = np.asarray(image_nobg)[:, :, -1]
+        x_nonzero = np.nonzero(arr.sum(axis=0))
+        y_nonzero = np.nonzero(arr.sum(axis=1))
+        x_min = int(x_nonzero[0].min())
+        y_min = int(y_nonzero[0].min())
+        x_max = int(x_nonzero[0].max())
+        y_max = int(y_nonzero[0].max())
+        input_image = sam_segment(predictor, input_image.convert('RGB'), x_min, y_min, x_max, y_max)
+    # Rescale and recenter
+    if rescale:
+        image_arr = np.array(input_image)
+        in_w, in_h = image_arr.shape[:2]
+        out_res = min(RES, max(in_w, in_h))
+        ret, mask = cv2.threshold(np.array(input_image.split()[-1]), 0, 255, cv2.THRESH_BINARY)
+        x, y, w, h = cv2.boundingRect(mask)
+        max_size = max(w, h)
+        ratio = 0.75
+        side_len = int(max_size / ratio)
+        padded_image = np.zeros((side_len, side_len, 4), dtype=np.uint8)
+        center = side_len // 2
+        padded_image[center - h // 2 : center - h // 2 + h, center - w // 2 : center - w // 2 + w] = image_arr[y : y + h, x : x + w]
+        rgba = Image.fromarray(padded_image).resize((out_res, out_res), Image.LANCZOS)
+
+        rgba_arr = np.array(rgba) / 255.0
+        rgb = rgba_arr[..., :3] * rgba_arr[..., -1:] + (1 - rgba_arr[..., -1:])
+        input_image = Image.fromarray((rgb * 255).astype(np.uint8))
+    else:
+        input_image = expand2square(input_image, (127, 127, 127, 0))
+    return input_image, input_image.resize((320, 320), Image.Resampling.LANCZOS)
+
+
+def load_wonder3d_pipeline(cfg):
+
+    pipeline = MVDiffusionImagePipeline.from_pretrained(
+    cfg.pretrained_model_name_or_path,
+    torch_dtype=weight_dtype
+    )
+
+    # pipeline.to('cuda:0')
+    pipeline.unet.enable_xformers_memory_efficient_attention()
+
+
+    if torch.cuda.is_available():
+        pipeline.to('cuda:0')
+    # sys.main_lock = threading.Lock()
+    return pipeline
+
+
+from mvdiffusion.data.single_image_dataset import SingleImageDataset
+
+
+def prepare_data(single_image, crop_size):
+    dataset = SingleImageDataset(root_dir='', num_views=6, img_wh=[256, 256], bg_color='white', crop_size=crop_size, single_image=single_image)
+    return dataset[0]
+
+scene = 'scene'
+
+def run_pipeline(pipeline, cfg, single_image, guidance_scale, steps, seed, crop_size, chk_group=None):
+    import pdb
+    global scene
+    # pdb.set_trace()
+
+    if chk_group is not None:
+        write_image = "Write Results" in chk_group
+
+    batch = prepare_data(single_image, crop_size)
+
+    pipeline.set_progress_bar_config(disable=True)
+    seed = int(seed)
+    generator = torch.Generator(device=pipeline.unet.device).manual_seed(seed)
+
+    # repeat  (2B, Nv, 3, H, W)
+    imgs_in = torch.cat([batch['imgs_in']] * 2, dim=0).to(weight_dtype)
+
+    # (2B, Nv, Nce)
+    camera_embeddings = torch.cat([batch['camera_embeddings']] * 2, dim=0).to(weight_dtype)
+
+    task_embeddings = torch.cat([batch['normal_task_embeddings'], batch['color_task_embeddings']], dim=0).to(weight_dtype)
+
+    camera_embeddings = torch.cat([camera_embeddings, task_embeddings], dim=-1).to(weight_dtype)
+
+    # (B*Nv, 3, H, W)
+    imgs_in = rearrange(imgs_in, "Nv C H W -> (Nv) C H W")
+    # (B*Nv, Nce)
+    # camera_embeddings = rearrange(camera_embeddings, "B Nv Nce -> (B Nv) Nce")
+
+    out = pipeline(
+        imgs_in,
+        # camera_embeddings,
+        generator=generator,
+        guidance_scale=guidance_scale,
+        num_inference_steps=steps,
+        output_type='pt',
+        num_images_per_prompt=1,
+        **cfg.pipe_validation_kwargs,
+    ).images
+
+    bsz = out.shape[0] // 2
+    normals_pred = out[:bsz]
+    images_pred = out[bsz:]
+    num_views = 6
+    if write_image:
+        VIEWS = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
+        cur_dir = os.path.join("./outputs", f"cropsize-{int(crop_size)}-cfg{guidance_scale:.1f}")
+
+        scene = 'scene'+datetime.now().strftime('@%Y%m%d-%H%M%S')
+        scene_dir = os.path.join(cur_dir, scene)
+        normal_dir = os.path.join(scene_dir, "normals")
+        masked_colors_dir = os.path.join(scene_dir, "masked_colors")
+        os.makedirs(normal_dir, exist_ok=True)
+        os.makedirs(masked_colors_dir, exist_ok=True)
+        for j in range(num_views):
+            view = VIEWS[j]
+            normal = normals_pred[j]
+            color = images_pred[j]
+
+            normal_filename = f"normals_000_{view}.png"
+            rgb_filename = f"rgb_000_{view}.png"
+            normal = save_image_to_disk(normal, os.path.join(normal_dir, normal_filename))
+            color = save_image_to_disk(color, os.path.join(scene_dir, rgb_filename))
+
+            # rm_normal = remove(normal)
+            # rm_color = remove(color)
+
+            # save_image_numpy(rm_normal, os.path.join(scene_dir, normal_filename))
+            # save_image_numpy(rm_color, os.path.join(masked_colors_dir, rgb_filename))
+
+    normals_pred = [save_image(normals_pred[i]) for i in range(bsz)]
+    images_pred = [save_image(images_pred[i]) for i in range(bsz)]
+
+    out = images_pred + normals_pred
+    return out
+
+
+def process_3d(mode, data_dir, guidance_scale, crop_size):
+    dir = None
+    global scene
+
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+    subprocess.run(
+        f'cd instant-nsr-pl && python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../{data_dir}/cropsize-{crop_size:.1f}-cfg{guidance_scale:.1f}/ dataset.scene={scene} && cd ..',
+        shell=True,
+    )
+    import glob
+    # import pdb
+
+    # pdb.set_trace()
+
+    obj_files = glob.glob(f'{cur_dir}/instant-nsr-pl/exp/{scene}/*/save/*.obj', recursive=True)
+    print(obj_files)
+    if obj_files:
+        dir = obj_files[0]
+    return dir
+
+
+@dataclass
+class TestConfig:
+    pretrained_model_name_or_path: str
+    pretrained_unet_path: str
+    revision: Optional[str]
+    validation_dataset: Dict
+    save_dir: str
+    seed: Optional[int]
+    validation_batch_size: int
+    dataloader_num_workers: int
+
+    local_rank: int
+
+    pipe_kwargs: Dict
+    pipe_validation_kwargs: Dict
+    unet_from_pretrained_kwargs: Dict
+    validation_guidance_scales: List[float]
+    validation_grid_nrow: int
+    camera_embedding_lr_mult: float
+
+    num_views: int
+    camera_embedding_type: str
+
+    pred_type: str  # joint, or ablation
+
+    enable_xformers_memory_efficient_attention: bool
+
+    cond_on_normals: bool
+    cond_on_colors: bool
+
+
+def run_demo():
+    from utils.misc import load_config
+    from omegaconf import OmegaConf
+
+    # parse YAML config to OmegaConf
+    cfg = load_config("./configs/mvdiffusion-joint-ortho-6views.yaml")
+    # print(cfg)
+    schema = OmegaConf.structured(TestConfig)
+    cfg = OmegaConf.merge(schema, cfg)
+
+    pipeline = load_wonder3d_pipeline(cfg)
+    torch.set_grad_enabled(False)
+    pipeline.to(f'cuda:{_GPU_ID}')
+
+    predictor = sam_init()
+
+    custom_theme = gr.themes.Soft(primary_hue="blue").set(
+        button_secondary_background_fill="*neutral_100", button_secondary_background_fill_hover="*neutral_200"
+    )
+    custom_css = '''#disp_image {
+        text-align: center; /* Horizontally center the content */
+    }'''
+
+    with gr.Blocks(title=_TITLE, theme=custom_theme, css=custom_css) as demo:
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown('# ' + _TITLE)
+        gr.Markdown(_DESCRIPTION)
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=1):
+                input_image = gr.Image(type='pil', image_mode='RGBA', height=320, label='Input image', tool=None)
+
+            with gr.Column(scale=1):
+                processed_image = gr.Image(
+                    type='pil',
+                    label="Processed Image",
+                    interactive=False,
+                    height=320,
+                    tool=None,
+                    image_mode='RGBA',
+                    elem_id="disp_image",
+                    visible=True,
+                )
+            # with gr.Column(scale=1):
+            #     ## add 3D Model
+            #     obj_3d = gr.Model3D(
+            #                         # clear_color=[0.0, 0.0, 0.0, 0.0], 
+            #                         label="3D Model", height=320, 
+            #                         # camera_position=[0,0,2.0]
+            #                         )
+                processed_image_highres = gr.Image(type='pil', image_mode='RGBA', visible=False, tool=None)
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=1):
+                example_folder = os.path.join(os.path.dirname(__file__), "./example_images")
+                example_fns = [os.path.join(example_folder, example) for example in os.listdir(example_folder)]
+                gr.Examples(
+                    examples=example_fns,
+                    inputs=[input_image],
+                    outputs=[input_image],
+                    cache_examples=False,
+                    label='Examples (click one of the images below to start)',
+                    examples_per_page=30,
+                )
+            with gr.Column(scale=1):
+                with gr.Accordion('Advanced options', open=True):
+                    with gr.Row():
+                        with gr.Column():
+                            input_processing = gr.CheckboxGroup(
+                                ['Background Removal'],
+                                label='Input Image Preprocessing',
+                                value=['Background Removal'],
+                                info='untick this, if masked image with alpha channel',
+                            )
+                        with gr.Column():
+                            output_processing = gr.CheckboxGroup(
+                                ['Write Results'], label='write the results in ./outputs folder', value=['Write Results']
+                            )
+                    with gr.Row():
+                        with gr.Column():
+                            scale_slider = gr.Slider(1, 5, value=1, step=1, label='Classifier Free Guidance Scale')
+                        with gr.Column():
+                            steps_slider = gr.Slider(15, 100, value=50, step=1, label='Number of Diffusion Inference Steps')
+                    with gr.Row():
+                        with gr.Column():
+                            seed = gr.Number(42, label='Seed')
+                        with gr.Column():
+                            crop_size = gr.Number(192, label='Crop size')
+
+                        mode = gr.Textbox('train', visible=False)
+                        data_dir = gr.Textbox('outputs', visible=False)
+                    # crop_size = 192
+                    # with gr.Row():
+                    #     method = gr.Radio(choices=['instant-nsr-pl', 'NeuS'], label='Method (Default: instant-nsr-pl)', value='instant-nsr-pl')
+                run_btn = gr.Button('Generate Normals and Colors', variant='primary', interactive=True)
+                # recon_btn = gr.Button('Reconstruct 3D model', variant='primary', interactive=True)
+                # gr.Markdown("<span style='color:red'>First click Generate button, then click Reconstruct button. Reconstruction may cost several minutes.</span>")
+        
+        with gr.Row():
+            view_1 = gr.Image(interactive=False, height=240, show_label=False)
+            view_2 = gr.Image(interactive=False, height=240, show_label=False)
+            view_3 = gr.Image(interactive=False, height=240, show_label=False)
+            view_4 = gr.Image(interactive=False, height=240, show_label=False)
+            view_5 = gr.Image(interactive=False, height=240, show_label=False)
+            view_6 = gr.Image(interactive=False, height=240, show_label=False)
+        with gr.Row():
+            normal_1 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_2 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_3 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_4 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_5 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_6 = gr.Image(interactive=False, height=240, show_label=False)
+
+        run_btn.click(
+            fn=partial(preprocess, predictor), inputs=[input_image, input_processing], outputs=[processed_image_highres, processed_image], queue=True
+        ).success(
+            fn=partial(run_pipeline, pipeline, cfg),
+            inputs=[processed_image_highres, scale_slider, steps_slider, seed, crop_size, output_processing],
+            outputs=[view_1, view_2, view_3, view_4, view_5, view_6, normal_1, normal_2, normal_3, normal_4, normal_5, normal_6],
+        )
+        # recon_btn.click(
+        #     process_3d, inputs=[mode, data_dir, scale_slider, crop_size], outputs=[obj_3d]
+        # )
+
+        demo.queue().launch(share=True, max_threads=80)
+
+
+if __name__ == '__main__':
+    fire.Fire(run_demo)

gradio_app_recon.py

@@ -0,0 +1,438 @@
+import os
+import torch
+import fire
+import gradio as gr
+from PIL import Image
+from functools import partial
+
+import cv2
+import time
+import numpy as np
+from rembg import remove
+from segment_anything import sam_model_registry, SamPredictor
+
+import os
+import sys
+import numpy
+import torch
+import rembg
+import threading
+import urllib.request
+from PIL import Image
+from typing import Dict, Optional, Tuple, List
+from dataclasses import dataclass
+import streamlit as st
+import huggingface_hub
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+from mvdiffusion.models.unet_mv2d_condition import UNetMV2DConditionModel
+from mvdiffusion.data.single_image_dataset import SingleImageDataset as MVDiffusionDataset
+from mvdiffusion.pipelines.pipeline_mvdiffusion_image import MVDiffusionImagePipeline
+from diffusers import AutoencoderKL, DDPMScheduler, DDIMScheduler
+from einops import rearrange
+import numpy as np
+import subprocess
+from datetime import datetime
+
+def save_image(tensor):
+    ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    # pdb.set_trace()
+    im = Image.fromarray(ndarr)
+    return ndarr
+
+
+def save_image_to_disk(tensor, fp):
+    ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    # pdb.set_trace()
+    im = Image.fromarray(ndarr)
+    im.save(fp)
+    return ndarr
+
+
+def save_image_numpy(ndarr, fp):
+    im = Image.fromarray(ndarr)
+    im.save(fp)
+
+
+weight_dtype = torch.float16
+
+_TITLE = '''Wonder3D: Single Image to 3D using Cross-Domain Diffusion'''
+_DESCRIPTION = '''
+<div>
+Generate consistent multi-view normals maps and color images.
+<a style="display:inline-block; margin-left: .5em" href='https://github.com/xxlong0/Wonder3D/'><img src='https://img.shields.io/github/stars/xxlong0/Wonder3D?style=social' /></a>
+</div>
+<div>
+The demo does not include the mesh reconstruction part, please visit <a href="https://github.com/xxlong0/Wonder3D/">our github repo</a> to get a textured mesh.
+</div>
+'''
+_GPU_ID = 0
+
+
+if not hasattr(Image, 'Resampling'):
+    Image.Resampling = Image
+
+
+def sam_init():
+    sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
+    model_type = "vit_h"
+
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda:{_GPU_ID}")
+    predictor = SamPredictor(sam)
+    return predictor
+
+
+def sam_segment(predictor, input_image, *bbox_coords):
+    bbox = np.array(bbox_coords)
+    image = np.asarray(input_image)
+
+    start_time = time.time()
+    predictor.set_image(image)
+
+    masks_bbox, scores_bbox, logits_bbox = predictor.predict(box=bbox, multimask_output=True)
+
+    print(f"SAM Time: {time.time() - start_time:.3f}s")
+    out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
+    out_image[:, :, :3] = image
+    out_image_bbox = out_image.copy()
+    out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255
+    torch.cuda.empty_cache()
+    return Image.fromarray(out_image_bbox, mode='RGBA')
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def preprocess(predictor, input_image, chk_group=None, segment=True, rescale=False):
+    RES = 1024
+    input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
+    if chk_group is not None:
+        segment = "Background Removal" in chk_group
+        rescale = "Rescale" in chk_group
+    if segment:
+        image_rem = input_image.convert('RGBA')
+        image_nobg = remove(image_rem, alpha_matting=True)
+        arr = np.asarray(image_nobg)[:, :, -1]
+        x_nonzero = np.nonzero(arr.sum(axis=0))
+        y_nonzero = np.nonzero(arr.sum(axis=1))
+        x_min = int(x_nonzero[0].min())
+        y_min = int(y_nonzero[0].min())
+        x_max = int(x_nonzero[0].max())
+        y_max = int(y_nonzero[0].max())
+        input_image = sam_segment(predictor, input_image.convert('RGB'), x_min, y_min, x_max, y_max)
+    # Rescale and recenter
+    if rescale:
+        image_arr = np.array(input_image)
+        in_w, in_h = image_arr.shape[:2]
+        out_res = min(RES, max(in_w, in_h))
+        ret, mask = cv2.threshold(np.array(input_image.split()[-1]), 0, 255, cv2.THRESH_BINARY)
+        x, y, w, h = cv2.boundingRect(mask)
+        max_size = max(w, h)
+        ratio = 0.75
+        side_len = int(max_size / ratio)
+        padded_image = np.zeros((side_len, side_len, 4), dtype=np.uint8)
+        center = side_len // 2
+        padded_image[center - h // 2 : center - h // 2 + h, center - w // 2 : center - w // 2 + w] = image_arr[y : y + h, x : x + w]
+        rgba = Image.fromarray(padded_image).resize((out_res, out_res), Image.LANCZOS)
+
+        rgba_arr = np.array(rgba) / 255.0
+        rgb = rgba_arr[..., :3] * rgba_arr[..., -1:] + (1 - rgba_arr[..., -1:])
+        input_image = Image.fromarray((rgb * 255).astype(np.uint8))
+    else:
+        input_image = expand2square(input_image, (127, 127, 127, 0))
+    return input_image, input_image.resize((320, 320), Image.Resampling.LANCZOS)
+
+
+def load_wonder3d_pipeline(cfg):
+
+    pipeline = MVDiffusionImagePipeline.from_pretrained(
+    cfg.pretrained_model_name_or_path,
+    torch_dtype=weight_dtype
+    )
+
+    # pipeline.to('cuda:0')
+    pipeline.unet.enable_xformers_memory_efficient_attention()
+
+
+    if torch.cuda.is_available():
+        pipeline.to('cuda:0')
+    # sys.main_lock = threading.Lock()
+    return pipeline
+
+
+from mvdiffusion.data.single_image_dataset import SingleImageDataset
+
+
+def prepare_data(single_image, crop_size):
+    dataset = SingleImageDataset(root_dir='', num_views=6, img_wh=[256, 256], bg_color='white', crop_size=crop_size, single_image=single_image)
+    return dataset[0]
+
+scene = 'scene'
+
+def run_pipeline(pipeline, cfg, single_image, guidance_scale, steps, seed, crop_size, chk_group=None):
+    import pdb
+    global scene
+    # pdb.set_trace()
+
+    if chk_group is not None:
+        write_image = "Write Results" in chk_group
+
+    batch = prepare_data(single_image, crop_size)
+
+    pipeline.set_progress_bar_config(disable=True)
+    seed = int(seed)
+    generator = torch.Generator(device=pipeline.unet.device).manual_seed(seed)
+
+    # repeat  (2B, Nv, 3, H, W)
+    imgs_in = torch.cat([batch['imgs_in']] * 2, dim=0).to(weight_dtype)
+
+    # (2B, Nv, Nce)
+    camera_embeddings = torch.cat([batch['camera_embeddings']] * 2, dim=0).to(weight_dtype)
+
+    task_embeddings = torch.cat([batch['normal_task_embeddings'], batch['color_task_embeddings']], dim=0).to(weight_dtype)
+
+    camera_embeddings = torch.cat([camera_embeddings, task_embeddings], dim=-1).to(weight_dtype)
+
+    # (B*Nv, 3, H, W)
+    imgs_in = rearrange(imgs_in, "Nv C H W -> (Nv) C H W")
+    # (B*Nv, Nce)
+    # camera_embeddings = rearrange(camera_embeddings, "B Nv Nce -> (B Nv) Nce")
+
+    out = pipeline(
+        imgs_in,
+        camera_embeddings,
+        generator=generator,
+        guidance_scale=guidance_scale,
+        num_inference_steps=steps,
+        output_type='pt',
+        num_images_per_prompt=1,
+        **cfg.pipe_validation_kwargs,
+    ).images
+
+    bsz = out.shape[0] // 2
+    normals_pred = out[:bsz]
+    images_pred = out[bsz:]
+    num_views = 6
+    if write_image:
+        VIEWS = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
+        cur_dir = os.path.join("./outputs", f"cropsize-{int(crop_size)}-cfg{guidance_scale:.1f}")
+
+        scene = 'scene'+datetime.now().strftime('@%Y%m%d-%H%M%S')
+        scene_dir = os.path.join(cur_dir, scene)
+        normal_dir = os.path.join(scene_dir, "normals")
+        masked_colors_dir = os.path.join(scene_dir, "masked_colors")
+        os.makedirs(normal_dir, exist_ok=True)
+        os.makedirs(masked_colors_dir, exist_ok=True)
+        for j in range(num_views):
+            view = VIEWS[j]
+            normal = normals_pred[j]
+            color = images_pred[j]
+
+            normal_filename = f"normals_000_{view}.png"
+            rgb_filename = f"rgb_000_{view}.png"
+            normal = save_image_to_disk(normal, os.path.join(normal_dir, normal_filename))
+            color = save_image_to_disk(color, os.path.join(scene_dir, rgb_filename))
+
+            rm_normal = remove(normal)
+            rm_color = remove(color)
+
+            save_image_numpy(rm_normal, os.path.join(scene_dir, normal_filename))
+            save_image_numpy(rm_color, os.path.join(masked_colors_dir, rgb_filename))
+
+    normals_pred = [save_image(normals_pred[i]) for i in range(bsz)]
+    images_pred = [save_image(images_pred[i]) for i in range(bsz)]
+
+    out = images_pred + normals_pred
+    return out
+
+
+def process_3d(mode, data_dir, guidance_scale, crop_size):
+    dir = None
+    global scene
+
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+    subprocess.run(
+        f'cd instant-nsr-pl && python launch.py --config configs/neuralangelo-ortho-wmask.yaml --gpu 0 --train dataset.root_dir=../{data_dir}/cropsize-{int(crop_size)}-cfg{guidance_scale:.1f}/ dataset.scene={scene} && cd ..',
+        shell=True,
+    )
+    import glob
+    # import pdb
+
+    # pdb.set_trace()
+
+    obj_files = glob.glob(f'{cur_dir}/instant-nsr-pl/exp/{scene}/*/save/*.obj', recursive=True)
+    print(obj_files)
+    if obj_files:
+        dir = obj_files[0]
+    return dir
+
+
+@dataclass
+class TestConfig:
+    pretrained_model_name_or_path: str
+    pretrained_unet_path: str
+    revision: Optional[str]
+    validation_dataset: Dict
+    save_dir: str
+    seed: Optional[int]
+    validation_batch_size: int
+    dataloader_num_workers: int
+
+    local_rank: int
+
+    pipe_kwargs: Dict
+    pipe_validation_kwargs: Dict
+    unet_from_pretrained_kwargs: Dict
+    validation_guidance_scales: List[float]
+    validation_grid_nrow: int
+    camera_embedding_lr_mult: float
+
+    num_views: int
+    camera_embedding_type: str
+
+    pred_type: str  # joint, or ablation
+
+    enable_xformers_memory_efficient_attention: bool
+
+    cond_on_normals: bool
+    cond_on_colors: bool
+
+
+def run_demo():
+    from utils.misc import load_config
+    from omegaconf import OmegaConf
+
+    # parse YAML config to OmegaConf
+    cfg = load_config("./configs/mvdiffusion-joint-ortho-6views.yaml")
+    # print(cfg)
+    schema = OmegaConf.structured(TestConfig)
+    cfg = OmegaConf.merge(schema, cfg)
+
+    pipeline = load_wonder3d_pipeline(cfg)
+    torch.set_grad_enabled(False)
+    pipeline.to(f'cuda:{_GPU_ID}')
+
+    predictor = sam_init()
+
+    custom_theme = gr.themes.Soft(primary_hue="blue").set(
+        button_secondary_background_fill="*neutral_100", button_secondary_background_fill_hover="*neutral_200"
+    )
+    custom_css = '''#disp_image {
+        text-align: center; /* Horizontally center the content */
+    }'''
+
+    with gr.Blocks(title=_TITLE, theme=custom_theme, css=custom_css) as demo:
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown('# ' + _TITLE)
+        gr.Markdown(_DESCRIPTION)
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=1):
+                input_image = gr.Image(type='pil', image_mode='RGBA', height=320, label='Input image', tool=None)
+
+            with gr.Column(scale=1):
+                processed_image = gr.Image(
+                    type='pil',
+                    label="Processed Image",
+                    interactive=False,
+                    height=320,
+                    tool=None,
+                    image_mode='RGBA',
+                    elem_id="disp_image",
+                    visible=True,
+                )
+            with gr.Column(scale=1):
+                ## add 3D Model
+                obj_3d = gr.Model3D(
+                                    # clear_color=[0.0, 0.0, 0.0, 0.0], 
+                                    label="3D Model", height=320, 
+                                    # camera_position=[0,0,2.0]
+                                    )
+                processed_image_highres = gr.Image(type='pil', image_mode='RGBA', visible=False, tool=None)
+        with gr.Row(variant='panel'):
+            with gr.Column(scale=1):
+                example_folder = os.path.join(os.path.dirname(__file__), "./example_images")
+                example_fns = [os.path.join(example_folder, example) for example in os.listdir(example_folder)]
+                gr.Examples(
+                    examples=example_fns,
+                    inputs=[input_image],
+                    outputs=[input_image],
+                    cache_examples=False,
+                    label='Examples (click one of the images below to start)',
+                    examples_per_page=30,
+                )
+            with gr.Column(scale=1):
+                with gr.Accordion('Advanced options', open=True):
+                    with gr.Row():
+                        with gr.Column():
+                            input_processing = gr.CheckboxGroup(
+                                ['Background Removal'],
+                                label='Input Image Preprocessing',
+                                value=['Background Removal'],
+                                info='untick this, if masked image with alpha channel',
+                            )
+                        with gr.Column():
+                            output_processing = gr.CheckboxGroup(
+                                ['Write Results'], label='write the results in ./outputs folder', value=['Write Results']
+                            )
+                    with gr.Row():
+                        with gr.Column():
+                            scale_slider = gr.Slider(1, 5, value=1, step=1, label='Classifier Free Guidance Scale')
+                        with gr.Column():
+                            steps_slider = gr.Slider(15, 100, value=50, step=1, label='Number of Diffusion Inference Steps')
+                    with gr.Row():
+                        with gr.Column():
+                            seed = gr.Number(42, label='Seed')
+                        with gr.Column():
+                            crop_size = gr.Number(192, label='Crop size')
+
+                        mode = gr.Textbox('train', visible=False)
+                        data_dir = gr.Textbox('outputs', visible=False)
+                    # crop_size = 192
+                    # with gr.Row():
+                    #     method = gr.Radio(choices=['instant-nsr-pl', 'NeuS'], label='Method (Default: instant-nsr-pl)', value='instant-nsr-pl')
+                # run_btn = gr.Button('Generate Normals and Colors', variant='primary', interactive=True)
+                run_btn = gr.Button('Reconstruct 3D model', variant='primary', interactive=True)
+                gr.Markdown("<span style='color:red'> Reconstruction may cost several minutes. Check results in instant-nsr-pl/exp/scene@{current-time}/ </span>")
+        
+        with gr.Row():
+            view_1 = gr.Image(interactive=False, height=240, show_label=False)
+            view_2 = gr.Image(interactive=False, height=240, show_label=False)
+            view_3 = gr.Image(interactive=False, height=240, show_label=False)
+            view_4 = gr.Image(interactive=False, height=240, show_label=False)
+            view_5 = gr.Image(interactive=False, height=240, show_label=False)
+            view_6 = gr.Image(interactive=False, height=240, show_label=False)
+        with gr.Row():
+            normal_1 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_2 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_3 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_4 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_5 = gr.Image(interactive=False, height=240, show_label=False)
+            normal_6 = gr.Image(interactive=False, height=240, show_label=False)
+
+        run_btn.click(
+            fn=partial(preprocess, predictor), inputs=[input_image, input_processing], outputs=[processed_image_highres, processed_image], queue=True
+        ).success(
+            fn=partial(run_pipeline, pipeline, cfg),
+            inputs=[processed_image_highres, scale_slider, steps_slider, seed, crop_size, output_processing],
+            outputs=[view_1, view_2, view_3, view_4, view_5, view_6, normal_1, normal_2, normal_3, normal_4, normal_5, normal_6],
+        ).success(
+            process_3d, inputs=[mode, data_dir, scale_slider, crop_size], outputs=[obj_3d]
+        )
+
+        demo.queue().launch(share=True, max_threads=80)
+
+
+if __name__ == '__main__':
+    fire.Fire(run_demo)

instant-nsr-pl/README.md

@@ -0,0 +1,122 @@
+# Instant Neural Surface Reconstruction
+
+This repository contains a concise and extensible implementation of NeRF and NeuS for neural surface reconstruction based on Instant-NGP and the Pytorch-Lightning framework. **Training on a NeRF-Synthetic scene takes ~5min for NeRF and ~10min for NeuS on a single RTX3090.**
+
+||NeRF in 5min|NeuS in 10 min|
+|---|---|---|
+|Rendering|![rendering-nerf](https://user-images.githubusercontent.com/19284678/199078178-b719676b-7e60-47f1-813b-c0b533f5480d.png)|![rendering-neus](https://user-images.githubusercontent.com/19284678/199078300-ebcf249d-b05e-431f-b035-da354705d8db.png)|
+|Mesh|![mesh-nerf](https://user-images.githubusercontent.com/19284678/199078661-b5cd569a-c22b-4220-9c11-d5fd13a52fb8.png)|![mesh-neus](https://user-images.githubusercontent.com/19284678/199078481-164e36a6-6d55-45cc-aaf3-795a114e4a38.png)|
+
+
+## Features
+**This repository aims to provide a highly efficient while customizable boilerplate for research projects based on NeRF or NeuS.**
+
+- acceleration techniques from [Instant-NGP](https://github.com/NVlabs/instant-ngp): multiresolution hash encoding and fully fused networks by [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn), occupancy grid pruning and rendering by [nerfacc](https://github.com/KAIR-BAIR/nerfacc)
+- out-of-the-box multi-GPU and mixed precision training by [PyTorch-Lightning](https://github.com/Lightning-AI/lightning)
+- hierarchical project layout that is designed to be easily customized and extended, flexible experiment configuration by [OmegaConf](https://github.com/omry/omegaconf)
+
+**Please subscribe to [#26](https://github.com/bennyguo/instant-nsr-pl/issues/26) for our latest findings on quality improvements!**
+
+## News
+
+🔥🔥🔥 Check out my new project on 3D content generation: https://github.com/threestudio-project/threestudio 🔥🔥🔥
+
+- 06/03/2023: Add an implementation of [Neuralangelo](https://research.nvidia.com/labs/dir/neuralangelo/). See [here](https://github.com/bennyguo/instant-nsr-pl#training-on-DTU) for details.
+- 03/31/2023: NeuS model now supports background modeling. You could try on the DTU dataset provided by [NeuS](https://drive.google.com/drive/folders/1Nlzejs4mfPuJYORLbDEUDWlc9IZIbU0C?usp=sharing) or [IDR](https://www.dropbox.com/sh/5tam07ai8ch90pf/AADniBT3dmAexvm_J1oL__uoa) following [the instruction here](https://github.com/bennyguo/instant-nsr-pl#training-on-DTU).
+- 02/11/2023: NeRF model now supports unbounded 360 scenes with learned background. You could try on [MipNeRF 360 data](http://storage.googleapis.com/gresearch/refraw360/360_v2.zip) following [the COLMAP configuration](https://github.com/bennyguo/instant-nsr-pl#training-on-custom-colmap-data).
+
+## Requirements
+**Note:**
+- To utilize multiresolution hash encoding or fully fused networks provided by tiny-cuda-nn, you should have least an RTX 2080Ti, see [https://github.com/NVlabs/tiny-cuda-nn#requirements](https://github.com/NVlabs/tiny-cuda-nn#requirements) for more details.
+- Multi-GPU training is currently not supported on Windows (see [#4](https://github.com/bennyguo/instant-nsr-pl/issues/4)).
+### Environments
+- Install PyTorch>=1.10 [here](https://pytorch.org/get-started/locally/) based the package management tool you used and your cuda version (older PyTorch versions may work but have not been tested)
+- Install tiny-cuda-nn PyTorch extension: `pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch`
+- `pip install -r requirements.txt`
+
+
+## Run
+### Training on NeRF-Synthetic
+Download the NeRF-Synthetic data [here](https://drive.google.com/drive/folders/128yBriW1IG_3NJ5Rp7APSTZsJqdJdfc1) and put it under `load/`. The file structure should be like `load/nerf_synthetic/lego`.
+
+Run the launch script with `--train`, specifying the config file, the GPU(s) to be used (GPU 0 will be used by default), and the scene name:
+```bash
+# train NeRF
+python launch.py --config configs/nerf-blender.yaml --gpu 0 --train dataset.scene=lego tag=example
+
+# train NeuS with mask
+python launch.py --config configs/neus-blender.yaml --gpu 0 --train dataset.scene=lego tag=example
+# train NeuS without mask
+python launch.py --config configs/neus-blender.yaml --gpu 0 --train dataset.scene=lego tag=example system.loss.lambda_mask=0.0
+```
+The code snapshots, checkpoints and experiment outputs are saved to `exp/[name]/[tag]@[timestamp]`, and tensorboard logs can be found at `runs/[name]/[tag]@[timestamp]`. You can change any configuration in the YAML file by specifying arguments without `--`, for example:
+```bash
+python launch.py --config configs/nerf-blender.yaml --gpu 0 --train dataset.scene=lego tag=iter50k seed=0 trainer.max_steps=50000
+```
+### Training on DTU
+Download preprocessed DTU data provided by [NeuS](https://drive.google.com/drive/folders/1Nlzejs4mfPuJYORLbDEUDWlc9IZIbU0C?usp=sharing) or [IDR](https://www.dropbox.com/sh/5tam07ai8ch90pf/AADniBT3dmAexvm_J1oL__uoa). In the provided config files we assume using NeuS DTU data. If you are using IDR DTU data, please set `dataset.cameras_file=cameras.npz`. You may also need to adjust `dataset.root_dir` to point to your downloaded data location.
+```bash
+# train NeuS on DTU without mask
+python launch.py --config configs/neus-dtu.yaml --gpu 0 --train
+# train NeuS on DTU with mask
+python launch.py --config configs/neus-dtu-wmask.yaml --gpu 0 --train
+# train NeuS on DTU with mask using tricks from Neuralangelo (experimental)
+python launch.py --config configs/neuralangelo-dtu-wmask.yaml --gpu 0 --train
+```
+Notes:
+- PSNR in the testing stage is meaningless, as we simply compare to pure white images in testing.
+- The results of Neuralangelo can't reach those in the original paper. Some potential improvements: more iterations; larger `system.geometry.xyz_encoding_config.update_steps`; larger `system.geometry.xyz_encoding_config.n_features_per_level`; larger `system.geometry.xyz_encoding_config.log2_hashmap_size`; adopting curvature loss.
+
+### Training on Custom COLMAP Data
+To get COLMAP data from custom images, you should have COLMAP installed (see [here](https://colmap.github.io/install.html) for installation instructions). Then put your images in the `images/` folder, and run `scripts/imgs2poses.py` specifying the path containing the `images/` folder. For example:
+```bash
+python scripts/imgs2poses.py ./load/bmvs_dog # images are in ./load/bmvs_dog/images
+```
+Existing data following this file structure also works as long as images are store in `images/` and there is a `sparse/` folder for the COLMAP output, for example [the data provided by MipNeRF 360](http://storage.googleapis.com/gresearch/refraw360/360_v2.zip). An optional `masks/` folder could be provided for object mask supervision. To train on COLMAP data, please refer to the example config files `config/*-colmap.yaml`. Some notes:
+- Adapt the `root_dir` and `img_wh` (or `img_downscale`) option in the config file to your data;
+- The scene is normalized so that cameras have a minimum distance `1.0` to the center of the scene. Setting `model.radius=1.0` works in most cases. If not, try setting a smaller radius that wraps tightly to your foreground object.
+- There are three choices to determine the scene center: `dataset.center_est_method=camera` uses the center of all camera positions as the scene center; `dataset.center_est_method=lookat` assumes the cameras are looking at the same point and calculates an approximate look-at point as the scene center; `dataset.center_est_method=point` uses the center of all points (reconstructed by COLMAP) that are bounded by cameras as the scene center. Please choose an appropriate method according to your capture.
+- PSNR in the testing stage is meaningless, as we simply compare to pure white images in testing.
+
+### Testing
+The training procedure are by default followed by testing, which computes metrics on test data, generates animations and exports the geometry as triangular meshes. If you want to do testing alone, just resume the pretrained model and replace `--train` with `--test`, for example:
+```bash
+python launch.py --config path/to/your/exp/config/parsed.yaml --resume path/to/your/exp/ckpt/epoch=0-step=20000.ckpt --gpu 0 --test
+```
+
+
+## Benchmarks
+All experiments are conducted on a single NVIDIA RTX3090.
+
+|PSNR|Chair|Drums|Ficus|Hotdog|Lego|Materials|Mic|Ship|Avg.|
+|---|---|---|---|---|---|---|---|---|---|
+|NeRF Paper|33.00|25.01|30.13|36.18|32.54|29.62|32.91|28.65|31.01|
+|NeRF Ours (20k)|34.80|26.04|33.89|37.42|35.33|29.46|35.22|31.17|32.92|
+|NeuS Ours (20k, with masks)|34.04|25.26|32.47|35.94|33.78|27.67|33.43|29.50|31.51|
+
+|Training Time (mm:ss)|Chair|Drums|Ficus|Hotdog|Lego|Materials|Mic|Ship|Avg.|
+|---|---|---|---|---|---|---|---|---|---|
+|NeRF Ours (20k)|04:34|04:35|04:18|04:46|04:39|04:35|04:26|05:41|04:42|
+|NeuS Ours (20k, with masks)|11:25|10:34|09:51|12:11|11:37|11:46|09:59|16:25|11:44|
+
+
+## TODO
+- [✅] Support more dataset formats, like COLMAP outputs and DTU
+- [✅] Support simple background model
+- [ ] Support GUI training and interaction
+- [ ] More illustrations about the framework
+
+## Related Projects
+- [ngp_pl](https://github.com/kwea123/ngp_pl): Great Instant-NGP implementation in PyTorch-Lightning! Background model and GUI supported.
+- [Instant-NSR](https://github.com/zhaofuq/Instant-NSR): NeuS implementation using multiresolution hash encoding.
+
+## Citation
+If you find this codebase useful, please consider citing:
+```
+@misc{instant-nsr-pl,
+    Author = {Yuan-Chen Guo},
+    Year = {2022},
+    Note = {https://github.com/bennyguo/instant-nsr-pl},
+    Title = {Instant Neural Surface Reconstruction}
+}
+```

instant-nsr-pl/configs/neuralangelo-ortho-wmask.yaml

@@ -0,0 +1,145 @@
+name: ${basename:${dataset.scene}}
+tag: ""
+seed: 42
+
+dataset:
+  name: ortho
+  root_dir: /home/xiaoxiao/Workplace/wonder3Dplus/outputs/joint-twice/aigc/cropsize-224-cfg1.0
+  cam_pose_dir: null
+  scene: scene_name
+  imSize: [1024, 1024]  # should use larger res, otherwise the exported mesh has wrong colors
+  camera_type: ortho
+  apply_mask: true
+  camera_params: null
+  view_weights: [1.0, 0.8, 0.2, 1.0, 0.4, 0.7]  #['front', 'front_right', 'right', 'back', 'left', 'front_left']
+  # view_weights: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] 
+
+model:
+  name: neus
+  radius: 1.0
+  num_samples_per_ray: 1024
+  train_num_rays: 256
+  max_train_num_rays: 8192
+  grid_prune: true
+  grid_prune_occ_thre: 0.001
+  dynamic_ray_sampling: true
+  batch_image_sampling: true
+  randomized: true
+  ray_chunk: 2048
+  cos_anneal_end: 20000
+  learned_background: false
+  background_color: black
+  variance:
+    init_val: 0.3
+    modulate: false
+  geometry:
+    name: volume-sdf
+    radius: ${model.radius}
+    feature_dim: 13
+    grad_type: finite_difference
+    finite_difference_eps: progressive
+    isosurface:
+      method: mc
+      resolution: 192
+      chunk: 2097152
+      threshold: 0.
+    xyz_encoding_config:
+      otype: ProgressiveBandHashGrid
+      n_levels: 10 # 12 modify
+      n_features_per_level: 2
+      log2_hashmap_size: 19
+      base_resolution: 32
+      per_level_scale: 1.3195079107728942
+      include_xyz: true
+      start_level: 4
+      start_step: 0
+      update_steps: 1000
+    mlp_network_config:
+      otype: VanillaMLP
+      activation: ReLU
+      output_activation: none
+      n_neurons: 64
+      n_hidden_layers: 1
+      sphere_init: true
+      sphere_init_radius: 0.5
+      weight_norm: true
+  texture:
+    name: volume-radiance
+    input_feature_dim: ${add:${model.geometry.feature_dim},3} # surface normal as additional input
+    dir_encoding_config:
+      otype: SphericalHarmonics
+      degree: 4
+    mlp_network_config:
+      otype: VanillaMLP
+      activation: ReLU
+      output_activation: none
+      n_neurons: 64
+      n_hidden_layers: 2
+    color_activation: sigmoid
+
+system:
+  name: ortho-neus-system
+  loss:
+    lambda_rgb_mse: 0.5
+    lambda_rgb_l1: 0.
+    lambda_mask: 1.0
+    lambda_eikonal: 0.2  # cannot be too large, will cause holes to thin objects
+    lambda_normal: 1.0  # cannot be too large
+    lambda_3d_normal_smooth: 1.0
+    # lambda_curvature: [0, 0.0, 1.e-4, 1000] # topology warmup
+    lambda_curvature: 0.
+    lambda_sparsity: 0.5
+    lambda_distortion: 0.0
+    lambda_distortion_bg: 0.0
+    lambda_opaque: 0.0
+    sparsity_scale: 100.0
+    geo_aware: true
+    rgb_p_ratio: 0.8
+    normal_p_ratio: 0.8
+    mask_p_ratio: 0.9
+  optimizer:
+    name: AdamW
+    args:
+      lr: 0.01
+      betas: [0.9, 0.99]
+      eps: 1.e-15
+    params:
+      geometry:
+        lr: 0.001
+      texture:
+        lr: 0.01
+      variance:
+        lr: 0.001
+  constant_steps: 500
+  scheduler:
+    name: SequentialLR
+    interval: step
+    milestones:
+      - ${system.constant_steps}
+    schedulers:
+      - name: ConstantLR
+        args:
+          factor: 1.0
+          total_iters: ${system.constant_steps}
+      - name: ExponentialLR
+        args:
+          gamma: ${calc_exp_lr_decay_rate:0.1,${sub:${trainer.max_steps},${system.constant_steps}}}
+
+checkpoint:
+  save_top_k: -1
+  every_n_train_steps: ${trainer.max_steps}
+
+export:
+  chunk_size: 2097152
+  export_vertex_color: True
+  ortho_scale: 1.35   #modify
+
+trainer:
+  max_steps: 3000
+  log_every_n_steps: 100
+  num_sanity_val_steps: 0
+  val_check_interval: 4000
+  limit_train_batches: 1.0
+  limit_val_batches: 2
+  enable_progress_bar: true
+  precision: 16

instant-nsr-pl/datasets/__init__.py

@@ -0,0 +1,16 @@
+datasets = {}
+
+
+def register(name):
+    def decorator(cls):
+        datasets[name] = cls
+        return cls
+    return decorator
+
+
+def make(name, config):
+    dataset = datasets[name](config)
+    return dataset
+
+
+from . import blender, colmap, dtu, ortho

instant-nsr-pl/datasets/blender.py

@@ -0,0 +1,135 @@
+import os
+import json
+import math
+import numpy as np
+from PIL import Image
+
+import torch
+from torch.utils.data import Dataset, DataLoader, IterableDataset
+import torchvision.transforms.functional as TF
+
+import pytorch_lightning as pl
+
+import datasets
+from models.ray_utils import get_ray_directions
+from utils.misc import get_rank
+
+
+class BlenderDatasetBase():
+    def setup(self, config, split):
+        self.config = config
+        self.split = split
+        self.rank = get_rank()
+
+        self.has_mask = True
+        self.apply_mask = True
+
+        with open(os.path.join(self.config.root_dir, f"transforms_{self.split}.json"), 'r') as f:
+            meta = json.load(f)
+
+        if 'w' in meta and 'h' in meta:
+            W, H = int(meta['w']), int(meta['h'])
+        else:
+            W, H = 800, 800
+
+        if 'img_wh' in self.config:
+            w, h = self.config.img_wh
+            assert round(W / w * h) == H
+        elif 'img_downscale' in self.config:
+            w, h = W // self.config.img_downscale, H // self.config.img_downscale
+        else:
+            raise KeyError("Either img_wh or img_downscale should be specified.")
+        
+        self.w, self.h = w, h
+        self.img_wh = (self.w, self.h)
+
+        self.near, self.far = self.config.near_plane, self.config.far_plane
+
+        self.focal = 0.5 * w / math.tan(0.5 * meta['camera_angle_x']) # scaled focal length
+
+        # ray directions for all pixels, same for all images (same H, W, focal)
+        self.directions = \
+            get_ray_directions(self.w, self.h, self.focal, self.focal, self.w//2, self.h//2).to(self.rank) # (h, w, 3)           
+
+        self.all_c2w, self.all_images, self.all_fg_masks = [], [], []
+
+        for i, frame in enumerate(meta['frames']):
+            c2w = torch.from_numpy(np.array(frame['transform_matrix'])[:3, :4])
+            self.all_c2w.append(c2w)
+
+            img_path = os.path.join(self.config.root_dir, f"{frame['file_path']}.png")
+            img = Image.open(img_path)
+            img = img.resize(self.img_wh, Image.BICUBIC)
+            img = TF.to_tensor(img).permute(1, 2, 0) # (4, h, w) => (h, w, 4)
+
+            self.all_fg_masks.append(img[..., -1]) # (h, w)
+            self.all_images.append(img[...,:3])
+
+        self.all_c2w, self.all_images, self.all_fg_masks = \
+            torch.stack(self.all_c2w, dim=0).float().to(self.rank), \
+            torch.stack(self.all_images, dim=0).float().to(self.rank), \
+            torch.stack(self.all_fg_masks, dim=0).float().to(self.rank)
+        
+
+class BlenderDataset(Dataset, BlenderDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __len__(self):
+        return len(self.all_images)
+    
+    def __getitem__(self, index):
+        return {
+            'index': index
+        }
+
+
+class BlenderIterableDataset(IterableDataset, BlenderDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __iter__(self):
+        while True:
+            yield {}
+
+
+@datasets.register('blender')
+class BlenderDataModule(pl.LightningDataModule):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+    
+    def setup(self, stage=None):
+        if stage in [None, 'fit']:
+            self.train_dataset = BlenderIterableDataset(self.config, self.config.train_split)
+        if stage in [None, 'fit', 'validate']:
+            self.val_dataset = BlenderDataset(self.config, self.config.val_split)
+        if stage in [None, 'test']:
+            self.test_dataset = BlenderDataset(self.config, self.config.test_split)
+        if stage in [None, 'predict']:
+            self.predict_dataset = BlenderDataset(self.config, self.config.train_split)
+
+    def prepare_data(self):
+        pass
+    
+    def general_loader(self, dataset, batch_size):
+        sampler = None
+        return DataLoader(
+            dataset, 
+            num_workers=os.cpu_count(), 
+            batch_size=batch_size,
+            pin_memory=True,
+            sampler=sampler
+        )
+    
+    def train_dataloader(self):
+        return self.general_loader(self.train_dataset, batch_size=1)
+
+    def val_dataloader(self):
+        return self.general_loader(self.val_dataset, batch_size=1)
+
+    def test_dataloader(self):
+        return self.general_loader(self.test_dataset, batch_size=1) 
+
+    def predict_dataloader(self):
+        return self.general_loader(self.predict_dataset, batch_size=1)       

instant-nsr-pl/datasets/colmap.py

@@ -0,0 +1,332 @@
+import os
+import math
+import numpy as np
+from PIL import Image
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, IterableDataset
+import torchvision.transforms.functional as TF
+
+import pytorch_lightning as pl
+
+import datasets
+from datasets.colmap_utils import \
+    read_cameras_binary, read_images_binary, read_points3d_binary
+from models.ray_utils import get_ray_directions
+from utils.misc import get_rank
+
+
+def get_center(pts):
+    center = pts.mean(0)
+    dis = (pts - center[None,:]).norm(p=2, dim=-1)
+    mean, std = dis.mean(), dis.std()
+    q25, q75 = torch.quantile(dis, 0.25), torch.quantile(dis, 0.75)
+    valid = (dis > mean - 1.5 * std) & (dis < mean + 1.5 * std) & (dis > mean - (q75 - q25) * 1.5) & (dis < mean + (q75 - q25) * 1.5)
+    center = pts[valid].mean(0)
+    return center
+
+def normalize_poses(poses, pts, up_est_method, center_est_method):
+    if center_est_method == 'camera':
+        # estimation scene center as the average of all camera positions
+        center = poses[...,3].mean(0)
+    elif center_est_method == 'lookat':
+        # estimation scene center as the average of the intersection of selected pairs of camera rays
+        cams_ori = poses[...,3]
+        cams_dir = poses[:,:3,:3] @ torch.as_tensor([0.,0.,-1.])
+        cams_dir = F.normalize(cams_dir, dim=-1)
+        A = torch.stack([cams_dir, -cams_dir.roll(1,0)], dim=-1)
+        b = -cams_ori + cams_ori.roll(1,0)
+        t = torch.linalg.lstsq(A, b).solution
+        center = (torch.stack([cams_dir, cams_dir.roll(1,0)], dim=-1) * t[:,None,:] + torch.stack([cams_ori, cams_ori.roll(1,0)], dim=-1)).mean((0,2))
+    elif center_est_method == 'point':
+        # first estimation scene center as the average of all camera positions
+        # later we'll use the center of all points bounded by the cameras as the final scene center
+        center = poses[...,3].mean(0)
+    else:
+        raise NotImplementedError(f'Unknown center estimation method: {center_est_method}')
+
+    if up_est_method == 'ground':
+        # estimate up direction as the normal of the estimated ground plane
+        # use RANSAC to estimate the ground plane in the point cloud
+        import pyransac3d as pyrsc
+        ground = pyrsc.Plane()
+        plane_eq, inliers = ground.fit(pts.numpy(), thresh=0.01) # TODO: determine thresh based on scene scale
+        plane_eq = torch.as_tensor(plane_eq) # A, B, C, D in Ax + By + Cz + D = 0
+        z = F.normalize(plane_eq[:3], dim=-1) # plane normal as up direction
+        signed_distance = (torch.cat([pts, torch.ones_like(pts[...,0:1])], dim=-1) * plane_eq).sum(-1)
+        if signed_distance.mean() < 0:
+            z = -z # flip the direction if points lie under the plane
+    elif up_est_method == 'camera':
+        # estimate up direction as the average of all camera up directions
+        z = F.normalize((poses[...,3] - center).mean(0), dim=0)
+    else:
+        raise NotImplementedError(f'Unknown up estimation method: {up_est_method}')
+
+    # new axis
+    y_ = torch.as_tensor([z[1], -z[0], 0.])
+    x = F.normalize(y_.cross(z), dim=0)
+    y = z.cross(x)
+
+    if center_est_method == 'point':
+        # rotation
+        Rc = torch.stack([x, y, z], dim=1)
+        R = Rc.T
+        poses_homo = torch.cat([poses, torch.as_tensor([[[0.,0.,0.,1.]]]).expand(poses.shape[0], -1, -1)], dim=1)
+        inv_trans = torch.cat([torch.cat([R, torch.as_tensor([[0.,0.,0.]]).T], dim=1), torch.as_tensor([[0.,0.,0.,1.]])], dim=0)
+        poses_norm = (inv_trans @ poses_homo)[:,:3]
+        pts = (inv_trans @ torch.cat([pts, torch.ones_like(pts[:,0:1])], dim=-1)[...,None])[:,:3,0]
+
+        # translation and scaling
+        poses_min, poses_max = poses_norm[...,3].min(0)[0], poses_norm[...,3].max(0)[0]
+        pts_fg = pts[(poses_min[0] < pts[:,0]) & (pts[:,0] < poses_max[0]) & (poses_min[1] < pts[:,1]) & (pts[:,1] < poses_max[1])]
+        center = get_center(pts_fg)
+        tc = center.reshape(3, 1)
+        t = -tc
+        poses_homo = torch.cat([poses_norm, torch.as_tensor([[[0.,0.,0.,1.]]]).expand(poses_norm.shape[0], -1, -1)], dim=1)
+        inv_trans = torch.cat([torch.cat([torch.eye(3), t], dim=1), torch.as_tensor([[0.,0.,0.,1.]])], dim=0)
+        poses_norm = (inv_trans @ poses_homo)[:,:3]
+        scale = poses_norm[...,3].norm(p=2, dim=-1).min()
+        poses_norm[...,3] /= scale
+        pts = (inv_trans @ torch.cat([pts, torch.ones_like(pts[:,0:1])], dim=-1)[...,None])[:,:3,0]
+        pts = pts / scale
+    else:
+        # rotation and translation
+        Rc = torch.stack([x, y, z], dim=1)
+        tc = center.reshape(3, 1)
+        R, t = Rc.T, -Rc.T @ tc
+        poses_homo = torch.cat([poses, torch.as_tensor([[[0.,0.,0.,1.]]]).expand(poses.shape[0], -1, -1)], dim=1)
+        inv_trans = torch.cat([torch.cat([R, t], dim=1), torch.as_tensor([[0.,0.,0.,1.]])], dim=0)
+        poses_norm = (inv_trans @ poses_homo)[:,:3] # (N_images, 4, 4)
+
+        # scaling
+        scale = poses_norm[...,3].norm(p=2, dim=-1).min()
+        poses_norm[...,3] /= scale
+
+        # apply the transformation to the point cloud
+        pts = (inv_trans @ torch.cat([pts, torch.ones_like(pts[:,0:1])], dim=-1)[...,None])[:,:3,0]
+        pts = pts / scale
+
+    return poses_norm, pts
+
+def create_spheric_poses(cameras, n_steps=120):
+    center = torch.as_tensor([0.,0.,0.], dtype=cameras.dtype, device=cameras.device)
+    mean_d = (cameras - center[None,:]).norm(p=2, dim=-1).mean()
+    mean_h = cameras[:,2].mean()
+    r = (mean_d**2 - mean_h**2).sqrt()
+    up = torch.as_tensor([0., 0., 1.], dtype=center.dtype, device=center.device)
+
+    all_c2w = []
+    for theta in torch.linspace(0, 2 * math.pi, n_steps):
+        cam_pos = torch.stack([r * theta.cos(), r * theta.sin(), mean_h])
+        l = F.normalize(center - cam_pos, p=2, dim=0)
+        s = F.normalize(l.cross(up), p=2, dim=0)
+        u = F.normalize(s.cross(l), p=2, dim=0)
+        c2w = torch.cat([torch.stack([s, u, -l], dim=1), cam_pos[:,None]], axis=1)
+        all_c2w.append(c2w)
+
+    all_c2w = torch.stack(all_c2w, dim=0)
+    
+    return all_c2w
+
+class ColmapDatasetBase():
+    # the data only has to be processed once
+    initialized = False
+    properties = {}
+
+    def setup(self, config, split):
+        self.config = config
+        self.split = split
+        self.rank = get_rank()
+
+        if not ColmapDatasetBase.initialized:
+            camdata = read_cameras_binary(os.path.join(self.config.root_dir, 'sparse/0/cameras.bin'))
+
+            H = int(camdata[1].height)
+            W = int(camdata[1].width)
+
+            if 'img_wh' in self.config:
+                w, h = self.config.img_wh
+                assert round(W / w * h) == H
+            elif 'img_downscale' in self.config:
+                w, h = int(W / self.config.img_downscale + 0.5), int(H / self.config.img_downscale + 0.5)
+            else:
+                raise KeyError("Either img_wh or img_downscale should be specified.")
+
+            img_wh = (w, h)
+            factor = w / W
+
+            if camdata[1].model == 'SIMPLE_RADIAL':
+                fx = fy = camdata[1].params[0] * factor
+                cx = camdata[1].params[1] * factor
+                cy = camdata[1].params[2] * factor
+            elif camdata[1].model in ['PINHOLE', 'OPENCV']:
+                fx = camdata[1].params[0] * factor
+                fy = camdata[1].params[1] * factor
+                cx = camdata[1].params[2] * factor
+                cy = camdata[1].params[3] * factor
+            else:
+                raise ValueError(f"Please parse the intrinsics for camera model {camdata[1].model}!")
+            
+            directions = get_ray_directions(w, h, fx, fy, cx, cy).to(self.rank)
+
+            imdata = read_images_binary(os.path.join(self.config.root_dir, 'sparse/0/images.bin'))
+
+            mask_dir = os.path.join(self.config.root_dir, 'masks')
+            has_mask = os.path.exists(mask_dir) # TODO: support partial masks
+            apply_mask = has_mask and self.config.apply_mask
+            
+            all_c2w, all_images, all_fg_masks = [], [], []
+
+            for i, d in enumerate(imdata.values()):
+                R = d.qvec2rotmat()
+                t = d.tvec.reshape(3, 1)
+                c2w = torch.from_numpy(np.concatenate([R.T, -R.T@t], axis=1)).float()
+                c2w[:,1:3] *= -1. # COLMAP => OpenGL
+                all_c2w.append(c2w)
+                if self.split in ['train', 'val']:
+                    img_path = os.path.join(self.config.root_dir, 'images', d.name)
+                    img = Image.open(img_path)
+                    img = img.resize(img_wh, Image.BICUBIC)
+                    img = TF.to_tensor(img).permute(1, 2, 0)[...,:3]
+                    img = img.to(self.rank) if self.config.load_data_on_gpu else img.cpu()
+                    if has_mask:
+                        mask_paths = [os.path.join(mask_dir, d.name), os.path.join(mask_dir, d.name[3:])]
+                        mask_paths = list(filter(os.path.exists, mask_paths))
+                        assert len(mask_paths) == 1
+                        mask = Image.open(mask_paths[0]).convert('L') # (H, W, 1)
+                        mask = mask.resize(img_wh, Image.BICUBIC)
+                        mask = TF.to_tensor(mask)[0]
+                    else:
+                        mask = torch.ones_like(img[...,0], device=img.device)
+                    all_fg_masks.append(mask) # (h, w)
+                    all_images.append(img)
+            
+            all_c2w = torch.stack(all_c2w, dim=0)   
+
+            pts3d = read_points3d_binary(os.path.join(self.config.root_dir, 'sparse/0/points3D.bin'))
+            pts3d = torch.from_numpy(np.array([pts3d[k].xyz for k in pts3d])).float()
+            all_c2w, pts3d = normalize_poses(all_c2w, pts3d, up_est_method=self.config.up_est_method, center_est_method=self.config.center_est_method)
+
+            ColmapDatasetBase.properties = {
+                'w': w,
+                'h': h,
+                'img_wh': img_wh,
+                'factor': factor,
+                'has_mask': has_mask,
+                'apply_mask': apply_mask,
+                'directions': directions,
+                'pts3d': pts3d,
+                'all_c2w': all_c2w,
+                'all_images': all_images,
+                'all_fg_masks': all_fg_masks
+            }
+
+            ColmapDatasetBase.initialized = True
+        
+        for k, v in ColmapDatasetBase.properties.items():
+            setattr(self, k, v)
+
+        if self.split == 'test':
+            self.all_c2w = create_spheric_poses(self.all_c2w[:,:,3], n_steps=self.config.n_test_traj_steps)
+            self.all_images = torch.zeros((self.config.n_test_traj_steps, self.h, self.w, 3), dtype=torch.float32)
+            self.all_fg_masks = torch.zeros((self.config.n_test_traj_steps, self.h, self.w), dtype=torch.float32)
+        else:
+            self.all_images, self.all_fg_masks = torch.stack(self.all_images, dim=0).float(), torch.stack(self.all_fg_masks, dim=0).float()
+
+        """
+        # for debug use
+        from models.ray_utils import get_rays
+        rays_o, rays_d = get_rays(self.directions.cpu(), self.all_c2w, keepdim=True)
+        pts_out = []
+        pts_out.append('\n'.join([' '.join([str(p) for p in l]) + ' 1.0 0.0 0.0' for l in rays_o[:,0,0].reshape(-1, 3).tolist()]))
+
+        t_vals = torch.linspace(0, 1, 8)
+        z_vals = 0.05 * (1 - t_vals) + 0.5 * t_vals
+
+        ray_pts = (rays_o[:,0,0][..., None, :] + z_vals[..., None] * rays_d[:,0,0][..., None, :])
+        pts_out.append('\n'.join([' '.join([str(p) for p in l]) + ' 0.0 1.0 0.0' for l in ray_pts.view(-1, 3).tolist()]))
+
+        ray_pts = (rays_o[:,0,0][..., None, :] + z_vals[..., None] * rays_d[:,self.h-1,0][..., None, :])
+        pts_out.append('\n'.join([' '.join([str(p) for p in l]) + ' 0.0 0.0 1.0' for l in ray_pts.view(-1, 3).tolist()]))
+
+        ray_pts = (rays_o[:,0,0][..., None, :] + z_vals[..., None] * rays_d[:,0,self.w-1][..., None, :])
+        pts_out.append('\n'.join([' '.join([str(p) for p in l]) + ' 0.0 1.0 1.0' for l in ray_pts.view(-1, 3).tolist()]))
+
+        ray_pts = (rays_o[:,0,0][..., None, :] + z_vals[..., None] * rays_d[:,self.h-1,self.w-1][..., None, :])
+        pts_out.append('\n'.join([' '.join([str(p) for p in l]) + ' 1.0 1.0 1.0' for l in ray_pts.view(-1, 3).tolist()]))
+        
+        open('cameras.txt', 'w').write('\n'.join(pts_out))
+        open('scene.txt', 'w').write('\n'.join([' '.join([str(p) for p in l]) + ' 0.0 0.0 0.0' for l in self.pts3d.view(-1, 3).tolist()]))
+
+        exit(1)
+        """
+
+        self.all_c2w = self.all_c2w.float().to(self.rank)
+        if self.config.load_data_on_gpu:
+            self.all_images = self.all_images.to(self.rank) 
+            self.all_fg_masks = self.all_fg_masks.to(self.rank)
+        
+
+class ColmapDataset(Dataset, ColmapDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __len__(self):
+        return len(self.all_images)
+    
+    def __getitem__(self, index):
+        return {
+            'index': index
+        }
+
+
+class ColmapIterableDataset(IterableDataset, ColmapDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __iter__(self):
+        while True:
+            yield {}
+
+
+@datasets.register('colmap')
+class ColmapDataModule(pl.LightningDataModule):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+    
+    def setup(self, stage=None):
+        if stage in [None, 'fit']:
+            self.train_dataset = ColmapIterableDataset(self.config, 'train')
+        if stage in [None, 'fit', 'validate']:
+            self.val_dataset = ColmapDataset(self.config, self.config.get('val_split', 'train'))
+        if stage in [None, 'test']:
+            self.test_dataset = ColmapDataset(self.config, self.config.get('test_split', 'test'))
+        if stage in [None, 'predict']:
+            self.predict_dataset = ColmapDataset(self.config, 'train')         
+
+    def prepare_data(self):
+        pass
+    
+    def general_loader(self, dataset, batch_size):
+        sampler = None
+        return DataLoader(
+            dataset, 
+            num_workers=os.cpu_count(), 
+            batch_size=batch_size,
+            pin_memory=True,
+            sampler=sampler
+        )
+    
+    def train_dataloader(self):
+        return self.general_loader(self.train_dataset, batch_size=1)
+
+    def val_dataloader(self):
+        return self.general_loader(self.val_dataset, batch_size=1)
+
+    def test_dataloader(self):
+        return self.general_loader(self.test_dataset, batch_size=1) 
+
+    def predict_dataloader(self):
+        return self.general_loader(self.predict_dataset, batch_size=1)       

instant-nsr-pl/datasets/colmap_utils.py

@@ -0,0 +1,295 @@
+# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     * Redistributions of source code must retain the above copyright
+#       notice, this list of conditions and the following disclaimer.
+#
+#     * Redistributions in binary form must reproduce the above copyright
+#       notice, this list of conditions and the following disclaimer in the
+#       documentation and/or other materials provided with the distribution.
+#
+#     * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of
+#       its contributors may be used to endorse or promote products derived
+#       from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+# POSSIBILITY OF SUCH DAMAGE.
+#
+# Author: Johannes L. Schoenberger (jsch at inf.ethz.ch)
+
+import os
+import collections
+import numpy as np
+import struct
+
+
+CameraModel = collections.namedtuple(
+    "CameraModel", ["model_id", "model_name", "num_params"])
+Camera = collections.namedtuple(
+    "Camera", ["id", "model", "width", "height", "params"])
+BaseImage = collections.namedtuple(
+    "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"])
+Point3D = collections.namedtuple(
+    "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"])
+
+class Image(BaseImage):
+    def qvec2rotmat(self):
+        return qvec2rotmat(self.qvec)
+
+
+CAMERA_MODELS = {
+    CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3),
+    CameraModel(model_id=1, model_name="PINHOLE", num_params=4),
+    CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4),
+    CameraModel(model_id=3, model_name="RADIAL", num_params=5),
+    CameraModel(model_id=4, model_name="OPENCV", num_params=8),
+    CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8),
+    CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12),
+    CameraModel(model_id=7, model_name="FOV", num_params=5),
+    CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4),
+    CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5),
+    CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12)
+}
+CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) \
+                         for camera_model in CAMERA_MODELS])
+
+
+def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
+    """Read and unpack the next bytes from a binary file.
+    :param fid:
+    :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
+    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
+    :param endian_character: Any of {@, =, <, >, !}
+    :return: Tuple of read and unpacked values.
+    """
+    data = fid.read(num_bytes)
+    return struct.unpack(endian_character + format_char_sequence, data)
+
+
+def read_cameras_text(path):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::WriteCamerasText(const std::string& path)
+        void Reconstruction::ReadCamerasText(const std::string& path)
+    """
+    cameras = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                camera_id = int(elems[0])
+                model = elems[1]
+                width = int(elems[2])
+                height = int(elems[3])
+                params = np.array(tuple(map(float, elems[4:])))
+                cameras[camera_id] = Camera(id=camera_id, model=model,
+                                            width=width, height=height,
+                                            params=params)
+    return cameras
+
+
+def read_cameras_binary(path_to_model_file):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::WriteCamerasBinary(const std::string& path)
+        void Reconstruction::ReadCamerasBinary(const std::string& path)
+    """
+    cameras = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_cameras = read_next_bytes(fid, 8, "Q")[0]
+        for camera_line_index in range(num_cameras):
+            camera_properties = read_next_bytes(
+                fid, num_bytes=24, format_char_sequence="iiQQ")
+            camera_id = camera_properties[0]
+            model_id = camera_properties[1]
+            model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name
+            width = camera_properties[2]
+            height = camera_properties[3]
+            num_params = CAMERA_MODEL_IDS[model_id].num_params
+            params = read_next_bytes(fid, num_bytes=8*num_params,
+                                     format_char_sequence="d"*num_params)
+            cameras[camera_id] = Camera(id=camera_id,
+                                        model=model_name,
+                                        width=width,
+                                        height=height,
+                                        params=np.array(params))
+        assert len(cameras) == num_cameras
+    return cameras
+
+
+def read_images_text(path):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadImagesText(const std::string& path)
+        void Reconstruction::WriteImagesText(const std::string& path)
+    """
+    images = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                image_id = int(elems[0])
+                qvec = np.array(tuple(map(float, elems[1:5])))
+                tvec = np.array(tuple(map(float, elems[5:8])))
+                camera_id = int(elems[8])
+                image_name = elems[9]
+                elems = fid.readline().split()
+                xys = np.column_stack([tuple(map(float, elems[0::3])),
+                                       tuple(map(float, elems[1::3]))])
+                point3D_ids = np.array(tuple(map(int, elems[2::3])))
+                images[image_id] = Image(
+                    id=image_id, qvec=qvec, tvec=tvec,
+                    camera_id=camera_id, name=image_name,
+                    xys=xys, point3D_ids=point3D_ids)
+    return images
+
+
+def read_images_binary(path_to_model_file):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadImagesBinary(const std::string& path)
+        void Reconstruction::WriteImagesBinary(const std::string& path)
+    """
+    images = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_reg_images = read_next_bytes(fid, 8, "Q")[0]
+        for image_index in range(num_reg_images):
+            binary_image_properties = read_next_bytes(
+                fid, num_bytes=64, format_char_sequence="idddddddi")
+            image_id = binary_image_properties[0]
+            qvec = np.array(binary_image_properties[1:5])
+            tvec = np.array(binary_image_properties[5:8])
+            camera_id = binary_image_properties[8]
+            image_name = ""
+            current_char = read_next_bytes(fid, 1, "c")[0]
+            while current_char != b"\x00":   # look for the ASCII 0 entry
+                image_name += current_char.decode("utf-8")
+                current_char = read_next_bytes(fid, 1, "c")[0]
+            num_points2D = read_next_bytes(fid, num_bytes=8,
+                                           format_char_sequence="Q")[0]
+            x_y_id_s = read_next_bytes(fid, num_bytes=24*num_points2D,
+                                       format_char_sequence="ddq"*num_points2D)
+            xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])),
+                                   tuple(map(float, x_y_id_s[1::3]))])
+            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
+            images[image_id] = Image(
+                id=image_id, qvec=qvec, tvec=tvec,
+                camera_id=camera_id, name=image_name,
+                xys=xys, point3D_ids=point3D_ids)
+    return images
+
+
+def read_points3D_text(path):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadPoints3DText(const std::string& path)
+        void Reconstruction::WritePoints3DText(const std::string& path)
+    """
+    points3D = {}
+    with open(path, "r") as fid:
+        while True:
+            line = fid.readline()
+            if not line:
+                break
+            line = line.strip()
+            if len(line) > 0 and line[0] != "#":
+                elems = line.split()
+                point3D_id = int(elems[0])
+                xyz = np.array(tuple(map(float, elems[1:4])))
+                rgb = np.array(tuple(map(int, elems[4:7])))
+                error = float(elems[7])
+                image_ids = np.array(tuple(map(int, elems[8::2])))
+                point2D_idxs = np.array(tuple(map(int, elems[9::2])))
+                points3D[point3D_id] = Point3D(id=point3D_id, xyz=xyz, rgb=rgb,
+                                               error=error, image_ids=image_ids,
+                                               point2D_idxs=point2D_idxs)
+    return points3D
+
+
+def read_points3d_binary(path_to_model_file):
+    """
+    see: src/base/reconstruction.cc
+        void Reconstruction::ReadPoints3DBinary(const std::string& path)
+        void Reconstruction::WritePoints3DBinary(const std::string& path)
+    """
+    points3D = {}
+    with open(path_to_model_file, "rb") as fid:
+        num_points = read_next_bytes(fid, 8, "Q")[0]
+        for point_line_index in range(num_points):
+            binary_point_line_properties = read_next_bytes(
+                fid, num_bytes=43, format_char_sequence="QdddBBBd")
+            point3D_id = binary_point_line_properties[0]
+            xyz = np.array(binary_point_line_properties[1:4])
+            rgb = np.array(binary_point_line_properties[4:7])
+            error = np.array(binary_point_line_properties[7])
+            track_length = read_next_bytes(
+                fid, num_bytes=8, format_char_sequence="Q")[0]
+            track_elems = read_next_bytes(
+                fid, num_bytes=8*track_length,
+                format_char_sequence="ii"*track_length)
+            image_ids = np.array(tuple(map(int, track_elems[0::2])))
+            point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))
+            points3D[point3D_id] = Point3D(
+                id=point3D_id, xyz=xyz, rgb=rgb,
+                error=error, image_ids=image_ids,
+                point2D_idxs=point2D_idxs)
+    return points3D
+
+
+def read_model(path, ext):
+    if ext == ".txt":
+        cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
+        images = read_images_text(os.path.join(path, "images" + ext))
+        points3D = read_points3D_text(os.path.join(path, "points3D") + ext)
+    else:
+        cameras = read_cameras_binary(os.path.join(path, "cameras" + ext))
+        images = read_images_binary(os.path.join(path, "images" + ext))
+        points3D = read_points3d_binary(os.path.join(path, "points3D") + ext)
+    return cameras, images, points3D
+
+
+def qvec2rotmat(qvec):
+    return np.array([
+        [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
+         2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
+         2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]],
+        [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
+         1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
+         2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],
+        [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
+         2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
+         1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]])
+
+
+def rotmat2qvec(R):
+    Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
+    K = np.array([
+        [Rxx - Ryy - Rzz, 0, 0, 0],
+        [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
+        [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
+        [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0
+    eigvals, eigvecs = np.linalg.eigh(K)
+    qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
+    if qvec[0] < 0:
+        qvec *= -1
+    return qvec

instant-nsr-pl/datasets/dtu.py

@@ -0,0 +1,201 @@
+import os
+import json
+import math
+import numpy as np
+from PIL import Image
+import cv2
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, IterableDataset
+import torchvision.transforms.functional as TF
+
+import pytorch_lightning as pl
+
+import datasets
+from models.ray_utils import get_ray_directions
+from utils.misc import get_rank
+
+
+def load_K_Rt_from_P(P=None):
+    out = cv2.decomposeProjectionMatrix(P)
+    K = out[0]
+    R = out[1]
+    t = out[2]
+
+    K = K / K[2, 2]
+    intrinsics = np.eye(4)
+    intrinsics[:3, :3] = K
+
+    pose = np.eye(4, dtype=np.float32)
+    pose[:3, :3] = R.transpose()
+    pose[:3, 3] = (t[:3] / t[3])[:, 0]
+
+    return intrinsics, pose
+
+def create_spheric_poses(cameras, n_steps=120):
+    center = torch.as_tensor([0.,0.,0.], dtype=cameras.dtype, device=cameras.device)
+    cam_center = F.normalize(cameras.mean(0), p=2, dim=-1) * cameras.mean(0).norm(2)
+    eigvecs = torch.linalg.eig(cameras.T @ cameras).eigenvectors
+    rot_axis = F.normalize(eigvecs[:,1].real.float(), p=2, dim=-1)
+    up = rot_axis
+    rot_dir = torch.cross(rot_axis, cam_center)
+    max_angle = (F.normalize(cameras, p=2, dim=-1) * F.normalize(cam_center, p=2, dim=-1)).sum(-1).acos().max()
+
+    all_c2w = []
+    for theta in torch.linspace(-max_angle, max_angle, n_steps):
+        cam_pos = cam_center * math.cos(theta) + rot_dir * math.sin(theta)
+        l = F.normalize(center - cam_pos, p=2, dim=0)
+        s = F.normalize(l.cross(up), p=2, dim=0)
+        u = F.normalize(s.cross(l), p=2, dim=0)
+        c2w = torch.cat([torch.stack([s, u, -l], dim=1), cam_pos[:,None]], axis=1)
+        all_c2w.append(c2w)
+
+    all_c2w = torch.stack(all_c2w, dim=0)
+    
+    return all_c2w
+
+class DTUDatasetBase():
+    def setup(self, config, split):
+        self.config = config
+        self.split = split
+        self.rank = get_rank()
+
+        cams = np.load(os.path.join(self.config.root_dir, self.config.cameras_file))
+
+        img_sample = cv2.imread(os.path.join(self.config.root_dir, 'image', '000000.png'))
+        H, W = img_sample.shape[0], img_sample.shape[1]
+
+        if 'img_wh' in self.config:
+            w, h = self.config.img_wh
+            assert round(W / w * h) == H
+        elif 'img_downscale' in self.config:
+            w, h = int(W / self.config.img_downscale + 0.5), int(H / self.config.img_downscale + 0.5)
+        else:
+            raise KeyError("Either img_wh or img_downscale should be specified.")
+
+        self.w, self.h = w, h
+        self.img_wh = (w, h)
+        self.factor = w / W
+
+        mask_dir = os.path.join(self.config.root_dir, 'mask')
+        self.has_mask = True
+        self.apply_mask = self.config.apply_mask
+        
+        self.directions = []
+        self.all_c2w, self.all_images, self.all_fg_masks = [], [], []
+
+        n_images = max([int(k.split('_')[-1]) for k in cams.keys()]) + 1
+
+        for i in range(n_images):
+            world_mat, scale_mat = cams[f'world_mat_{i}'], cams[f'scale_mat_{i}']
+            P = (world_mat @ scale_mat)[:3,:4]
+            K, c2w = load_K_Rt_from_P(P)
+            fx, fy, cx, cy = K[0,0] * self.factor, K[1,1] * self.factor, K[0,2] * self.factor, K[1,2] * self.factor
+            directions = get_ray_directions(w, h, fx, fy, cx, cy)
+            self.directions.append(directions)
+            
+            c2w = torch.from_numpy(c2w).float()
+
+            # blender follows opengl camera coordinates (right up back)
+            # NeuS DTU data coordinate system (right down front) is different from blender
+            # https://github.com/Totoro97/NeuS/issues/9
+            # for c2w, flip the sign of input camera coordinate yz
+            c2w_ = c2w.clone()
+            c2w_[:3,1:3] *= -1. # flip input sign
+            self.all_c2w.append(c2w_[:3,:4])         
+
+            if self.split in ['train', 'val']:
+                img_path = os.path.join(self.config.root_dir, 'image', f'{i:06d}.png')
+                img = Image.open(img_path)
+                img = img.resize(self.img_wh, Image.BICUBIC)
+                img = TF.to_tensor(img).permute(1, 2, 0)[...,:3]
+
+                mask_path = os.path.join(mask_dir, f'{i:03d}.png')
+                mask = Image.open(mask_path).convert('L') # (H, W, 1)
+                mask = mask.resize(self.img_wh, Image.BICUBIC)
+                mask = TF.to_tensor(mask)[0]
+
+                self.all_fg_masks.append(mask) # (h, w)
+                self.all_images.append(img)
+
+        self.all_c2w = torch.stack(self.all_c2w, dim=0)
+
+        if self.split == 'test':
+            self.all_c2w = create_spheric_poses(self.all_c2w[:,:,3], n_steps=self.config.n_test_traj_steps)
+            self.all_images = torch.zeros((self.config.n_test_traj_steps, self.h, self.w, 3), dtype=torch.float32)
+            self.all_fg_masks = torch.zeros((self.config.n_test_traj_steps, self.h, self.w), dtype=torch.float32)
+            self.directions = self.directions[0]
+        else:
+            self.all_images, self.all_fg_masks = torch.stack(self.all_images, dim=0), torch.stack(self.all_fg_masks, dim=0)  
+            self.directions = torch.stack(self.directions, dim=0)
+
+        self.directions = self.directions.float().to(self.rank)
+        self.all_c2w, self.all_images, self.all_fg_masks = \
+            self.all_c2w.float().to(self.rank), \
+            self.all_images.float().to(self.rank), \
+            self.all_fg_masks.float().to(self.rank)
+        
+
+class DTUDataset(Dataset, DTUDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __len__(self):
+        return len(self.all_images)
+    
+    def __getitem__(self, index):
+        return {
+            'index': index
+        }
+
+
+class DTUIterableDataset(IterableDataset, DTUDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __iter__(self):
+        while True:
+            yield {}
+
+
+@datasets.register('dtu')
+class DTUDataModule(pl.LightningDataModule):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+    
+    def setup(self, stage=None):
+        if stage in [None, 'fit']:
+            self.train_dataset = DTUIterableDataset(self.config, 'train')
+        if stage in [None, 'fit', 'validate']:
+            self.val_dataset = DTUDataset(self.config, self.config.get('val_split', 'train'))
+        if stage in [None, 'test']:
+            self.test_dataset = DTUDataset(self.config, self.config.get('test_split', 'test'))
+        if stage in [None, 'predict']:
+            self.predict_dataset = DTUDataset(self.config, 'train')    
+
+    def prepare_data(self):
+        pass
+    
+    def general_loader(self, dataset, batch_size):
+        sampler = None
+        return DataLoader(
+            dataset, 
+            num_workers=os.cpu_count(), 
+            batch_size=batch_size,
+            pin_memory=True,
+            sampler=sampler
+        )
+    
+    def train_dataloader(self):
+        return self.general_loader(self.train_dataset, batch_size=1)
+
+    def val_dataloader(self):
+        return self.general_loader(self.val_dataset, batch_size=1)
+
+    def test_dataloader(self):
+        return self.general_loader(self.test_dataset, batch_size=1) 
+
+    def predict_dataloader(self):
+        return self.general_loader(self.predict_dataset, batch_size=1)       

instant-nsr-pl/datasets/fixed_poses/000_back_RT.txt

@@ -0,0 +1,3 @@
+-1.000000238418579102e+00 0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00
+0.000000000000000000e+00 -1.343588564850506373e-07 1.000000119209289551e+00 1.746665105883948854e-07
+0.000000000000000000e+00 1.000000119209289551e+00 -1.343588564850506373e-07 -1.300000071525573730e+00

instant-nsr-pl/datasets/fixed_poses/000_back_left_RT.txt

@@ -0,0 +1,3 @@
+-7.071069478988647461e-01 -7.071068286895751953e-01 0.000000000000000000e+00 -1.192092895507812500e-07
+0.000000000000000000e+00 -7.587616579485256807e-08 1.000000119209289551e+00 9.863901340168013121e-08
+-7.071068286895751953e-01 7.071068286895751953e-01 -7.587616579485256807e-08 -1.838477730751037598e+00

instant-nsr-pl/datasets/fixed_poses/000_back_right_RT.txt

@@ -0,0 +1,3 @@
+-7.071069478988647461e-01 7.071068286895751953e-01 0.000000000000000000e+00 1.192092895507812500e-07
+0.000000000000000000e+00 -7.587616579485256807e-08 1.000000119209289551e+00 9.863901340168013121e-08
+7.071068286895751953e-01 7.071068286895751953e-01 -7.587616579485256807e-08 -1.838477730751037598e+00

instant-nsr-pl/datasets/fixed_poses/000_front_RT.txt

@@ -0,0 +1,3 @@
+1.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00
+0.000000000000000000e+00 -1.343588564850506373e-07 1.000000119209289551e+00 -1.746665105883948854e-07
+0.000000000000000000e+00 -1.000000119209289551e+00 -1.343588564850506373e-07 -1.300000071525573730e+00

instant-nsr-pl/datasets/fixed_poses/000_front_left_RT.txt

@@ -0,0 +1,3 @@
+7.071067690849304199e-01 -7.071068286895751953e-01 0.000000000000000000e+00 -1.192092895507812500e-07
+0.000000000000000000e+00 -7.587616579485256807e-08 1.000000119209289551e+00 -9.863901340168013121e-08
+-7.071068286895751953e-01 -7.071068286895751953e-01 -7.587616579485256807e-08 -1.838477730751037598e+00

instant-nsr-pl/datasets/fixed_poses/000_front_right_RT.txt

@@ -0,0 +1,3 @@
+7.071067690849304199e-01 7.071068286895751953e-01 0.000000000000000000e+00 1.192092895507812500e-07
+0.000000000000000000e+00 -7.587616579485256807e-08 1.000000119209289551e+00 -9.863901340168013121e-08
+7.071068286895751953e-01 -7.071068286895751953e-01 -7.587616579485256807e-08 -1.838477730751037598e+00

instant-nsr-pl/datasets/fixed_poses/000_left_RT.txt

@@ -0,0 +1,3 @@
+-2.220446049250313081e-16 -1.000000000000000000e+00 0.000000000000000000e+00 -2.886579758146288598e-16
+0.000000000000000000e+00 -2.220446049250313081e-16 1.000000000000000000e+00 0.000000000000000000e+00
+-1.000000000000000000e+00 0.000000000000000000e+00 -2.220446049250313081e-16 -1.299999952316284180e+00

instant-nsr-pl/datasets/fixed_poses/000_right_RT.txt

@@ -0,0 +1,3 @@
+-2.220446049250313081e-16 1.000000000000000000e+00 0.000000000000000000e+00 2.886579758146288598e-16
+0.000000000000000000e+00 -2.220446049250313081e-16 1.000000000000000000e+00 0.000000000000000000e+00
+1.000000000000000000e+00 0.000000000000000000e+00 -2.220446049250313081e-16 -1.299999952316284180e+00

instant-nsr-pl/datasets/fixed_poses/000_top_RT.txt

@@ -0,0 +1,3 @@
+1.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00
+0.000000000000000000e+00 1.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00
+0.000000000000000000e+00 0.000000000000000000e+00 1.000000000000000000e+00 -1.299999952316284180e+00

instant-nsr-pl/datasets/ortho.py

@@ -0,0 +1,287 @@
+import os
+import json
+import math
+import numpy as np
+from PIL import Image
+import cv2
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, IterableDataset
+import torchvision.transforms.functional as TF
+
+import pytorch_lightning as pl
+
+import datasets
+from models.ray_utils import get_ortho_ray_directions_origins, get_ortho_rays, get_ray_directions
+from utils.misc import get_rank
+
+from glob import glob
+import PIL.Image
+
+
+def camNormal2worldNormal(rot_c2w, camNormal):
+    H,W,_ = camNormal.shape
+    normal_img = np.matmul(rot_c2w[None, :, :], camNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
+
+    return normal_img
+
+def worldNormal2camNormal(rot_w2c, worldNormal):
+    H,W,_ = worldNormal.shape
+    normal_img = np.matmul(rot_w2c[None, :, :], worldNormal.reshape(-1,3)[:, :, None]).reshape([H, W, 3])
+
+    return normal_img
+
+def trans_normal(normal, RT_w2c, RT_w2c_target):
+
+    normal_world = camNormal2worldNormal(np.linalg.inv(RT_w2c[:3,:3]), normal)
+    normal_target_cam = worldNormal2camNormal(RT_w2c_target[:3,:3], normal_world)
+
+    return normal_target_cam
+
+def img2normal(img):
+    return (img/255.)*2-1
+
+def normal2img(normal):
+    return np.uint8((normal*0.5+0.5)*255)
+
+def norm_normalize(normal, dim=-1):
+
+    normal = normal/(np.linalg.norm(normal, axis=dim, keepdims=True)+1e-6)
+
+    return normal
+
+def RT_opengl2opencv(RT):
+     # Build the coordinate transform matrix from world to computer vision camera
+    # R_world2cv = R_bcam2cv@R_world2bcam
+    # T_world2cv = R_bcam2cv@T_world2bcam
+
+    R = RT[:3, :3]
+    t = RT[:3, 3]
+
+    R_bcam2cv = np.asarray([[1, 0, 0], [0, -1, 0], [0, 0, -1]], np.float32)
+
+    R_world2cv = R_bcam2cv @ R
+    t_world2cv = R_bcam2cv @ t
+
+    RT = np.concatenate([R_world2cv,t_world2cv[:,None]],1)
+
+    return RT
+
+def normal_opengl2opencv(normal):
+    H,W,C = np.shape(normal)
+    # normal_img = np.reshape(normal, (H*W,C))
+    R_bcam2cv = np.array([1, -1, -1], np.float32)
+    normal_cv = normal * R_bcam2cv[None, None, :]
+
+    print(np.shape(normal_cv))
+
+    return normal_cv
+
+def inv_RT(RT):
+    RT_h = np.concatenate([RT, np.array([[0,0,0,1]])], axis=0)
+    RT_inv = np.linalg.inv(RT_h)
+
+    return RT_inv[:3, :]
+
+
+def load_a_prediction(root_dir, test_object, imSize, view_types, load_color=False, cam_pose_dir=None,
+                         normal_system='front', erode_mask=True, camera_type='ortho', cam_params=None):
+
+    all_images = []
+    all_normals = []
+    all_normals_world = []
+    all_masks = []
+    all_color_masks = []
+    all_poses = []
+    all_w2cs = []
+    directions = []
+    ray_origins = []
+
+    RT_front = np.loadtxt(glob(os.path.join(cam_pose_dir, '*_%s_RT.txt'%( 'front')))[0])   # world2cam matrix
+    RT_front_cv = RT_opengl2opencv(RT_front)   # convert normal from opengl to opencv
+    for idx, view in enumerate(view_types):
+        print(os.path.join(root_dir,test_object))
+        normal_filepath = os.path.join(root_dir, test_object, 'normals_000_%s.png'%( view))
+        # Load key frame
+        if load_color:  # use bgr
+            image =np.array(PIL.Image.open(normal_filepath.replace("normals", "rgb")).resize(imSize))[:, :, :3]
+
+        normal = np.array(PIL.Image.open(normal_filepath).resize(imSize))
+        mask = normal[:, :, 3]
+        normal = normal[:, :, :3]
+
+        color_mask = np.array(PIL.Image.open(os.path.join(root_dir,test_object, 'masked_colors/rgb_000_%s.png'%( view))).resize(imSize))[:, :, 3]
+        invalid_color_mask = color_mask < 255*0.5
+        threshold =  np.ones_like(image[:, :, 0]) * 250
+        invalid_white_mask = (image[:, :, 0] > threshold) & (image[:, :, 1] > threshold) & (image[:, :, 2] > threshold)
+        invalid_color_mask_final = invalid_color_mask & invalid_white_mask
+        color_mask = (1 - invalid_color_mask_final) > 0
+
+        # if erode_mask:
+        #     kernel = np.ones((3, 3), np.uint8)
+        #     mask = cv2.erode(mask, kernel, iterations=1)
+
+        RT = np.loadtxt(os.path.join(cam_pose_dir, '000_%s_RT.txt'%( view)))  # world2cam matrix
+
+        normal = img2normal(normal)
+
+        normal[mask==0] = [0,0,0]
+        mask = mask> (0.5*255)
+        if load_color:
+            all_images.append(image)
+        
+        all_masks.append(mask)
+        all_color_masks.append(color_mask)
+        RT_cv = RT_opengl2opencv(RT)   # convert normal from opengl to opencv
+        all_poses.append(inv_RT(RT_cv))   # cam2world
+        all_w2cs.append(RT_cv)
+
+        # whether to 
+        normal_cam_cv = normal_opengl2opencv(normal)
+
+        if normal_system == 'front':
+            print("the loaded normals are defined in the system of front view")
+            normal_world = camNormal2worldNormal(inv_RT(RT_front_cv)[:3, :3], normal_cam_cv)
+        elif normal_system == 'self':
+            print("the loaded normals are in their independent camera systems")
+            normal_world = camNormal2worldNormal(inv_RT(RT_cv)[:3, :3], normal_cam_cv)
+        all_normals.append(normal_cam_cv)
+        all_normals_world.append(normal_world)
+
+        if camera_type == 'ortho':
+            origins, dirs = get_ortho_ray_directions_origins(W=imSize[0], H=imSize[1])
+        elif camera_type == 'pinhole':
+            dirs = get_ray_directions(W=imSize[0], H=imSize[1],
+                                                 fx=cam_params[0], fy=cam_params[1], cx=cam_params[2], cy=cam_params[3])
+            origins = dirs # occupy a position
+        else:
+            raise Exception("not support camera type")
+        ray_origins.append(origins)
+        directions.append(dirs)
+        
+        
+        if not load_color:
+            all_images = [normal2img(x) for x in all_normals_world]
+
+
+    return np.stack(all_images), np.stack(all_masks), np.stack(all_normals), \
+        np.stack(all_normals_world), np.stack(all_poses), np.stack(all_w2cs), np.stack(ray_origins), np.stack(directions), np.stack(all_color_masks)
+
+
+class OrthoDatasetBase():
+    def setup(self, config, split):
+        self.config = config
+        self.split = split
+        self.rank = get_rank()
+
+        self.data_dir = self.config.root_dir
+        self.object_name = self.config.scene
+        self.scene = self.config.scene
+        self.imSize = self.config.imSize
+        self.load_color = True
+        self.img_wh = [self.imSize[0], self.imSize[1]]
+        self.w = self.img_wh[0]
+        self.h = self.img_wh[1]
+        self.camera_type = self.config.camera_type
+        self.camera_params = self.config.camera_params  # [fx, fy, cx, cy]
+        
+        self.view_types = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
+
+        self.view_weights = torch.from_numpy(np.array(self.config.view_weights)).float().to(self.rank).view(-1)
+        self.view_weights = self.view_weights.view(-1,1,1).repeat(1, self.h, self.w)
+
+        if self.config.cam_pose_dir is None:
+            self.cam_pose_dir = "./datasets/fixed_poses"
+        else:
+            self.cam_pose_dir = self.config.cam_pose_dir
+            
+        self.images_np, self.masks_np, self.normals_cam_np, self.normals_world_np, \
+            self.pose_all_np, self.w2c_all_np, self.origins_np, self.directions_np, self.rgb_masks_np = load_a_prediction(
+                self.data_dir, self.object_name, self.imSize, self.view_types,
+                self.load_color, self.cam_pose_dir, normal_system='front', 
+                camera_type=self.camera_type, cam_params=self.camera_params)
+
+        self.has_mask = True
+        self.apply_mask = self.config.apply_mask
+
+        self.all_c2w = torch.from_numpy(self.pose_all_np)
+        self.all_images = torch.from_numpy(self.images_np) / 255.
+        self.all_fg_masks = torch.from_numpy(self.masks_np)
+        self.all_rgb_masks = torch.from_numpy(self.rgb_masks_np)
+        self.all_normals_world = torch.from_numpy(self.normals_world_np)
+        self.origins = torch.from_numpy(self.origins_np)
+        self.directions = torch.from_numpy(self.directions_np)
+
+        self.directions = self.directions.float().to(self.rank)
+        self.origins = self.origins.float().to(self.rank)
+        self.all_rgb_masks = self.all_rgb_masks.float().to(self.rank)
+        self.all_c2w, self.all_images, self.all_fg_masks, self.all_normals_world = \
+            self.all_c2w.float().to(self.rank), \
+            self.all_images.float().to(self.rank), \
+            self.all_fg_masks.float().to(self.rank), \
+            self.all_normals_world.float().to(self.rank)
+        
+
+class OrthoDataset(Dataset, OrthoDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __len__(self):
+        return len(self.all_images)
+    
+    def __getitem__(self, index):
+        return {
+            'index': index
+        }
+
+
+class OrthoIterableDataset(IterableDataset, OrthoDatasetBase):
+    def __init__(self, config, split):
+        self.setup(config, split)
+
+    def __iter__(self):
+        while True:
+            yield {}
+
+
+@datasets.register('ortho')
+class OrthoDataModule(pl.LightningDataModule):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+    
+    def setup(self, stage=None):
+        if stage in [None, 'fit']:
+            self.train_dataset = OrthoIterableDataset(self.config, 'train')
+        if stage in [None, 'fit', 'validate']:
+            self.val_dataset = OrthoDataset(self.config, self.config.get('val_split', 'train'))
+        if stage in [None, 'test']:
+            self.test_dataset = OrthoDataset(self.config, self.config.get('test_split', 'test'))
+        if stage in [None, 'predict']:
+            self.predict_dataset = OrthoDataset(self.config, 'train')    
+
+    def prepare_data(self):
+        pass
+    
+    def general_loader(self, dataset, batch_size):
+        sampler = None
+        return DataLoader(
+            dataset, 
+            num_workers=os.cpu_count(), 
+            batch_size=batch_size,
+            pin_memory=True,
+            sampler=sampler
+        )
+    
+    def train_dataloader(self):
+        return self.general_loader(self.train_dataset, batch_size=1)
+
+    def val_dataloader(self):
+        return self.general_loader(self.val_dataset, batch_size=1)
+
+    def test_dataloader(self):
+        return self.general_loader(self.test_dataset, batch_size=1) 
+
+    def predict_dataloader(self):
+        return self.general_loader(self.predict_dataset, batch_size=1)