vid2avatar baseline

.gitattributes

@@ -33,3 +33,46 @@ 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/exstrimalik.gif filter=lfs diff=lfs merge=lfs -text
+assets/martial.gif filter=lfs diff=lfs merge=lfs -text
+assets/parkinglot_360.gif filter=lfs diff=lfs merge=lfs -text
+assets/roger.gif filter=lfs diff=lfs merge=lfs -text
+assets/teaser.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0000.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0001.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0002.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0003.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0004.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0005.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0006.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0007.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0010.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0011.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0012.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0013.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0014.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0015.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0016.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0017.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0018.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0020.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0021.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0022.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0023.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0024.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0025.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0027.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0028.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0029.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0030.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0031.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0032.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0033.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0034.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0035.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0036.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0037.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0038.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0039.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0040.png filter=lfs diff=lfs merge=lfs -text
+data/parkinglot/image/0041.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,19 @@
+.vscode/
+env/
+__pycache__/
+*.ply
+*.npz
+*.npy
+*.pkl
+outputs
+*.obj
+*.ipynb
+*.so
+data
+code/UNKNOWN.egg-info
+code/dist
+code/build
+visualization/imgui.ini
+export
+preprocessing/raw_data
+preprocessing/romp

LICENSE

@@ -0,0 +1,399 @@
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+	wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More_considerations
+     for the public: 
+	wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

README.md

@@ -1,11 +1,74 @@
----
-title: IF3D
-emoji: 👁
-colorFrom: blue
-colorTo: pink
-sdk: docker
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Vid2Avatar: 3D Avatar Reconstruction from Videos in the Wild via Self-supervised Scene Decomposition
+## [Paper](https://arxiv.org/abs/2302.11566) | [Video Youtube](https://youtu.be/EGi47YeIeGQ) | [Project Page](https://moygcc.github.io/vid2avatar/) | [SynWild Data](https://synwild.ait.ethz.ch/)
+
+
+Official Repository for CVPR 2023 paper [*Vid2Avatar: 3D Avatar Reconstruction from Videos in the Wild via Self-supervised Scene Decomposition*](https://arxiv.org/abs/2302.11566). 
+
+<img src="assets/teaser.png" width="800" height="223"/> 
+
+## Getting Started
+* Clone this repo: `git clone https://github.com/MoyGcc/vid2avatar`
+* Create a python virtual environment and activate. `conda create -n v2a python=3.7` and `conda activate v2a`
+* Install dependenices. `cd vid2avatar`, `pip install -r requirement.txt` and `cd code; python setup.py develop`
+* Install [Kaolin](https://kaolin.readthedocs.io/en/v0.10.0/notes/installation.html). We use version 0.10.0.
+* Download [SMPL model](https://smpl.is.tue.mpg.de/download.php) (1.0.0 for Python 2.7 (10 shape PCs)) and move them to the corresponding places:
+```
+mkdir code/lib/smpl/smpl_model/
+mv /path/to/smpl/models/basicModel_f_lbs_10_207_0_v1.0.0.pkl code/lib/smpl/smpl_model/SMPL_FEMALE.pkl
+mv /path/to/smpl/models/basicmodel_m_lbs_10_207_0_v1.0.0.pkl code/lib/smpl/smpl_model/SMPL_MALE.pkl
+```
+## Download preprocessed demo data
+You can quickly start trying out Vid2Avatar with a preprocessed demo sequence including the pre-trained checkpoint. This can be downloaded from [Google drive](https://drive.google.com/drive/u/1/folders/1AUtKSmib7CvpWBCFO6mQ9spVrga_CTU4) which is originally a video clip provided by [NeuMan](https://github.com/apple/ml-neuman). Put this preprocessed demo data under the folder `data/` and put the folder `checkpoints` under `outputs/parkinglot/`.
+
+## Training
+Before training, make sure that the `metaninfo` in the data config file `/code/confs/dataset/video.yaml` does match the expected training video. You can also continue the training by changing the flag `is_continue` in the model config file `code/confs/model/model_w_bg`. And then run:
+```
+cd code
+python train.py
+```
+The training usually takes 24-48 hours. The validation results can be found at `outputs/`.
+## Test
+Run the following command to obtain the final outputs. By default, this loads the latest checkpoint.
+```
+cd code
+python test.py
+```
+## 3D Visualization
+We use [AITViewer](https://github.com/eth-ait/aitviewer) to visualize the human models in 3D. First install AITViewer: `pip install aitviewer imgui==1.4.1`, and then run the following command to visualize the canonical mesh (--mode static) or deformed mesh sequence (--mode dynamic):
+```
+cd visualization 
+python vis.py --mode {MODE} --path {PATH}
+```
+<p align="center">
+  <img src="assets/parkinglot_360.gif" width="623" height="346"/> 
+</p>
+
+## Play on custom video
+* We use [ROMP](https://github.com/Arthur151/ROMP#installation) to obtain initial SMPL shape and poses: `pip install --upgrade simple-romp`
+* Install [OpenPose](https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/installation/0_index.md) as well as the python bindings.
+* Put the video frames under the folder `preprocessing/raw_data/{SEQUENCE_NAME}/frames`
+* Modify the preprocessing script `preprocessing/run_preprocessing.sh` accordingly: the data source, sequence name, and the gender. The data source is by default "custom" which will estimate camera intrinsics. If the camera intrinsics are known, it's better if the true camera parameters can be given.
+* Run preprocessing: `cd preprocessing` and `bash run_preprocessing.sh`. The processed data will be stored in `data/`. The intermediate outputs of the preprocessing can be found at `preprocessing/raw_data/{SEQUENCE_NAME}/`
+* Launch training and test in the same way as above. The `metainfo` in the data config file `/code/confs/dataset/video.yaml` should be changed according to the custom video.
+
+<p align="center">
+  <img src="assets/roger.gif" width="240" height="270"/>  <img src="assets/exstrimalik.gif" width="240" height="270"/> <img src="assets/martial.gif" width="240" height="270"/>
+</p>
+
+## Acknowledgement
+We have used codes from other great research work, including [VolSDF](https://github.com/lioryariv/volsdf), [NeRF++](https://github.com/Kai-46/nerfplusplus), [SMPL-X](https://github.com/vchoutas/smplx), [Anim-NeRF](https://github.com/JanaldoChen/Anim-NeRF), [I M Avatar](https://github.com/zhengyuf/IMavatar) and [SNARF](https://github.com/xuchen-ethz/snarf). We sincerely thank the authors for their awesome work! We also thank the authors of [ICON](https://github.com/YuliangXiu/ICON) and [SelfRecon](https://github.com/jby1993/SelfReconCode) for discussing experiment.
+
+## Related Works 
+Here are more recent related human body reconstruction projects from our team:
+* [Jiang and Chen et. al. - InstantAvatar: Learning Avatars from Monocular Video in 60 Seconds](https://github.com/tijiang13/InstantAvatar)
+* [Shen and Guo et. al. - X-Avatar: Expressive Human Avatars](https://skype-line.github.io/projects/X-Avatar/)
+* [Yin et. al. - Hi4D: 4D Instance Segmentation of Close Human Interaction](https://yifeiyin04.github.io/Hi4D/)
+
+```
+@inproceedings{guo2023vid2avatar,
+      title={Vid2Avatar: 3D Avatar Reconstruction from Videos in the Wild via Self-supervised Scene Decomposition},
+      author={Guo, Chen and Jiang, Tianjian and Chen, Xu and Song, Jie and Hilliges, Otmar},    
+      booktitle = {Computer Vision and Pattern Recognition (CVPR)},
+      year      = {2023}
+    }
+```

assets/smpl_init.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93541fbf3eb32ade3201565d4b0d793c851aa4173b58a46f1e62f3da292037ce
+size 2415862

code/check_cuda.py

@@ -0,0 +1,11 @@
+import torch
+
+print("Number of GPUs:", torch.cuda.device_count())
+
+print("Torch version:",torch.__version__)
+
+print("Is CUDA enabled?",torch.cuda.is_available())
+
+print(torch.cuda.device_count())
+
+# pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html

code/confs/base.yaml

@@ -0,0 +1,13 @@
+hydra:
+    run:
+        dir: "../outputs/${exp}/${run}"
+
+defaults:
+  - model: model_w_bg
+  - dataset: video
+  - _self_
+
+seed: 42
+project_name: "model_w_bg"
+exp: ${dataset.train.type}
+run: ${dataset.metainfo.subject}

code/confs/dataset/video.yaml

@@ -0,0 +1,37 @@
+metainfo:
+    gender: 'male'
+    data_dir : C:\Users\leob3\vid2avatar\data\parkinglot
+    subject: "parkinglot"
+    start_frame: 0
+    end_frame: 42
+
+train:
+    type: "Video"
+    batch_size: 1
+    drop_last: False
+    shuffle: True
+    worker: 8
+
+    num_sample : 512 
+
+valid:
+    type: "VideoVal"
+    image_id: 0
+    batch_size: 1
+    drop_last: False
+    shuffle: False
+    worker: 8
+
+    num_sample : -1
+    pixel_per_batch: 2048 
+
+test:
+    type: "VideoTest"
+    image_id: 0
+    batch_size: 1
+    drop_last: False
+    shuffle: False
+    worker: 8
+
+    num_sample : -1
+    pixel_per_batch: 2048

code/confs/model/model_w_bg.yaml

@@ -0,0 +1,77 @@
+learning_rate : 5.0e-4
+sched_milestones : [200,500]
+sched_factor : 0.5
+smpl_init: True
+is_continue: False
+use_body_parsing: False
+with_bkgd: True
+using_inpainting: False
+use_smpl_deformer: True
+use_bbox_sampler: False
+
+implicit_network:
+    feature_vector_size: 256
+    d_in: 3
+    d_out: 1
+    dims: [ 256, 256, 256, 256, 256, 256, 256, 256 ]
+    init: 'geometry'
+    bias: 0.6
+    skip_in: [4]
+    weight_norm: True
+    embedder_mode: 'fourier'
+    multires: 6
+    cond: 'smpl'
+    scene_bounding_sphere: 3.0
+rendering_network:
+    feature_vector_size: 256
+    mode: "pose"
+    d_in:  14
+    d_out: 3
+    dims: [ 256, 256, 256, 256]
+    weight_norm: True
+    multires_view: -1
+bg_implicit_network:
+    feature_vector_size: 256
+    d_in: 4
+    d_out: 1
+    dims: [ 256, 256, 256, 256, 256, 256, 256, 256 ]
+    init: 'none'
+    bias: 0.0
+    skip_in: [4]
+    weight_norm: False
+    embedder_mode: 'fourier'
+    multires: 10
+    cond: 'frame'
+    dim_frame_encoding: 32
+bg_rendering_network:
+    feature_vector_size: 256
+    mode: 'nerf_frame_encoding'
+    d_in: 3
+    d_out: 3
+    dims: [128]
+    weight_norm: False
+    multires_view: 4
+    dim_frame_encoding: 32
+shadow_network:
+    d_in: 3
+    d_out: 1
+    dims: [128, 128]
+    weight_norm: False
+density:
+    params_init: {beta: 0.1}
+    beta_min: 0.0001
+ray_sampler:
+    near: 0.0
+    N_samples: 64
+    N_samples_eval: 128
+    N_samples_extra: 32
+    eps: 0.1
+    beta_iters: 10
+    max_total_iters: 5
+    N_samples_inverse_sphere: 32
+    add_tiny: 1.0e-6
+loss:
+    eikonal_weight : 0.1
+    bce_weight: 5.0e-3
+    opacity_sparse_weight: 3.0e-3
+    in_shape_weight: 1.0e-2

code/lib/datasets/__init__.py

@@ -0,0 +1,26 @@
+from .dataset import Dataset, ValDataset, TestDataset
+from torch.utils.data import DataLoader
+
+def find_dataset_using_name(name):
+    mapping = {
+        "Video": Dataset,
+        "VideoVal": ValDataset,
+        "VideoTest": TestDataset,
+    }
+    cls = mapping.get(name, None)
+    if cls is None:
+        raise ValueError(f"Fail to find dataset {name}") 
+    return cls
+
+
+def create_dataset(metainfo, split):
+    dataset_cls = find_dataset_using_name(split.type)
+    dataset = dataset_cls(metainfo, split)
+    return DataLoader(
+        dataset,
+        batch_size=split.batch_size,
+        drop_last=split.drop_last,
+        shuffle=split.shuffle,
+        num_workers=split.worker,
+        pin_memory=True
+    )

code/lib/datasets/dataset.py

@@ -0,0 +1,175 @@
+import os
+import glob
+import hydra
+import cv2
+import numpy as np
+import torch
+from lib.utils import utils
+
+
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self, metainfo, split):
+        root = os.path.join("../data", metainfo.data_dir)
+        root = hydra.utils.to_absolute_path(root)
+
+        self.start_frame = metainfo.start_frame
+        self.end_frame = metainfo.end_frame 
+        self.skip_step = 1
+        self.images, self.img_sizes = [], []
+        self.training_indices = list(range(metainfo.start_frame, metainfo.end_frame, self.skip_step))
+
+        # images
+        img_dir = os.path.join(root, "image")
+        self.img_paths = sorted(glob.glob(f"{img_dir}/*.png"))
+
+        # only store the image paths to avoid OOM
+        self.img_paths = [self.img_paths[i] for i in self.training_indices]
+        self.img_size = cv2.imread(self.img_paths[0]).shape[:2]
+        self.n_images = len(self.img_paths)
+
+        # coarse projected SMPL masks, only for sampling
+        mask_dir = os.path.join(root, "mask")
+        self.mask_paths = sorted(glob.glob(f"{mask_dir}/*.png"))
+        self.mask_paths = [self.mask_paths[i] for i in self.training_indices]
+
+        self.shape = np.load(os.path.join(root, "mean_shape.npy"))
+        self.poses = np.load(os.path.join(root, 'poses.npy'))[self.training_indices]
+        self.trans = np.load(os.path.join(root, 'normalize_trans.npy'))[self.training_indices]
+        # cameras
+        camera_dict = np.load(os.path.join(root, "cameras_normalize.npz"))
+        scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in self.training_indices]
+        world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in self.training_indices]
+
+        self.scale = 1 / scale_mats[0][0, 0]
+
+        self.intrinsics_all = []
+        self.pose_all = []
+        for scale_mat, world_mat in zip(scale_mats, world_mats):
+            P = world_mat @ scale_mat
+            P = P[:3, :4]
+            intrinsics, pose = utils.load_K_Rt_from_P(None, P)
+            self.intrinsics_all.append(torch.from_numpy(intrinsics).float())
+            self.pose_all.append(torch.from_numpy(pose).float())
+        assert len(self.intrinsics_all) == len(self.pose_all)
+
+        # other properties
+        self.num_sample = split.num_sample
+        self.sampling_strategy = "weighted"
+
+    def __len__(self):
+        return self.n_images
+
+    def __getitem__(self, idx):
+        # normalize RGB
+        img = cv2.imread(self.img_paths[idx])
+        # preprocess: BGR -> RGB -> Normalize
+
+        img = img[:, :, ::-1] / 255
+
+        mask = cv2.imread(self.mask_paths[idx])
+        # preprocess: BGR -> Gray -> Mask
+        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY) > 0
+
+        img_size = self.img_size
+
+        uv = np.mgrid[:img_size[0], :img_size[1]].astype(np.int32)
+        uv = np.flip(uv, axis=0).copy().transpose(1, 2, 0).astype(np.float32)
+
+        smpl_params = torch.zeros([86]).float()
+        smpl_params[0] = torch.from_numpy(np.asarray(self.scale)).float() 
+
+        smpl_params[1:4] = torch.from_numpy(self.trans[idx]).float()
+        smpl_params[4:76] = torch.from_numpy(self.poses[idx]).float()
+        smpl_params[76:] = torch.from_numpy(self.shape).float()
+
+        if self.num_sample > 0:
+            data = {
+                "rgb": img,
+                "uv": uv,
+                "object_mask": mask,
+            }
+
+            samples, index_outside = utils.weighted_sampling(data, img_size, self.num_sample)
+            inputs = {
+                "uv": samples["uv"].astype(np.float32),
+                "intrinsics": self.intrinsics_all[idx],
+                "pose": self.pose_all[idx],
+                "smpl_params": smpl_params,
+                'index_outside': index_outside,
+                "idx": idx
+            }
+            images = {"rgb": samples["rgb"].astype(np.float32)}
+            return inputs, images
+        else:
+            inputs = {
+                "uv": uv.reshape(-1, 2).astype(np.float32),
+                "intrinsics": self.intrinsics_all[idx],
+                "pose": self.pose_all[idx],
+                "smpl_params": smpl_params,
+                "idx": idx
+            }
+            images = {
+                "rgb": img.reshape(-1, 3).astype(np.float32),
+                "img_size": self.img_size
+            }
+            return inputs, images
+
+class ValDataset(torch.utils.data.Dataset):
+    def __init__(self, metainfo, split):
+        self.dataset = Dataset(metainfo, split)
+        self.img_size = self.dataset.img_size
+
+        self.total_pixels = np.prod(self.img_size)
+        self.pixel_per_batch = split.pixel_per_batch
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, idx):
+        image_id = int(np.random.choice(len(self.dataset), 1))  
+        self.data = self.dataset[image_id]
+        inputs, images = self.data
+
+        inputs = {
+            "uv": inputs["uv"],
+            "intrinsics": inputs['intrinsics'],
+            "pose": inputs['pose'],
+            "smpl_params": inputs["smpl_params"],
+            'image_id': image_id,
+            "idx": inputs['idx']
+        }
+        images = {
+            "rgb": images["rgb"],
+            "img_size": images["img_size"],
+            'pixel_per_batch': self.pixel_per_batch,
+            'total_pixels': self.total_pixels
+        }
+        return inputs, images
+
+class TestDataset(torch.utils.data.Dataset):
+    def __init__(self, metainfo, split):
+        self.dataset = Dataset(metainfo, split)
+
+        self.img_size = self.dataset.img_size
+
+        self.total_pixels = np.prod(self.img_size)
+        self.pixel_per_batch = split.pixel_per_batch
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        data = self.dataset[idx]
+
+        inputs, images = data
+        inputs = {
+            "uv": inputs["uv"],
+            "intrinsics": inputs['intrinsics'],
+            "pose": inputs['pose'],
+            "smpl_params": inputs["smpl_params"],
+            "idx": inputs['idx']
+        }
+        images = {
+            "rgb": images["rgb"],
+            "img_size": images["img_size"]
+        }
+        return inputs, images, self.pixel_per_batch, self.total_pixels, idx

code/lib/libmise/mise.pyx

@@ -0,0 +1,370 @@
+# distutils: language = c++
+cimport cython
+from libc.stdint cimport int32_t, int64_t
+from cython.operator cimport dereference as dref
+from libcpp.vector cimport vector
+from libcpp.map cimport map
+from libc.math cimport isnan, NAN
+import numpy as np
+
+
+cdef struct Vector3D:
+    int x, y, z
+
+
+cdef struct Voxel:
+    Vector3D loc
+    unsigned int level
+    bint is_leaf
+    unsigned long children[2][2][2]
+
+
+cdef struct GridPoint:
+    Vector3D loc
+    double value
+    bint known
+
+
+cdef inline unsigned long vec_to_idx(Vector3D coord, long resolution):
+    cdef unsigned long idx
+    idx = resolution * resolution * coord.x + resolution * coord.y + coord.z
+    return idx
+
+
+cdef class MISE:
+    cdef vector[Voxel] voxels
+    cdef vector[GridPoint] grid_points
+    cdef map[long, long] grid_point_hash
+    cdef readonly int resolution_0
+    cdef readonly int depth
+    cdef readonly double threshold
+    cdef readonly int voxel_size_0
+    cdef readonly int resolution
+
+    def __cinit__(self, int resolution_0, int depth, double threshold):
+        self.resolution_0 = resolution_0
+        self.depth = depth
+        self.threshold = threshold
+        self.voxel_size_0 = (1 << depth)
+        self.resolution = resolution_0 * self.voxel_size_0
+
+        # Create initial voxels
+        self.voxels.reserve(resolution_0 * resolution_0 * resolution_0)
+        
+        cdef Voxel voxel
+        cdef GridPoint point
+        cdef Vector3D loc
+        cdef int i, j, k
+        for i in range(resolution_0):
+            for j in range(resolution_0): 
+                for  k in range (resolution_0):
+                    loc = Vector3D(
+                        i * self.voxel_size_0, 
+                        j * self.voxel_size_0,
+                        k * self.voxel_size_0,
+                    )
+                    voxel = Voxel(
+                        loc=loc,
+                        level=0,
+                        is_leaf=True,
+                    )
+
+                    assert(self.voxels.size() == vec_to_idx(Vector3D(i, j, k), resolution_0))
+                    self.voxels.push_back(voxel)
+
+        # Create initial grid points
+        self.grid_points.reserve((resolution_0 + 1) * (resolution_0 + 1) * (resolution_0 + 1))
+        for i in range(resolution_0 + 1):
+            for j in range(resolution_0 + 1):
+                for k in range(resolution_0 + 1):
+                    loc = Vector3D(
+                        i * self.voxel_size_0, 
+                        j * self.voxel_size_0,
+                        k * self.voxel_size_0,
+                    )
+                    assert(self.grid_points.size() == vec_to_idx(Vector3D(i, j, k), resolution_0 + 1))
+                    self.add_grid_point(loc)
+
+    def update(self, int64_t[:, :] points, double[:] values):
+        """Update points and set their values. Also determine all active voxels and subdivide them."""
+        assert(points.shape[0] == values.shape[0])
+        assert(points.shape[1] == 3)
+        cdef Vector3D loc
+        cdef long idx
+        cdef int i
+
+        # Find all indices of point and set value
+        for i in range(points.shape[0]):
+            loc = Vector3D(points[i, 0], points[i, 1], points[i, 2])
+            idx = self.get_grid_point_idx(loc)
+            if idx == -1:
+                raise ValueError('Point not in grid!')
+            self.grid_points[idx].value = values[i]
+            self.grid_points[idx].known = True
+        # Subdivide activate voxels and add new points
+        self.subdivide_voxels()
+
+    def query(self):
+        """Query points to evaluate."""
+        # Find all points with unknown value
+        cdef vector[Vector3D] points
+        cdef int n_unknown = 0
+        for p in self.grid_points:
+            if not p.known:
+                n_unknown += 1 
+
+        points.reserve(n_unknown)
+        for p in self.grid_points:
+            if not p.known:
+                points.push_back(p.loc)
+
+        # Convert to numpy
+        points_np = np.zeros((points.size(), 3), dtype=np.int64)
+        cdef int64_t[:, :] points_view = points_np
+        for i in range(points.size()):
+            points_view[i, 0] = points[i].x
+            points_view[i, 1] = points[i].y
+            points_view[i, 2] = points[i].z
+
+        return points_np
+
+    def to_dense(self):
+        """Output dense matrix at highest resolution."""
+        out_array = np.full((self.resolution + 1,) * 3, np.nan)
+        cdef double[:, :, :] out_view = out_array
+        cdef GridPoint point
+        cdef int i, j, k
+        
+        for point in self.grid_points:
+            # Take voxel for which points is upper left corner
+            # assert(point.known)
+            out_view[point.loc.x, point.loc.y, point.loc.z] = point.value
+
+        # Complete along x axis
+        for i in range(1, self.resolution + 1):
+            for j in range(self.resolution + 1):
+                for k in range(self.resolution + 1):
+                    if isnan(out_view[i, j, k]):
+                        out_view[i, j, k] = out_view[i-1, j, k]
+
+        # Complete along y axis
+        for i in range(self.resolution + 1):
+            for j in range(1, self.resolution + 1):
+                for k in range(self.resolution + 1):
+                    if isnan(out_view[i, j, k]):
+                        out_view[i, j, k] = out_view[i, j-1, k]
+
+
+        # Complete along z axis
+        for i in range(self.resolution + 1):
+            for j in range(self.resolution + 1):
+                for k in range(1, self.resolution + 1):
+                    if isnan(out_view[i, j, k]):
+                        out_view[i, j, k] = out_view[i, j, k-1]
+                    assert(not isnan(out_view[i, j, k]))
+        return out_array
+
+    def get_points(self):
+        points_np = np.zeros((self.grid_points.size(), 3), dtype=np.int64)
+        values_np = np.zeros((self.grid_points.size()), dtype=np.float64)
+
+        cdef long[:, :] points_view = points_np
+        cdef double[:] values_view = values_np
+        cdef Vector3D loc
+        cdef int i
+
+        for i in range(self.grid_points.size()):
+            loc = self.grid_points[i].loc
+            points_view[i, 0] = loc.x
+            points_view[i, 1] = loc.y
+            points_view[i, 2] = loc.z
+            values_view[i] = self.grid_points[i].value
+
+        return points_np, values_np
+
+    cdef void subdivide_voxels(self) except +:
+        cdef vector[bint] next_to_positive
+        cdef vector[bint] next_to_negative
+        cdef int i, j, k
+        cdef long idx
+        cdef Vector3D loc, adj_loc
+
+        # Initialize vectors
+        next_to_positive.resize(self.voxels.size(), False)
+        next_to_negative.resize(self.voxels.size(), False)
+    
+        # Iterate over grid points and mark voxels active
+        # TODO: can move this to update operation and add attibute to voxel
+        for grid_point in self.grid_points:
+            loc = grid_point.loc
+            if not grid_point.known:
+                continue
+
+            # Iterate over the 8 adjacent voxels
+            for i in range(-1, 1):
+                for j in range(-1, 1):
+                    for k in range(-1, 1):
+                        adj_loc = Vector3D(
+                            x=loc.x + i,
+                            y=loc.y + j,
+                            z=loc.z + k,
+                        )
+                        idx = self.get_voxel_idx(adj_loc)
+                        if idx == -1:
+                            continue
+            
+                        if grid_point.value >= self.threshold:
+                            next_to_positive[idx] = True
+                        if grid_point.value <= self.threshold:
+                            next_to_negative[idx] = True
+
+        cdef int n_subdivide = 0
+        
+        for idx in range(self.voxels.size()):
+            if not self.voxels[idx].is_leaf or self.voxels[idx].level == self.depth:
+                continue
+            if next_to_positive[idx] and next_to_negative[idx]:
+                n_subdivide += 1
+
+        self.voxels.reserve(self.voxels.size() + 8 * n_subdivide)
+        self.grid_points.reserve(self.voxels.size() + 19 * n_subdivide)
+
+        for idx in range(self.voxels.size()):
+            if not self.voxels[idx].is_leaf or self.voxels[idx].level == self.depth:
+                continue
+            if next_to_positive[idx] and next_to_negative[idx]:
+                self.subdivide_voxel(idx)
+
+    cdef void subdivide_voxel(self, long idx):
+        cdef Voxel voxel
+        cdef GridPoint point
+        cdef Vector3D loc0 = self.voxels[idx].loc
+        cdef Vector3D loc
+        cdef int new_level = self.voxels[idx].level + 1
+        cdef int new_size = 1 << (self.depth - new_level)
+        assert(new_level <= self.depth)
+        assert(1 <= new_size <= self.voxel_size_0)
+
+        # Current voxel is not leaf anymore
+        self.voxels[idx].is_leaf = False
+        # Add new voxels        
+        cdef int i, j, k
+        for i in range(2):
+            for j in range(2):
+                for k in range(2):
+                    loc = Vector3D(
+                        x=loc0.x + i * new_size,
+                        y=loc0.y + j * new_size,
+                        z=loc0.z + k * new_size,
+                    )
+                    voxel = Voxel(
+                        loc=loc, 
+                        level=new_level,
+                        is_leaf=True
+                    )
+
+                    self.voxels[idx].children[i][j][k] = self.voxels.size()
+                    self.voxels.push_back(voxel)
+
+        # Add new grid points
+        for i in range(3):
+            for j in range(3):
+                for k in range(3):
+                    loc = Vector3D(
+                        loc0.x + i * new_size,
+                        loc0.y + j * new_size,
+                        loc0.z + k * new_size,
+                    )
+
+                    # Only add new grid points
+                    if self.get_grid_point_idx(loc) == -1:
+                        self.add_grid_point(loc)
+
+
+    @cython.cdivision(True) 
+    cdef long get_voxel_idx(self, Vector3D loc) except +:
+        """Utility function for getting voxel index corresponding to 3D coordinates."""
+        # Shorthands
+        cdef long resolution = self.resolution
+        cdef long resolution_0 = self.resolution_0
+        cdef long depth = self.depth
+        cdef long voxel_size_0 = self.voxel_size_0
+
+        # Return -1 if point lies outside bounds
+        if not (0 <= loc.x < resolution and 0<= loc.y < resolution and 0 <= loc.z < resolution):
+            return -1
+        
+        # Coordinates in coarse voxel grid
+        cdef Vector3D loc0 = Vector3D(
+            x=loc.x >> depth,
+            y=loc.y >> depth,
+            z=loc.z >> depth,
+        )       
+
+        # Initial voxels
+        cdef int idx = vec_to_idx(loc0, resolution_0)
+        cdef Voxel voxel = self.voxels[idx]
+        assert(voxel.loc.x == loc0.x * voxel_size_0)
+        assert(voxel.loc.y == loc0.y * voxel_size_0)
+        assert(voxel.loc.z == loc0.z * voxel_size_0)
+
+        # Relative coordinates
+        cdef Vector3D loc_rel = Vector3D(
+            x=loc.x - (loc0.x << depth),
+            y=loc.y - (loc0.y << depth),
+            z=loc.z - (loc0.z << depth),
+        ) 
+
+        cdef Vector3D loc_offset
+        cdef long voxel_size = voxel_size_0
+
+        while not voxel.is_leaf:
+            voxel_size = voxel_size >> 1
+            assert(voxel_size >= 1)
+
+            # Determine child
+            loc_offset = Vector3D(
+                x=1 if (loc_rel.x >= voxel_size) else 0,
+                y=1 if (loc_rel.y >= voxel_size) else 0,
+                z=1 if (loc_rel.z >= voxel_size) else 0,
+            )
+            # New voxel
+            idx = voxel.children[loc_offset.x][loc_offset.y][loc_offset.z]
+            voxel = self.voxels[idx]
+
+            # New relative coordinates
+            loc_rel = Vector3D(
+                x=loc_rel.x - loc_offset.x * voxel_size,
+                y=loc_rel.y - loc_offset.y * voxel_size,
+                z=loc_rel.z - loc_offset.z * voxel_size,
+            ) 
+
+            assert(0<= loc_rel.x < voxel_size)
+            assert(0<= loc_rel.y < voxel_size)
+            assert(0<= loc_rel.z < voxel_size)
+
+
+        # Return idx
+        return idx
+
+
+    cdef inline void add_grid_point(self, Vector3D loc):
+        cdef GridPoint point = GridPoint(
+            loc=loc,
+            value=0.,
+            known=False,
+        )
+        self.grid_point_hash[vec_to_idx(loc, self.resolution + 1)] = self.grid_points.size()
+        self.grid_points.push_back(point)
+
+    cdef inline int get_grid_point_idx(self, Vector3D loc):
+        p_idx = self.grid_point_hash.find(vec_to_idx(loc, self.resolution + 1))
+        if p_idx == self.grid_point_hash.end():
+            return -1
+
+        cdef int idx = dref(p_idx).second
+        assert(self.grid_points[idx].loc.x == loc.x)
+        assert(self.grid_points[idx].loc.y == loc.y)
+        assert(self.grid_points[idx].loc.z == loc.z)
+
+        return idx

code/lib/model/body_model_params.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+
+class BodyModelParams(nn.Module):
+    def __init__(self, num_frames, model_type='smpl'):
+        super(BodyModelParams, self).__init__()
+        self.num_frames = num_frames
+        self.model_type = model_type
+        self.params_dim = {
+            'betas': 10,
+            'global_orient': 3,
+            'transl': 3,
+        }
+        if model_type == 'smpl':
+            self.params_dim.update({
+                'body_pose': 69,
+            })
+        else:
+            assert ValueError(f'Unknown model type {model_type}, exiting!')
+        
+        self.param_names = self.params_dim.keys()
+        
+        for param_name in self.param_names:
+            if param_name == 'betas':
+                param = nn.Embedding(1, self.params_dim[param_name])
+                param.weight.data.fill_(0)
+                param.weight.requires_grad = False
+                setattr(self, param_name, param)
+            else:
+                param = nn.Embedding(num_frames, self.params_dim[param_name])
+                param.weight.data.fill_(0)
+                param.weight.requires_grad = False
+                setattr(self, param_name, param)
+    
+    def init_parameters(self, param_name, data, requires_grad=False):
+        getattr(self, param_name).weight.data = data[..., :self.params_dim[param_name]]
+        getattr(self, param_name).weight.requires_grad = requires_grad
+
+    def set_requires_grad(self, param_name, requires_grad=True):
+        getattr(self, param_name).weight.requires_grad = requires_grad
+
+    def forward(self, frame_ids):
+        params = {}
+        for param_name in self.param_names:
+            if param_name == 'betas':
+                params[param_name] = getattr(self, param_name)(torch.zeros_like(frame_ids))
+            else:
+                params[param_name] = getattr(self, param_name)(frame_ids)
+        return params

code/lib/model/deformer.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn.functional as F
+from .smpl import SMPLServer
+from pytorch3d import ops
+
+class SMPLDeformer():
+    def __init__(self, max_dist=0.1, K=1, gender='female', betas=None):
+        super().__init__()
+
+        self.max_dist = max_dist
+        self.K = K
+        self.smpl = SMPLServer(gender=gender)
+        smpl_params_canoical = self.smpl.param_canonical.clone()
+        smpl_params_canoical[:, 76:] = torch.tensor(betas).float().to(self.smpl.param_canonical.device)
+        cano_scale, cano_transl, cano_thetas, cano_betas = torch.split(smpl_params_canoical, [1, 3, 72, 10], dim=1)
+        smpl_output = self.smpl(cano_scale, cano_transl, cano_thetas, cano_betas)
+        self.smpl_verts = smpl_output['smpl_verts']
+        self.smpl_weights = smpl_output['smpl_weights']
+    def forward(self, x, smpl_tfs, return_weights=True, inverse=False, smpl_verts=None):
+        if x.shape[0] == 0: return x
+        if smpl_verts is None:
+            weights, outlier_mask = self.query_skinning_weights_smpl_multi(x[None], smpl_verts=self.smpl_verts[0], smpl_weights=self.smpl_weights)
+        else:
+            weights, outlier_mask = self.query_skinning_weights_smpl_multi(x[None], smpl_verts=smpl_verts[0], smpl_weights=self.smpl_weights)
+        if return_weights:
+            return weights
+
+        x_transformed = skinning(x.unsqueeze(0), weights, smpl_tfs, inverse=inverse)
+
+        return x_transformed.squeeze(0), outlier_mask
+    def forward_skinning(self, xc, cond, smpl_tfs):
+        weights, _ = self.query_skinning_weights_smpl_multi(xc, smpl_verts=self.smpl_verts[0], smpl_weights=self.smpl_weights)
+        x_transformed = skinning(xc, weights, smpl_tfs, inverse=False)
+
+        return x_transformed
+
+    def query_skinning_weights_smpl_multi(self, pts, smpl_verts, smpl_weights):
+
+        distance_batch, index_batch, neighbor_points = ops.knn_points(pts, smpl_verts.unsqueeze(0),
+                                                                      K=self.K, return_nn=True)
+        distance_batch = torch.clamp(distance_batch, max=4)
+        weights_conf = torch.exp(-distance_batch)
+        distance_batch = torch.sqrt(distance_batch)
+        weights_conf = weights_conf / weights_conf.sum(-1, keepdim=True)
+        index_batch = index_batch[0]
+        weights = smpl_weights[:, index_batch, :]
+        weights = torch.sum(weights * weights_conf.unsqueeze(-1), dim=-2).detach()
+
+        outlier_mask = (distance_batch[..., 0] > self.max_dist)[0]
+        return weights, outlier_mask
+
+    def query_weights(self, xc):
+        weights = self.forward(xc, None, return_weights=True, inverse=False)
+        return weights
+
+    def forward_skinning_normal(self, xc, normal, cond, tfs, inverse = False):
+        if normal.ndim == 2:
+            normal = normal.unsqueeze(0)
+        w = self.query_weights(xc[0], cond)
+
+        p_h = F.pad(normal, (0, 1), value=0)
+
+        if inverse:
+            # p:num_point, n:num_bone, i,j: num_dim+1
+            tf_w = torch.einsum('bpn,bnij->bpij', w.double(), tfs.double())
+            p_h = torch.einsum('bpij,bpj->bpi', tf_w.inverse(), p_h.double()).float()
+        else:
+            p_h = torch.einsum('bpn, bnij, bpj->bpi', w.double(), tfs.double(), p_h.double()).float()
+        
+        return p_h[:, :, :3]
+
+def skinning(x, w, tfs, inverse=False):
+    """Linear blend skinning
+    Args:
+        x (tensor): canonical points. shape: [B, N, D]
+        w (tensor): conditional input. [B, N, J]
+        tfs (tensor): bone transformation matrices. shape: [B, J, D+1, D+1]
+    Returns:
+        x (tensor): skinned points. shape: [B, N, D]
+    """
+    x_h = F.pad(x, (0, 1), value=1.0)
+
+    if inverse:
+        # p:n_point, n:n_bone, i,k: n_dim+1
+        w_tf = torch.einsum("bpn,bnij->bpij", w, tfs)
+        x_h = torch.einsum("bpij,bpj->bpi", w_tf.inverse(), x_h)
+    else:
+        x_h = torch.einsum("bpn,bnij,bpj->bpi", w, tfs, x_h)
+    return x_h[:, :, :3]

code/lib/model/density.py

@@ -0,0 +1,46 @@
+import torch.nn as nn
+import torch
+
+class Density(nn.Module):
+    def __init__(self, params_init={}):
+        super().__init__()
+        for p in params_init:
+            param = nn.Parameter(torch.tensor(params_init[p]))
+            setattr(self, p, param)
+
+    def forward(self, sdf, beta=None):
+        return self.density_func(sdf, beta=beta)
+
+
+class LaplaceDensity(Density):  # alpha * Laplace(loc=0, scale=beta).cdf(-sdf)
+    def __init__(self, params_init={}, beta_min=0.0001):
+        super().__init__(params_init=params_init)
+        self.beta_min = torch.tensor(beta_min).cuda()
+
+    def density_func(self, sdf, beta=None):
+        if beta is None:
+            beta = self.get_beta()
+
+        alpha = 1 / beta
+        return alpha * (0.5 + 0.5 * sdf.sign() * torch.expm1(-sdf.abs() / beta))
+
+    def get_beta(self):
+        beta = self.beta.abs() + self.beta_min
+        return beta
+
+
+class AbsDensity(Density):  # like NeRF++
+    def density_func(self, sdf, beta=None):
+        return torch.abs(sdf)
+
+
+class SimpleDensity(Density):  # like NeRF
+    def __init__(self, params_init={}, noise_std=1.0):
+        super().__init__(params_init=params_init)
+        self.noise_std = noise_std
+
+    def density_func(self, sdf, beta=None):
+        if self.training and self.noise_std > 0.0:
+            noise = torch.randn(sdf.shape).cuda() * self.noise_std
+            sdf = sdf + noise
+        return torch.relu(sdf)

code/lib/model/embedders.py

@@ -0,0 +1,50 @@
+import torch
+
+class Embedder:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+        self.create_embedding_fn()
+
+    def create_embedding_fn(self):
+        embed_fns = []
+        d = self.kwargs['input_dims']
+        out_dim = 0
+        if self.kwargs['include_input']:
+            embed_fns.append(lambda x: x)
+            out_dim += d
+
+        max_freq = self.kwargs['max_freq_log2']
+        N_freqs = self.kwargs['num_freqs']
+
+        if self.kwargs['log_sampling']:
+            freq_bands = 2. ** torch.linspace(0., max_freq, N_freqs)
+        else:
+            freq_bands = torch.linspace(2.**0., 2.**max_freq, N_freqs)
+
+        for freq in freq_bands:
+            for p_fn in self.kwargs['periodic_fns']:
+                embed_fns.append(lambda x, p_fn=p_fn,
+                                 freq=freq: p_fn(x * freq))
+                out_dim += d
+
+        self.embed_fns = embed_fns
+        self.out_dim = out_dim
+
+    def embed(self, inputs):
+        return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
+
+def get_embedder(multires, input_dims=3, mode='fourier'):
+    embed_kwargs = {
+        'include_input': True,
+        'input_dims': input_dims,
+        'max_freq_log2': multires-1,
+        'num_freqs': multires,
+        'log_sampling': True,
+        'periodic_fns': [torch.sin, torch.cos],
+    }
+    if mode == 'fourier':
+        embedder_obj = Embedder(**embed_kwargs)
+
+
+    def embed(x, eo=embedder_obj): return eo.embed(x)
+    return embed, embedder_obj.out_dim

code/lib/model/loss.py

@@ -0,0 +1,64 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class Loss(nn.Module):
+    def __init__(self, opt):
+        super().__init__()
+        self.eikonal_weight = opt.eikonal_weight
+        self.bce_weight = opt.bce_weight
+        self.opacity_sparse_weight = opt.opacity_sparse_weight
+        self.in_shape_weight = opt.in_shape_weight
+        self.eps = 1e-6
+        self.milestone = 200
+        self.l1_loss = nn.L1Loss(reduction='mean')
+        self.l2_loss = nn.MSELoss(reduction='mean')
+    
+    # L1 reconstruction loss for RGB values
+    def get_rgb_loss(self, rgb_values, rgb_gt):
+        rgb_loss = self.l1_loss(rgb_values, rgb_gt)
+        return rgb_loss
+    
+    # Eikonal loss introduced in IGR
+    def get_eikonal_loss(self, grad_theta):
+        eikonal_loss = ((grad_theta.norm(2, dim=-1) - 1)**2).mean()
+        return eikonal_loss
+
+    # BCE loss for clear boundary
+    def get_bce_loss(self, acc_map):
+        binary_loss = -1 * (acc_map * (acc_map + self.eps).log() + (1-acc_map) * (1 - acc_map + self.eps).log()).mean() * 2
+        return binary_loss
+
+    # Global opacity sparseness regularization 
+    def get_opacity_sparse_loss(self, acc_map, index_off_surface):
+        opacity_sparse_loss = self.l1_loss(acc_map[index_off_surface], torch.zeros_like(acc_map[index_off_surface]))
+        return opacity_sparse_loss
+
+    # Optional: This loss helps to stablize the training in the very beginning
+    def get_in_shape_loss(self, acc_map, index_in_surface):
+        in_shape_loss = self.l1_loss(acc_map[index_in_surface], torch.ones_like(acc_map[index_in_surface]))
+        return in_shape_loss
+
+    def forward(self, model_outputs, ground_truth):
+        nan_filter = ~torch.any(model_outputs['rgb_values'].isnan(), dim=1)
+        rgb_gt = ground_truth['rgb'][0].cuda()
+        rgb_loss = self.get_rgb_loss(model_outputs['rgb_values'][nan_filter], rgb_gt[nan_filter])
+        eikonal_loss = self.get_eikonal_loss(model_outputs['grad_theta'])
+        bce_loss = self.get_bce_loss(model_outputs['acc_map'])
+        opacity_sparse_loss = self.get_opacity_sparse_loss(model_outputs['acc_map'], model_outputs['index_off_surface'])
+        in_shape_loss = self.get_in_shape_loss(model_outputs['acc_map'], model_outputs['index_in_surface'])
+        curr_epoch_for_loss = min(self.milestone, model_outputs['epoch']) # will not increase after the milestone
+
+        loss = rgb_loss + \
+               self.eikonal_weight * eikonal_loss + \
+               self.bce_weight * bce_loss + \
+               self.opacity_sparse_weight * (1 + curr_epoch_for_loss ** 2 / 40) * opacity_sparse_loss + \
+               self.in_shape_weight * (1 - curr_epoch_for_loss / self.milestone) * in_shape_loss
+        return {
+            'loss': loss,
+            'rgb_loss': rgb_loss,
+            'eikonal_loss': eikonal_loss,
+            'bce_loss': bce_loss,
+            'opacity_sparse_loss': opacity_sparse_loss,
+            'in_shape_loss': in_shape_loss,
+        }

code/lib/model/networks.py

@@ -0,0 +1,178 @@
+import torch.nn as nn
+import torch
+import numpy as np
+from .embedders import get_embedder
+
+class ImplicitNet(nn.Module):
+    def __init__(self, opt):
+        super().__init__()
+
+        dims = [opt.d_in] + list(
+            opt.dims) + [opt.d_out + opt.feature_vector_size]
+        self.num_layers = len(dims)
+        self.skip_in = opt.skip_in
+        self.embed_fn = None
+        self.opt = opt
+
+        if opt.multires > 0:
+            embed_fn, input_ch = get_embedder(opt.multires, input_dims=opt.d_in, mode=opt.embedder_mode)
+            self.embed_fn = embed_fn
+            dims[0] = input_ch
+        self.cond = opt.cond   
+        if self.cond == 'smpl':
+            self.cond_layer = [0]
+            self.cond_dim = 69
+        elif self.cond == 'frame':
+            self.cond_layer = [0]
+            self.cond_dim = opt.dim_frame_encoding
+        self.dim_pose_embed = 0
+        if self.dim_pose_embed > 0:
+            self.lin_p0 = nn.Linear(self.cond_dim, self.dim_pose_embed)
+            self.cond_dim = self.dim_pose_embed
+        for l in range(0, self.num_layers - 1):
+            if l + 1 in self.skip_in:
+                out_dim = dims[l + 1] - dims[0]
+            else:
+                out_dim = dims[l + 1]
+            
+            if self.cond != 'none' and l in self.cond_layer:
+                lin = nn.Linear(dims[l] + self.cond_dim, out_dim)
+            else:
+                lin = nn.Linear(dims[l], out_dim)
+            if opt.init == 'geometry':
+                if l == self.num_layers - 2:
+                    torch.nn.init.normal_(lin.weight,
+                                          mean=np.sqrt(np.pi) /
+                                          np.sqrt(dims[l]),
+                                          std=0.0001)
+                    torch.nn.init.constant_(lin.bias, -opt.bias)
+                elif opt.multires > 0 and l == 0:
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
+                    torch.nn.init.normal_(lin.weight[:, :3], 0.0,
+                                          np.sqrt(2) / np.sqrt(out_dim))
+                elif opt.multires > 0 and l in self.skip_in:
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    torch.nn.init.normal_(lin.weight, 0.0,
+                                          np.sqrt(2) / np.sqrt(out_dim))
+                    torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):],
+                                            0.0)
+                else:
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    torch.nn.init.normal_(lin.weight, 0.0,
+                                          np.sqrt(2) / np.sqrt(out_dim))
+            if opt.init == 'zero':
+                init_val = 1e-5
+                if l == self.num_layers - 2:
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    torch.nn.init.uniform_(lin.weight, -init_val, init_val)
+            if opt.weight_norm:
+                lin = nn.utils.weight_norm(lin)
+            setattr(self, "lin" + str(l), lin)
+        self.softplus = nn.Softplus(beta=100)
+
+    def forward(self, input, cond, current_epoch=None):
+        if input.ndim == 2: input = input.unsqueeze(0)
+
+        num_batch, num_point, num_dim = input.shape
+
+        if num_batch * num_point == 0: return input
+
+        input = input.reshape(num_batch * num_point, num_dim)
+
+        if self.cond != 'none':
+            num_batch, num_cond = cond[self.cond].shape
+
+            input_cond = cond[self.cond].unsqueeze(1).expand(num_batch, num_point, num_cond)
+
+            input_cond = input_cond.reshape(num_batch * num_point, num_cond)
+
+            if self.dim_pose_embed:
+                input_cond = self.lin_p0(input_cond)
+
+        if self.embed_fn is not None:
+            input = self.embed_fn(input)
+
+        x = input
+
+        for l in range(0, self.num_layers - 1):
+            lin = getattr(self, "lin" + str(l))
+            if self.cond != 'none' and l in self.cond_layer:
+                x = torch.cat([x, input_cond], dim=-1)
+            if l in self.skip_in:
+                x = torch.cat([x, input], 1) / np.sqrt(2)
+            x = lin(x)
+            if l < self.num_layers - 2:
+                x = self.softplus(x)
+        
+        x = x.reshape(num_batch, num_point, -1)
+
+        return x
+
+    def gradient(self, x, cond):
+        x.requires_grad_(True)
+        y = self.forward(x, cond)[:, :1]
+        d_output = torch.ones_like(y, requires_grad=False, device=y.device)
+        gradients = torch.autograd.grad(outputs=y,
+                                        inputs=x,
+                                        grad_outputs=d_output,
+                                        create_graph=True,
+                                        retain_graph=True,
+                                        only_inputs=True)[0]
+        return gradients.unsqueeze(1)
+
+
+class RenderingNet(nn.Module):
+    def __init__(self, opt):
+        super().__init__()
+
+        self.mode = opt.mode
+        dims = [opt.d_in + opt.feature_vector_size] + list(
+            opt.dims) + [opt.d_out]
+
+        self.embedview_fn = None
+        if opt.multires_view > 0:
+            embedview_fn, input_ch = get_embedder(opt.multires_view)
+            self.embedview_fn = embedview_fn
+            dims[0] += (input_ch - 3)
+        if self.mode == 'nerf_frame_encoding':
+            dims[0] += opt.dim_frame_encoding
+        if self.mode == 'pose':
+            self.dim_cond_embed = 8 
+            self.cond_dim = 69 # dimension of the body pose, global orientation excluded.
+            # lower the condition dimension
+            self.lin_pose = torch.nn.Linear(self.cond_dim, self.dim_cond_embed)
+        self.num_layers = len(dims)
+        for l in range(0, self.num_layers - 1):
+            out_dim = dims[l + 1]
+            lin = nn.Linear(dims[l], out_dim)
+            if opt.weight_norm:
+                lin = nn.utils.weight_norm(lin)
+            setattr(self, "lin" + str(l), lin)
+        self.relu = nn.ReLU()
+        self.sigmoid = nn.Sigmoid()
+        
+    def forward(self, points, normals, view_dirs, body_pose, feature_vectors, frame_latent_code=None):
+        if self.embedview_fn is not None:
+            if self.mode == 'nerf_frame_encoding':
+                view_dirs = self.embedview_fn(view_dirs)
+
+        if self.mode == 'nerf_frame_encoding':
+            frame_latent_code = frame_latent_code.expand(view_dirs.shape[0], -1)
+            rendering_input = torch.cat([view_dirs, frame_latent_code, feature_vectors], dim=-1)
+        elif self.mode == 'pose':
+            num_points = points.shape[0]
+            body_pose = body_pose.unsqueeze(1).expand(-1, num_points, -1).reshape(num_points, -1)
+            body_pose = self.lin_pose(body_pose)
+            rendering_input = torch.cat([points, normals, body_pose, feature_vectors], dim=-1)
+        else:
+            raise NotImplementedError
+
+        x = rendering_input
+        for l in range(0, self.num_layers - 1):
+            lin = getattr(self, "lin" + str(l))
+            x = lin(x)
+            if l < self.num_layers - 2:
+                x = self.relu(x)
+        x = self.sigmoid(x)
+        return x

code/lib/model/ray_sampler.py

@@ -0,0 +1,234 @@
+import abc
+import torch
+from lib.utils import utils
+
+class RaySampler(metaclass=abc.ABCMeta):
+    def __init__(self,near, far):
+        self.near = near
+        self.far = far
+
+    @abc.abstractmethod
+    def get_z_vals(self, ray_dirs, cam_loc, model):
+        pass
+
+class UniformSampler(RaySampler):
+    """Samples uniformly in the range [near, far]
+    """
+    def __init__(self, scene_bounding_sphere, near, N_samples, take_sphere_intersection=False, far=-1):
+        super().__init__(near, 2.0 * scene_bounding_sphere if far == -1 else far)  # default far is 2*R
+        self.N_samples = N_samples
+        self.scene_bounding_sphere = scene_bounding_sphere
+        self.take_sphere_intersection = take_sphere_intersection
+
+    def get_z_vals(self, ray_dirs, cam_loc, model):
+        if not self.take_sphere_intersection:
+            near, far = self.near * torch.ones(ray_dirs.shape[0], 1).cuda(), self.far * torch.ones(ray_dirs.shape[0], 1).cuda()
+        else:
+            sphere_intersections = utils.get_sphere_intersections(cam_loc, ray_dirs, r=self.scene_bounding_sphere)
+            near = self.near * torch.ones(ray_dirs.shape[0], 1).cuda()
+            far = sphere_intersections[:,1:]
+
+        t_vals = torch.linspace(0., 1., steps=self.N_samples).cuda()
+        z_vals = near * (1. - t_vals) + far * (t_vals)
+
+        if model.training:
+            # get intervals between samples
+            mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
+            upper = torch.cat([mids, z_vals[..., -1:]], -1)
+            lower = torch.cat([z_vals[..., :1], mids], -1)
+            # stratified samples in those intervals
+            t_rand = torch.rand(z_vals.shape).cuda()
+
+            z_vals = lower + (upper - lower) * t_rand
+
+        return z_vals
+
+class ErrorBoundSampler(RaySampler):
+    def __init__(self, scene_bounding_sphere, near, N_samples, N_samples_eval, N_samples_extra,
+                 eps, beta_iters, max_total_iters,
+                 inverse_sphere_bg=False, N_samples_inverse_sphere=0, add_tiny=0.0):
+        super().__init__(near, 2.0 * scene_bounding_sphere)
+        self.N_samples = N_samples
+        self.N_samples_eval = N_samples_eval
+        self.uniform_sampler = UniformSampler(scene_bounding_sphere, near, N_samples_eval, take_sphere_intersection=inverse_sphere_bg)
+
+        self.N_samples_extra = N_samples_extra
+
+        self.eps = eps
+        self.beta_iters = beta_iters
+        self.max_total_iters = max_total_iters
+        self.scene_bounding_sphere = scene_bounding_sphere
+        self.add_tiny = add_tiny
+
+        self.inverse_sphere_bg = inverse_sphere_bg
+        if inverse_sphere_bg:
+            N_samples_inverse_sphere = 32
+            self.inverse_sphere_sampler = UniformSampler(1.0, 0.0, N_samples_inverse_sphere, False, far=1.0)
+
+    def get_z_vals(self, ray_dirs, cam_loc, model, cond, smpl_tfs, eval_mode, smpl_verts):
+        beta0 = model.density.get_beta().detach()
+
+        # Start with uniform sampling
+        z_vals = self.uniform_sampler.get_z_vals(ray_dirs, cam_loc, model)
+        samples, samples_idx = z_vals, None
+
+        # Get maximum beta from the upper bound (Lemma 2)
+        dists = z_vals[:, 1:] - z_vals[:, :-1]
+        bound = (1.0 / (4.0 * torch.log(torch.tensor(self.eps + 1.0)))) * (dists ** 2.).sum(-1)
+        beta = torch.sqrt(bound)
+
+        total_iters, not_converge = 0, True
+
+        # VolSDF Algorithm 1
+        while not_converge and total_iters < self.max_total_iters:
+            points = cam_loc.unsqueeze(1) + samples.unsqueeze(2) * ray_dirs.unsqueeze(1)
+            points_flat = points.reshape(-1, 3)
+            # Calculating the SDF only for the new sampled points
+            model.implicit_network.eval()
+            with torch.no_grad():
+                samples_sdf = model.sdf_func_with_smpl_deformer(points_flat, cond, smpl_tfs, smpl_verts=smpl_verts)[0] 
+            model.implicit_network.train()
+            if samples_idx is not None:
+                sdf_merge = torch.cat([sdf.reshape(-1, z_vals.shape[1] - samples.shape[1]),
+                                       samples_sdf.reshape(-1, samples.shape[1])], -1)
+                sdf = torch.gather(sdf_merge, 1, samples_idx).reshape(-1, 1)
+            else:
+                sdf = samples_sdf
+
+
+            # Calculating the bound d* (Theorem 1)
+            d = sdf.reshape(z_vals.shape)
+            dists = z_vals[:, 1:] - z_vals[:, :-1]
+            a, b, c = dists, d[:, :-1].abs(), d[:, 1:].abs()
+            first_cond = a.pow(2) + b.pow(2) <= c.pow(2)
+            second_cond = a.pow(2) + c.pow(2) <= b.pow(2)
+            d_star = torch.zeros(z_vals.shape[0], z_vals.shape[1] - 1).cuda()
+            d_star[first_cond] = b[first_cond]
+            d_star[second_cond] = c[second_cond]
+            s = (a + b + c) / 2.0
+            area_before_sqrt = s * (s - a) * (s - b) * (s - c)
+            mask = ~first_cond & ~second_cond & (b + c - a > 0)
+            d_star[mask] = (2.0 * torch.sqrt(area_before_sqrt[mask])) / (a[mask])
+            d_star = (d[:, 1:].sign() * d[:, :-1].sign() == 1) * d_star  # Fixing the sign
+
+
+            # Updating beta using line search
+            curr_error = self.get_error_bound(beta0, model, sdf, z_vals, dists, d_star)
+            beta[curr_error <= self.eps] = beta0
+            beta_min, beta_max = beta0.unsqueeze(0).repeat(z_vals.shape[0]), beta
+            for j in range(self.beta_iters):
+                beta_mid = (beta_min + beta_max) / 2.
+                curr_error = self.get_error_bound(beta_mid.unsqueeze(-1), model, sdf, z_vals, dists, d_star)
+                beta_max[curr_error <= self.eps] = beta_mid[curr_error <= self.eps]
+                beta_min[curr_error > self.eps] = beta_mid[curr_error > self.eps]
+            beta = beta_max
+
+
+            # Upsample more points
+            density = model.density(sdf.reshape(z_vals.shape), beta=beta.unsqueeze(-1))
+
+            dists = torch.cat([dists, torch.tensor([1e10]).cuda().unsqueeze(0).repeat(dists.shape[0], 1)], -1)
+            free_energy = dists * density
+            shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1).cuda(), free_energy[:, :-1]], dim=-1)
+            alpha = 1 - torch.exp(-free_energy)
+            transmittance = torch.exp(-torch.cumsum(shifted_free_energy, dim=-1))
+            weights = alpha * transmittance  # probability of the ray hits something here
+
+            # Check if we are done and this is the last sampling
+            total_iters += 1
+            not_converge = beta.max() > beta0
+
+            if not_converge and total_iters < self.max_total_iters:
+                ''' Sample more points proportional to the current error bound'''
+
+                N = self.N_samples_eval
+
+                bins = z_vals
+                error_per_section = torch.exp(-d_star / beta.unsqueeze(-1)) * (dists[:,:-1] ** 2.) / (4 * beta.unsqueeze(-1) ** 2)
+                error_integral = torch.cumsum(error_per_section, dim=-1)
+                bound_opacity = (torch.clamp(torch.exp(error_integral),max=1.e6) - 1.0) * transmittance[:,:-1]
+
+                pdf = bound_opacity + self.add_tiny
+                pdf = pdf / torch.sum(pdf, -1, keepdim=True)
+                cdf = torch.cumsum(pdf, -1)
+                cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
+
+            else:
+                ''' Sample the final sample set to be used in the volume rendering integral '''
+
+                N = self.N_samples
+
+                bins = z_vals
+                pdf = weights[..., :-1]
+                pdf = pdf + 1e-5  # prevent nans
+                pdf = pdf / torch.sum(pdf, -1, keepdim=True)
+                cdf = torch.cumsum(pdf, -1)
+                cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)  # (batch, len(bins))
+
+
+            # Invert CDF
+            if (not_converge and total_iters < self.max_total_iters) or (not model.training):
+                u = torch.linspace(0., 1., steps=N).cuda().unsqueeze(0).repeat(cdf.shape[0], 1)
+            else:
+                u = torch.rand(list(cdf.shape[:-1]) + [N]).cuda()
+            u = u.contiguous()
+
+            inds = torch.searchsorted(cdf, u, right=True)
+            below = torch.max(torch.zeros_like(inds - 1), inds - 1)
+            above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
+            inds_g = torch.stack([below, above], -1)  # (batch, N_samples, 2)
+
+            matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
+            cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
+            bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
+
+            denom = (cdf_g[..., 1] - cdf_g[..., 0])
+            denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
+            t = (u - cdf_g[..., 0]) / denom
+            samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
+
+
+            # Adding samples if we not converged
+            if not_converge and total_iters < self.max_total_iters:
+                z_vals, samples_idx = torch.sort(torch.cat([z_vals, samples], -1), -1)
+
+
+        z_samples = samples
+
+        near, far = self.near * torch.ones(ray_dirs.shape[0], 1).cuda(), self.far * torch.ones(ray_dirs.shape[0],1).cuda()
+        if self.inverse_sphere_bg: # if inverse sphere then need to add the far sphere intersection
+            far = utils.get_sphere_intersections(cam_loc, ray_dirs, r=self.scene_bounding_sphere)[:,1:]
+
+        if self.N_samples_extra > 0:
+            if model.training:
+                sampling_idx = torch.randperm(z_vals.shape[1])[:self.N_samples_extra]
+            else:
+                sampling_idx = torch.linspace(0, z_vals.shape[1]-1, self.N_samples_extra).long()
+            z_vals_extra = torch.cat([near, far, z_vals[:,sampling_idx]], -1)
+        else:
+            z_vals_extra = torch.cat([near, far], -1)
+
+        z_vals, _ = torch.sort(torch.cat([z_samples, z_vals_extra], -1), -1)
+
+        # add some of the near surface points
+        idx = torch.randint(z_vals.shape[-1], (z_vals.shape[0],)).cuda()
+        z_samples_eik = torch.gather(z_vals, 1, idx.unsqueeze(-1))
+
+        if self.inverse_sphere_bg:
+            z_vals_inverse_sphere = self.inverse_sphere_sampler.get_z_vals(ray_dirs, cam_loc, model)
+            z_vals_inverse_sphere = z_vals_inverse_sphere * (1./self.scene_bounding_sphere)
+            z_vals = (z_vals, z_vals_inverse_sphere)
+
+        return z_vals, z_samples_eik
+
+    def get_error_bound(self, beta, model, sdf, z_vals, dists, d_star):
+        density = model.density(sdf.reshape(z_vals.shape), beta=beta)
+        shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1).cuda(), dists * density[:, :-1]], dim=-1)
+        integral_estimation = torch.cumsum(shifted_free_energy, dim=-1)
+        error_per_section = torch.exp(-d_star / beta) * (dists ** 2.) / (4 * beta ** 2)
+        error_integral = torch.cumsum(error_per_section, dim=-1)
+        bound_opacity = (torch.clamp(torch.exp(error_integral), max=1.e6) - 1.0) * torch.exp(-integral_estimation[:, :-1])
+
+        return bound_opacity.max(-1)[0]
+
+

code/lib/model/sampler.py

@@ -0,0 +1,29 @@
+import torch
+
+
+class PointInSpace:
+    def __init__(self, global_sigma=0.5, local_sigma=0.01):
+        self.global_sigma = global_sigma
+        self.local_sigma = local_sigma
+
+    def get_points(self, pc_input=None, local_sigma=None, global_ratio=0.125):
+        """Sample one point near each of the given point + 1/8 uniformly. 
+        Args:
+            pc_input (tensor): sampling centers. shape: [B, N, D]
+        Returns:
+            samples (tensor): sampled points. shape: [B, N + N / 8, D]
+        """
+
+        batch_size, sample_size, dim = pc_input.shape
+        if local_sigma is None:
+            sample_local = pc_input + (torch.randn_like(pc_input) * self.local_sigma)
+        else:
+            sample_local = pc_input + (torch.randn_like(pc_input) * local_sigma)
+        sample_global = (
+            torch.rand(batch_size, int(sample_size * global_ratio), dim, device=pc_input.device)
+            * (self.global_sigma * 2)
+        ) - self.global_sigma
+
+        sample = torch.cat([sample_local, sample_global], dim=1)
+
+        return sample

code/lib/model/smpl.py

@@ -0,0 +1,94 @@
+import torch
+import hydra
+import numpy as np
+from ..smpl.body_models import SMPL
+
+class SMPLServer(torch.nn.Module):
+
+    def __init__(self, gender='neutral', betas=None, v_template=None):
+        super().__init__()
+
+
+        self.smpl = SMPL(model_path=hydra.utils.to_absolute_path('lib/smpl/smpl_model'),
+                         gender=gender,
+                         batch_size=1,
+                         use_hands=False,
+                         use_feet_keypoints=False,
+                         dtype=torch.float32).cuda()
+
+        self.bone_parents = self.smpl.bone_parents.astype(int)
+        self.bone_parents[0] = -1
+        self.bone_ids = []
+        self.faces = self.smpl.faces
+        for i in range(24): self.bone_ids.append([self.bone_parents[i], i])
+
+        if v_template is not None:
+            self.v_template = torch.tensor(v_template).float().cuda()
+        else:
+            self.v_template = None
+
+        if betas is not None:
+            self.betas = torch.tensor(betas).float().cuda()
+        else:
+            self.betas = None
+
+        # define the canonical pose
+        param_canonical = torch.zeros((1, 86),dtype=torch.float32).cuda()
+        param_canonical[0, 0] = 1
+        param_canonical[0, 9] = np.pi / 6
+        param_canonical[0, 12] = -np.pi / 6
+        if self.betas is not None and self.v_template is None:
+            param_canonical[0,-10:] = self.betas
+        self.param_canonical = param_canonical
+
+        output = self.forward(*torch.split(self.param_canonical, [1, 3, 72, 10], dim=1), absolute=True)
+        self.verts_c = output['smpl_verts']
+        self.joints_c = output['smpl_jnts']
+        self.tfs_c_inv = output['smpl_tfs'].squeeze(0).inverse()
+
+
+    def forward(self, scale, transl, thetas, betas, absolute=False):
+        """return SMPL output from params
+        Args:
+            scale : scale factor. shape: [B, 1]
+            transl: translation. shape: [B, 3]
+            thetas: pose. shape: [B, 72]
+            betas: shape. shape: [B, 10]
+            absolute (bool): if true return smpl_tfs wrt thetas=0. else wrt thetas=thetas_canonical. 
+        Returns:
+            smpl_verts: vertices. shape: [B, 6893. 3]
+            smpl_tfs: bone transformations. shape: [B, 24, 4, 4]
+            smpl_jnts: joint positions. shape: [B, 25, 3]
+        """
+
+        output = {}
+
+        # ignore betas if v_template is provided
+        if self.v_template is not None:
+            betas = torch.zeros_like(betas)
+
+        
+        smpl_output = self.smpl.forward(betas=betas,
+                                        transl=torch.zeros_like(transl),
+                                        body_pose=thetas[:, 3:],
+                                        global_orient=thetas[:, :3],
+                                        return_verts=True,
+                                        return_full_pose=True,
+                                        v_template=self.v_template)
+
+        verts = smpl_output.vertices.clone()
+        output['smpl_verts'] = verts * scale.unsqueeze(1) + transl.unsqueeze(1) * scale.unsqueeze(1)
+
+        joints = smpl_output.joints.clone()
+        output['smpl_jnts'] = joints * scale.unsqueeze(1) + transl.unsqueeze(1) * scale.unsqueeze(1)
+
+        tf_mats = smpl_output.T.clone()
+        tf_mats[:, :, :3, :] = tf_mats[:, :, :3, :] * scale.unsqueeze(1).unsqueeze(1)
+        tf_mats[:, :, :3, 3] = tf_mats[:, :, :3, 3] + transl.unsqueeze(1) * scale.unsqueeze(1)
+
+        if not absolute:
+            tf_mats = torch.einsum('bnij,njk->bnik', tf_mats, self.tfs_c_inv)
+        
+        output['smpl_tfs'] = tf_mats
+        output['smpl_weights'] = smpl_output.weights
+        return output

code/lib/model/v2a.py

@@ -0,0 +1,368 @@
+from .networks import ImplicitNet, RenderingNet
+from .density import LaplaceDensity, AbsDensity
+from .ray_sampler import ErrorBoundSampler
+from .deformer import SMPLDeformer
+from .smpl import SMPLServer
+
+from .sampler import PointInSpace
+
+from ..utils import utils
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.autograd import grad
+import hydra
+import kaolin
+from kaolin.ops.mesh import index_vertices_by_faces
+class V2A(nn.Module):
+    def __init__(self, opt, betas_path, gender, num_training_frames):
+        super().__init__()
+
+        # Foreground networks
+        self.implicit_network = ImplicitNet(opt.implicit_network)
+        self.rendering_network = RenderingNet(opt.rendering_network)
+
+        # Background networks
+        self.bg_implicit_network = ImplicitNet(opt.bg_implicit_network)
+        self.bg_rendering_network = RenderingNet(opt.bg_rendering_network)
+
+        # Frame latent encoder
+        self.frame_latent_encoder = nn.Embedding(num_training_frames, opt.bg_rendering_network.dim_frame_encoding)
+        self.sampler = PointInSpace()
+
+        betas = np.load(betas_path)
+        self.use_smpl_deformer = opt.use_smpl_deformer
+        self.gender = gender
+        if self.use_smpl_deformer:
+            self.deformer = SMPLDeformer(betas=betas, gender=self.gender)
+        
+        # pre-defined bounding sphere
+        self.sdf_bounding_sphere = 3.0
+        
+        # threshold for the out-surface points
+        self.threshold = 0.05
+        
+        self.density = LaplaceDensity(**opt.density)
+        self.bg_density = AbsDensity()
+
+        self.ray_sampler = ErrorBoundSampler(self.sdf_bounding_sphere, inverse_sphere_bg=True, **opt.ray_sampler)
+        self.smpl_server = SMPLServer(gender=self.gender, betas=betas)
+
+        if opt.smpl_init:
+            smpl_model_state = torch.load(hydra.utils.to_absolute_path('../assets/smpl_init.pth'))
+            self.implicit_network.load_state_dict(smpl_model_state["model_state_dict"])
+
+        self.smpl_v_cano = self.smpl_server.verts_c
+        self.smpl_f_cano = torch.tensor(self.smpl_server.smpl.faces.astype(np.int64), device=self.smpl_v_cano.device)
+
+        self.mesh_v_cano = self.smpl_server.verts_c
+        self.mesh_f_cano = torch.tensor(self.smpl_server.smpl.faces.astype(np.int64), device=self.smpl_v_cano.device)
+        self.mesh_face_vertices = index_vertices_by_faces(self.mesh_v_cano, self.mesh_f_cano)
+
+    def sdf_func_with_smpl_deformer(self, x, cond, smpl_tfs, smpl_verts):
+        """ sdf_func_with_smpl_deformer method 
+        Used to compute SDF values for input points using the SMPL deformer and the implicit network. 
+        It handles the deforming of points, network inference, feature extraction, and handling of outlier points.
+        """
+        if hasattr(self, "deformer"):
+            x_c, outlier_mask = self.deformer.forward(x, smpl_tfs, return_weights=False, inverse=True, smpl_verts=smpl_verts)
+            output = self.implicit_network(x_c, cond)[0]
+            sdf = output[:, 0:1]
+            feature = output[:, 1:]
+            if not self.training:
+                sdf[outlier_mask] = 4. # set a large SDF value for outlier points
+            
+        return sdf, x_c, feature
+    
+    def check_off_in_surface_points_cano_mesh(self, x_cano, N_samples, threshold=0.05):
+        """check_off_in_surface_points_cano_mesh method 
+        Used to check whether points are off the surface or within the surface of a canonical mesh. 
+        It calculates distances, signs, and signed distances to determine the position of points with respect to the mesh surface. 
+        The method plays a role in identifying points that might be considered outliers or outside the reconstructed avatar's surface.
+        """
+
+        distance, _, _ = kaolin.metrics.trianglemesh.point_to_mesh_distance(x_cano.unsqueeze(0).contiguous(), self.mesh_face_vertices)
+
+        distance = torch.sqrt(distance) # kaolin outputs squared distance
+        sign = kaolin.ops.mesh.check_sign(self.mesh_v_cano, self.mesh_f_cano, x_cano.unsqueeze(0)).float()
+        sign = 1 - 2 * sign   # -1 for off-surface, 1 for in-surface
+        signed_distance = sign * distance
+        batch_size = x_cano.shape[0] // N_samples  
+        signed_distance = signed_distance.reshape(batch_size, N_samples, 1)   # The distances are reshaped to match the batch size and the number of samples
+
+        minimum = torch.min(signed_distance, 1)[0]
+        index_off_surface = (minimum > threshold).squeeze(1)
+        index_in_surface = (minimum <= 0.).squeeze(1)
+        return index_off_surface, index_in_surface   # Indexes of off-surface points and in-surface points
+
+    def forward(self, input):
+        # Parse model input, prepares the necessary input data and SMPL parameters for subsequent calculations
+        torch.set_grad_enabled(True)
+        intrinsics = input["intrinsics"]
+        pose = input["pose"]
+        uv = input["uv"]
+
+        scale = input['smpl_params'][:, 0]
+        smpl_pose = input["smpl_pose"]
+        smpl_shape = input["smpl_shape"]
+        smpl_trans = input["smpl_trans"]
+        smpl_output = self.smpl_server(scale, smpl_trans, smpl_pose, smpl_shape)   #  invokes the SMPL model to obtain the transformations for pose and shape changes
+
+        smpl_tfs = smpl_output['smpl_tfs']
+
+        cond = {'smpl': smpl_pose[:, 3:]/np.pi}
+        if self.training:
+            if input['current_epoch'] < 20 or input['current_epoch'] % 20 == 0:   # set the pose to zero for the first 20 epochs
+                cond = {'smpl': smpl_pose[:, 3:] * 0.}  
+        ray_dirs, cam_loc = utils.get_camera_params(uv, pose, intrinsics)   # get the ray directions and camera location
+        batch_size, num_pixels, _ = ray_dirs.shape
+
+        cam_loc = cam_loc.unsqueeze(1).repeat(1, num_pixels, 1).reshape(-1, 3)   # reshape to match the batch size and the number of pixels
+        ray_dirs = ray_dirs.reshape(-1, 3)    # reshape to match the batch size and the number of pixels
+
+        z_vals, _ = self.ray_sampler.get_z_vals(ray_dirs, cam_loc, self, cond, smpl_tfs, eval_mode=True, smpl_verts=smpl_output['smpl_verts'])   # get the z values for each pixel  
+
+        z_vals, z_vals_bg = z_vals   # unpack the z values for the foreground and the background
+        z_max = z_vals[:,-1]   # get the maximum z value
+        z_vals = z_vals[:,:-1]   # get the z values for the foreground
+        N_samples = z_vals.shape[1]   # get the number of samples
+
+        points = cam_loc.unsqueeze(1) + z_vals.unsqueeze(2) * ray_dirs.unsqueeze(1)   # 3D points along the rays are calculated by adding z_vals scaled by ray directions to the camera location. The result is stored in the points tensor of shape (batch_size * num_pixels, N_samples, 3)
+        points_flat = points.reshape(-1, 3)   # The points tensor is reshaped into a flattened tensor points_flat of shape (batch_size * num_pixels * N_samples, 3)
+
+        dirs = ray_dirs.unsqueeze(1).repeat(1,N_samples,1)   # The dirs tensor is created by repeating ray_dirs for each sample along the rays. The resulting tensor has shape (batch_size * num_pixels, N_samples, 3)
+        sdf_output, canonical_points, feature_vectors = self.sdf_func_with_smpl_deformer(points_flat, cond, smpl_tfs, smpl_output['smpl_verts'])   # The sdf_func_with_smpl_deformer method is called to compute the signed distance functions (SDF) for the points
+
+        sdf_output = sdf_output.unsqueeze(1)   # The sdf_output tensor is reshaped by unsqueezing along the first dimension
+
+        if self.training:
+            index_off_surface, index_in_surface = self.check_off_in_surface_points_cano_mesh(canonical_points, N_samples, threshold=self.threshold)
+            canonical_points = canonical_points.reshape(num_pixels, N_samples, 3)   
+
+            canonical_points = canonical_points.reshape(-1, 3)   # The canonical points tensor flattened to shape (-1, 3) 
+
+            # sample canonical SMPL surface pnts for the eikonal loss
+            smpl_verts_c = self.smpl_server.verts_c.repeat(batch_size, 1,1)   # The canonical SMPL surface vertices are repeated across the batch dimension
+            
+            indices = torch.randperm(smpl_verts_c.shape[1])[:num_pixels].cuda()   # Random indices are generated to select a subset of vertices for sampling. The number of selected vertices is num_pixels
+            verts_c = torch.index_select(smpl_verts_c, 1, indices)   # The selected vertices are gathered from smpl_verts_c, resulting in the tensor verts_c.
+            sample = self.sampler.get_points(verts_c, global_ratio=0.)   # The get_points method of the sampler class is called to sample points around the canonical SMPL surface points. The global_ratio is set to 0.0, indicating local sampling
+
+            sample.requires_grad_()   # The sampled points are marked as requiring gradients 
+            local_pred = self.implicit_network(sample, cond)[..., 0:1]   # The sampled points (sample) are passed through the implicit network along with the conditioning (cond). The local prediction (SDF) for each sampled point is extracted using [..., 0:1]
+            grad_theta = gradient(sample, local_pred)   # compute gradients with respect to the sampled points and their local predictions (local_pred).
+
+            differentiable_points = canonical_points   # The differentiable_points tensor is assigned the value of canonical_points
+
+        else:
+            differentiable_points = canonical_points.reshape(num_pixels, N_samples, 3).reshape(-1, 3)
+            grad_theta = None
+
+        sdf_output = sdf_output.reshape(num_pixels, N_samples, 1).reshape(-1, 1)   # flattened to shape (num_pixels * N_samples, )
+        z_vals = z_vals
+        view = -dirs.reshape(-1, 3)   # The view vector is calculated as the negation of the reshaped dirs, giving the view directions for points along the rays.
+
+        if differentiable_points.shape[0] > 0:   # If there are differentiable points (indicating that gradient information is available)
+            fg_rgb_flat, others = self.get_rbg_value(points_flat, differentiable_points, view,
+                                                     cond, smpl_tfs, feature_vectors=feature_vectors, is_training=self.training)   # The returned values include fg_rgb_flat (foreground RGB values) and others (other calculated values, including normals)                 
+            normal_values = others['normals']   # The normal values are extracted from the others dictionary
+
+        if 'image_id' in input.keys():
+            frame_latent_code = self.frame_latent_encoder(input['image_id'])
+        else:
+            frame_latent_code = self.frame_latent_encoder(input['idx'])
+
+        fg_rgb = fg_rgb_flat.reshape(-1, N_samples, 3)
+        normal_values = normal_values.reshape(-1, N_samples, 3)
+        weights, bg_transmittance = self.volume_rendering(z_vals, z_max, sdf_output)
+
+        fg_rgb_values = torch.sum(weights.unsqueeze(-1) * fg_rgb, 1)
+
+        # Background rendering
+        if input['idx'] is not None:
+            N_bg_samples = z_vals_bg.shape[1]
+            z_vals_bg = torch.flip(z_vals_bg, dims=[-1, ])  # 1--->0
+
+            bg_dirs = ray_dirs.unsqueeze(1).repeat(1,N_bg_samples,1)
+            bg_locs = cam_loc.unsqueeze(1).repeat(1,N_bg_samples,1)
+
+            bg_points = self.depth2pts_outside(bg_locs, bg_dirs, z_vals_bg)  # [..., N_samples, 4]
+            bg_points_flat = bg_points.reshape(-1, 4)
+            bg_dirs_flat = bg_dirs.reshape(-1, 3)
+            bg_output = self.bg_implicit_network(bg_points_flat, {'frame': frame_latent_code})[0]
+            bg_sdf = bg_output[:, :1]
+            bg_feature_vectors = bg_output[:, 1:]
+            
+            bg_rendering_output = self.bg_rendering_network(None, None, bg_dirs_flat, None, bg_feature_vectors, frame_latent_code)
+            if bg_rendering_output.shape[-1] == 4:
+                bg_rgb_flat = bg_rendering_output[..., :-1]
+                shadow_r = bg_rendering_output[..., -1]
+                bg_rgb = bg_rgb_flat.reshape(-1, N_bg_samples, 3)
+                shadow_r = shadow_r.reshape(-1, N_bg_samples, 1)
+                bg_rgb = (1 - shadow_r) * bg_rgb
+            else:
+                bg_rgb_flat = bg_rendering_output
+                bg_rgb = bg_rgb_flat.reshape(-1, N_bg_samples, 3)
+            bg_weights = self.bg_volume_rendering(z_vals_bg, bg_sdf)
+            bg_rgb_values = torch.sum(bg_weights.unsqueeze(-1) * bg_rgb, 1)
+        else:
+            bg_rgb_values = torch.ones_like(fg_rgb_values, device=fg_rgb_values.device)
+
+        # Composite foreground and background
+        bg_rgb_values = bg_transmittance.unsqueeze(-1) * bg_rgb_values
+        rgb_values = fg_rgb_values + bg_rgb_values
+
+        normal_values = torch.sum(weights.unsqueeze(-1) * normal_values, 1)
+
+        if self.training:
+            output = {
+                'points': points,
+                'rgb_values': rgb_values,
+                'normal_values': normal_values,
+                'index_outside': input['index_outside'],
+                'index_off_surface': index_off_surface,
+                'index_in_surface': index_in_surface,
+                'acc_map': torch.sum(weights, -1),
+                'sdf_output': sdf_output,
+                'grad_theta': grad_theta,
+                'epoch': input['current_epoch'],
+            }
+        else:
+            fg_output_rgb = fg_rgb_values + bg_transmittance.unsqueeze(-1) * torch.ones_like(fg_rgb_values, device=fg_rgb_values.device)
+            output = {
+                'acc_map': torch.sum(weights, -1),
+                'rgb_values': rgb_values,
+                'fg_rgb_values': fg_output_rgb,
+                'normal_values': normal_values,
+                'sdf_output': sdf_output,
+            }
+        return output
+
+    def get_rbg_value(self, x, points, view_dirs, cond, tfs, feature_vectors, is_training=True):
+        pnts_c = points
+        others = {}
+
+        _, gradients, feature_vectors = self.forward_gradient(x, pnts_c, cond, tfs, create_graph=is_training, retain_graph=is_training)
+        # ensure the gradient is normalized
+        normals = nn.functional.normalize(gradients, dim=-1, eps=1e-6)
+        fg_rendering_output = self.rendering_network(pnts_c, normals, view_dirs, cond['smpl'],
+                                                     feature_vectors)
+        
+        rgb_vals = fg_rendering_output[:, :3]
+        others['normals'] = normals
+        return rgb_vals, others
+
+    def forward_gradient(self, x, pnts_c, cond, tfs, create_graph=True, retain_graph=True):
+        if pnts_c.shape[0] == 0:
+            return pnts_c.detach()
+        pnts_c.requires_grad_(True)
+
+        pnts_d = self.deformer.forward_skinning(pnts_c.unsqueeze(0), None, tfs).squeeze(0)
+        num_dim = pnts_d.shape[-1]
+        grads = []
+        for i in range(num_dim):
+            d_out = torch.zeros_like(pnts_d, requires_grad=False, device=pnts_d.device)
+            d_out[:, i] = 1
+            grad = torch.autograd.grad(
+                outputs=pnts_d,
+                inputs=pnts_c,
+                grad_outputs=d_out,
+                create_graph=create_graph,
+                retain_graph=True if i < num_dim - 1 else retain_graph,
+                only_inputs=True)[0]
+            grads.append(grad)
+        grads = torch.stack(grads, dim=-2)
+        grads_inv = grads.inverse()
+
+        output = self.implicit_network(pnts_c, cond)[0]
+        sdf = output[:, :1]
+        
+        feature = output[:, 1:]
+        d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)
+        gradients = torch.autograd.grad(
+            outputs=sdf,
+            inputs=pnts_c,
+            grad_outputs=d_output,
+            create_graph=create_graph,
+            retain_graph=retain_graph,
+            only_inputs=True)[0]
+
+        return grads.reshape(grads.shape[0], -1), torch.nn.functional.normalize(torch.einsum('bi,bij->bj', gradients, grads_inv), dim=1), feature
+
+    def volume_rendering(self, z_vals, z_max, sdf):
+        density_flat = self.density(sdf)
+        density = density_flat.reshape(-1, z_vals.shape[1]) # (batch_size * num_pixels) x N_samples
+
+        # included also the dist from the sphere intersection
+        dists = z_vals[:, 1:] - z_vals[:, :-1]
+        dists = torch.cat([dists, z_max.unsqueeze(-1) - z_vals[:, -1:]], -1)
+
+        # LOG SPACE
+        free_energy = dists * density
+        shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1).cuda(), free_energy], dim=-1)  # add 0 for transperancy 1 at t_0
+        alpha = 1 - torch.exp(-free_energy)  # probability of it is not empty here
+        transmittance = torch.exp(-torch.cumsum(shifted_free_energy, dim=-1))  # probability of everything is empty up to now
+        fg_transmittance = transmittance[:, :-1]
+        weights = alpha * fg_transmittance  # probability of the ray hits something here
+        bg_transmittance = transmittance[:, -1]  # factor to be multiplied with the bg volume rendering
+
+        return weights, bg_transmittance
+
+    def bg_volume_rendering(self, z_vals_bg, bg_sdf):
+        bg_density_flat = self.bg_density(bg_sdf)
+        bg_density = bg_density_flat.reshape(-1, z_vals_bg.shape[1]) # (batch_size * num_pixels) x N_samples
+
+        bg_dists = z_vals_bg[:, :-1] - z_vals_bg[:, 1:]
+        bg_dists = torch.cat([bg_dists, torch.tensor([1e10]).cuda().unsqueeze(0).repeat(bg_dists.shape[0], 1)], -1)
+
+        # LOG SPACE
+        bg_free_energy = bg_dists * bg_density
+        bg_shifted_free_energy = torch.cat([torch.zeros(bg_dists.shape[0], 1).cuda(), bg_free_energy[:, :-1]], dim=-1)  # shift one step
+        bg_alpha = 1 - torch.exp(-bg_free_energy)  # probability of it is not empty here
+        bg_transmittance = torch.exp(-torch.cumsum(bg_shifted_free_energy, dim=-1))  # probability of everything is empty up to now
+        bg_weights = bg_alpha * bg_transmittance # probability of the ray hits something here
+
+        return bg_weights
+    
+    def depth2pts_outside(self, ray_o, ray_d, depth):
+
+        '''
+        ray_o, ray_d: [..., 3]
+        depth: [...]; inverse of distance to sphere origin
+        '''
+
+        o_dot_d = torch.sum(ray_d * ray_o, dim=-1)
+        under_sqrt = o_dot_d ** 2 - ((ray_o ** 2).sum(-1) - self.sdf_bounding_sphere ** 2)
+        d_sphere = torch.sqrt(under_sqrt) - o_dot_d
+        p_sphere = ray_o + d_sphere.unsqueeze(-1) * ray_d
+        p_mid = ray_o - o_dot_d.unsqueeze(-1) * ray_d
+        p_mid_norm = torch.norm(p_mid, dim=-1)
+
+        rot_axis = torch.cross(ray_o, p_sphere, dim=-1)
+        rot_axis = rot_axis / torch.norm(rot_axis, dim=-1, keepdim=True)
+        phi = torch.asin(p_mid_norm / self.sdf_bounding_sphere)
+        theta = torch.asin(p_mid_norm * depth)  # depth is inside [0, 1]
+        rot_angle = (phi - theta).unsqueeze(-1)  # [..., 1]
+
+        # now rotate p_sphere
+        # Rodrigues formula: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
+        p_sphere_new = p_sphere * torch.cos(rot_angle) + \
+                       torch.cross(rot_axis, p_sphere, dim=-1) * torch.sin(rot_angle) + \
+                       rot_axis * torch.sum(rot_axis * p_sphere, dim=-1, keepdim=True) * (1. - torch.cos(rot_angle))
+        p_sphere_new = p_sphere_new / torch.norm(p_sphere_new, dim=-1, keepdim=True)
+        pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)
+
+        return pts
+
+def gradient(inputs, outputs):
+
+    d_points = torch.ones_like(outputs, requires_grad=False, device=outputs.device)
+    points_grad = grad(
+        outputs=outputs,
+        inputs=inputs,
+        grad_outputs=d_points,
+        create_graph=True,
+        retain_graph=True,
+        only_inputs=True)[0][:, :, -3:]
+    return points_grad

code/lib/smpl/body_models.py

@@ -0,0 +1,365 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import os
+import os.path as osp
+
+
+import pickle
+
+import numpy as np
+
+from collections import namedtuple
+
+import torch
+import torch.nn as nn
+
+from .lbs import (
+    lbs, vertices2joints, blend_shapes)
+
+from .vertex_ids import vertex_ids as VERTEX_IDS
+from .utils import Struct, to_np, to_tensor
+from .vertex_joint_selector import VertexJointSelector
+
+
+ModelOutput = namedtuple('ModelOutput',
+                         ['vertices','faces', 'joints', 'full_pose', 'betas',
+                          'global_orient',
+                          'body_pose', 'expression',
+                          'left_hand_pose', 'right_hand_pose',
+                          'jaw_pose', 'T', 'T_weighted', 'weights'])
+ModelOutput.__new__.__defaults__ = (None,) * len(ModelOutput._fields)
+
+class SMPL(nn.Module):
+
+    NUM_JOINTS = 23
+    NUM_BODY_JOINTS = 23
+    NUM_BETAS = 10
+
+    def __init__(self, model_path, data_struct=None,
+                 create_betas=True,
+                 betas=None,
+                 create_global_orient=True,
+                 global_orient=None,
+                 create_body_pose=True,
+                 body_pose=None,
+                 create_transl=True,
+                 transl=None,
+                 dtype=torch.float32,
+                 batch_size=1,
+                 joint_mapper=None, gender='neutral',
+                 vertex_ids=None,
+                 pose_blend=True,
+                 **kwargs):
+        ''' SMPL model constructor
+
+            Parameters
+            ----------
+            model_path: str
+                The path to the folder or to the file where the model
+                parameters are stored
+            data_struct: Strct
+                A struct object. If given, then the parameters of the model are
+                read from the object. Otherwise, the model tries to read the
+                parameters from the given `model_path`. (default = None)
+            create_global_orient: bool, optional
+                Flag for creating a member variable for the global orientation
+                of the body. (default = True)
+            global_orient: torch.tensor, optional, Bx3
+                The default value for the global orientation variable.
+                (default = None)
+            create_body_pose: bool, optional
+                Flag for creating a member variable for the pose of the body.
+                (default = True)
+            body_pose: torch.tensor, optional, Bx(Body Joints * 3)
+                The default value for the body pose variable.
+                (default = None)
+            create_betas: bool, optional
+                Flag for creating a member variable for the shape space
+                (default = True).
+            betas: torch.tensor, optional, Bx10
+                The default value for the shape member variable.
+                (default = None)
+            create_transl: bool, optional
+                Flag for creating a member variable for the translation
+                of the body. (default = True)
+            transl: torch.tensor, optional, Bx3
+                The default value for the transl variable.
+                (default = None)
+            dtype: torch.dtype, optional
+                The data type for the created variables
+            batch_size: int, optional
+                The batch size used for creating the member variables
+            joint_mapper: object, optional
+                An object that re-maps the joints. Useful if one wants to
+                re-order the SMPL joints to some other convention (e.g. MSCOCO)
+                (default = None)
+            gender: str, optional
+                Which gender to load
+            vertex_ids: dict, optional
+                A dictionary containing the indices of the extra vertices that
+                will be selected
+        '''
+
+        self.gender = gender
+        self.pose_blend = pose_blend
+
+        if data_struct is None:
+            if osp.isdir(model_path):
+                model_fn = 'SMPL_{}.{ext}'.format(gender.upper(), ext='pkl')
+                smpl_path = os.path.join(model_path, model_fn)
+            else:
+                smpl_path = model_path
+            assert osp.exists(smpl_path), 'Path {} does not exist!'.format(
+                smpl_path)
+
+            with open(smpl_path, 'rb') as smpl_file:
+                data_struct = Struct(**pickle.load(smpl_file,encoding='latin1'))
+        super(SMPL, self).__init__()
+        self.batch_size = batch_size
+
+        if vertex_ids is None:
+            # SMPL and SMPL-H share the same topology, so any extra joints can
+            # be drawn from the same place
+            vertex_ids = VERTEX_IDS['smplh']
+
+        self.dtype = dtype
+
+        self.joint_mapper = joint_mapper
+
+        self.vertex_joint_selector = VertexJointSelector(
+            vertex_ids=vertex_ids, **kwargs)
+
+        self.faces = data_struct.f
+        self.register_buffer('faces_tensor',
+                             to_tensor(to_np(self.faces, dtype=np.int64),
+                                       dtype=torch.long))
+
+        if create_betas:
+            if betas is None:
+                default_betas = torch.zeros([batch_size, self.NUM_BETAS],
+                                            dtype=dtype)
+            else:
+                if 'torch.Tensor' in str(type(betas)):
+                    default_betas = betas.clone().detach()
+                else:
+                    default_betas = torch.tensor(betas,
+                                                 dtype=dtype)
+
+            self.register_parameter('betas', nn.Parameter(default_betas,
+                                                          requires_grad=True))
+
+        # The tensor that contains the global rotation of the model
+        # It is separated from the pose of the joints in case we wish to
+        # optimize only over one of them
+        if create_global_orient:
+            if global_orient is None:
+                default_global_orient = torch.zeros([batch_size, 3],
+                                                    dtype=dtype)
+            else:
+                if 'torch.Tensor' in str(type(global_orient)):
+                    default_global_orient = global_orient.clone().detach()
+                else:
+                    default_global_orient = torch.tensor(global_orient,
+                                                         dtype=dtype)
+
+            global_orient = nn.Parameter(default_global_orient,
+                                         requires_grad=True)
+            self.register_parameter('global_orient', global_orient)
+
+        if create_body_pose:
+            if body_pose is None:
+                default_body_pose = torch.zeros(
+                    [batch_size, self.NUM_BODY_JOINTS * 3], dtype=dtype)
+            else:
+                if 'torch.Tensor' in str(type(body_pose)):
+                    default_body_pose = body_pose.clone().detach()
+                else:
+                    default_body_pose = torch.tensor(body_pose,
+                                                     dtype=dtype)
+            self.register_parameter(
+                'body_pose',
+                nn.Parameter(default_body_pose, requires_grad=True))
+
+        if create_transl:
+            if transl is None:
+                default_transl = torch.zeros([batch_size, 3],
+                                             dtype=dtype,
+                                             requires_grad=True)
+            else:
+                default_transl = torch.tensor(transl, dtype=dtype)
+            self.register_parameter(
+                'transl',
+                nn.Parameter(default_transl, requires_grad=True))
+
+        # The vertices of the template model
+        self.register_buffer('v_template',
+                             to_tensor(to_np(data_struct.v_template),
+                                       dtype=dtype))
+
+        # The shape components
+        shapedirs = data_struct.shapedirs[:, :, :self.NUM_BETAS]
+        # The shape components
+        self.register_buffer(
+            'shapedirs',
+            to_tensor(to_np(shapedirs), dtype=dtype))
+
+
+        j_regressor = to_tensor(to_np(
+            data_struct.J_regressor), dtype=dtype)
+        self.register_buffer('J_regressor', j_regressor)
+
+        # if self.gender == 'neutral':
+        #     joint_regressor = to_tensor(to_np(
+        #     data_struct.cocoplus_regressor), dtype=dtype).permute(1,0)
+        #     self.register_buffer('joint_regressor', joint_regressor)
+
+        # Pose blend shape basis: 6890 x 3 x 207, reshaped to 6890*3 x 207
+        num_pose_basis = data_struct.posedirs.shape[-1]
+        # 207 x 20670
+        posedirs = np.reshape(data_struct.posedirs, [-1, num_pose_basis]).T
+        self.register_buffer('posedirs',
+                             to_tensor(to_np(posedirs), dtype=dtype))
+
+        # indices of parents for each joints
+        parents = to_tensor(to_np(data_struct.kintree_table[0])).long()
+        parents[0] = -1
+        self.register_buffer('parents', parents)
+
+        self.bone_parents = to_np(data_struct.kintree_table[0])
+
+        self.register_buffer('lbs_weights',
+                             to_tensor(to_np(data_struct.weights), dtype=dtype))
+
+    def create_mean_pose(self, data_struct):
+        pass
+
+    @torch.no_grad()
+    def reset_params(self, **params_dict):
+        for param_name, param in self.named_parameters():
+            if param_name in params_dict:
+                param[:] = torch.tensor(params_dict[param_name])
+            else:
+                param.fill_(0)
+
+    def get_T_hip(self, betas=None):
+        v_shaped = self.v_template + blend_shapes(betas, self.shapedirs)
+        J = vertices2joints(self.J_regressor, v_shaped)
+        T_hip = J[0,0]
+        return T_hip
+
+    def get_num_verts(self):
+        return self.v_template.shape[0]
+
+    def get_num_faces(self):
+        return self.faces.shape[0]
+
+    def extra_repr(self):
+        return 'Number of betas: {}'.format(self.NUM_BETAS)
+
+    def forward(self, betas=None, body_pose=None, global_orient=None,
+                transl=None, return_verts=True, return_full_pose=False,displacement=None,v_template=None,
+                **kwargs):
+        ''' Forward pass for the SMPL model
+
+            Parameters
+            ----------
+            global_orient: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable and use it as the global
+                rotation of the body. Useful if someone wishes to predicts this
+                with an external model. (default=None)
+            betas: torch.tensor, optional, shape Bx10
+                If given, ignore the member variable `betas` and use it
+                instead. For example, it can used if shape parameters
+                `betas` are predicted from some external model.
+                (default=None)
+            body_pose: torch.tensor, optional, shape Bx(J*3)
+                If given, ignore the member variable `body_pose` and use it
+                instead. For example, it can used if someone predicts the
+                pose of the body joints are predicted from some external model.
+                It should be a tensor that contains joint rotations in
+                axis-angle format. (default=None)
+            transl: torch.tensor, optional, shape Bx3
+                If given, ignore the member variable `transl` and use it
+                instead. For example, it can used if the translation
+                `transl` is predicted from some external model.
+                (default=None)
+            return_verts: bool, optional
+                Return the vertices. (default=True)
+            return_full_pose: bool, optional
+                Returns the full axis-angle pose vector (default=False)
+
+            Returns
+            -------
+        '''
+        # If no shape and pose parameters are passed along, then use the
+        # ones from the module
+        global_orient = (global_orient if global_orient is not None else
+                         self.global_orient)
+        body_pose = body_pose if body_pose is not None else self.body_pose
+        betas = betas if betas is not None else self.betas
+
+        apply_trans = transl is not None or hasattr(self, 'transl')
+        if transl is None and hasattr(self, 'transl'):
+            transl = self.transl
+
+        full_pose = torch.cat([global_orient, body_pose], dim=1)
+
+        # if betas.shape[0] != self.batch_size:
+        #     num_repeats = int(self.batch_size / betas.shape[0])
+        #     betas = betas.expand(num_repeats, -1)
+
+        if v_template is None:
+            v_template = self.v_template
+
+        if displacement is not None:
+            vertices, joints_smpl, T_weighted, W, T = lbs(betas, full_pose, v_template+displacement,
+                                   self.shapedirs, self.posedirs,
+                                   self.J_regressor, self.parents,
+                                   self.lbs_weights, dtype=self.dtype,pose_blend=self.pose_blend)
+        else:
+            vertices, joints_smpl,T_weighted, W, T = lbs(betas, full_pose, v_template,
+                                        self.shapedirs, self.posedirs,
+                                        self.J_regressor, self.parents,
+                                        self.lbs_weights, dtype=self.dtype,pose_blend=self.pose_blend)
+
+        # if self.gender is not 'neutral':
+        joints = self.vertex_joint_selector(vertices, joints_smpl)
+        # else:
+        # joints = torch.matmul(vertices.permute(0,2,1),self.joint_regressor).permute(0,2,1)
+        # Map the joints to the current dataset
+        if self.joint_mapper is not None:
+            joints = self.joint_mapper(joints)
+
+        if apply_trans:
+            joints_smpl = joints_smpl + transl.unsqueeze(dim=1)
+            joints = joints + transl.unsqueeze(dim=1)
+            vertices = vertices + transl.unsqueeze(dim=1)
+
+        output = ModelOutput(vertices=vertices if return_verts else None,
+                             faces=self.faces,
+                             global_orient=global_orient,
+                             body_pose=body_pose,
+                             joints=joints_smpl,
+                             betas=self.betas,
+                             full_pose=full_pose if return_full_pose else None,
+                             T=T, T_weighted=T_weighted, weights=W)
+        return output

code/lib/smpl/lbs.py

@@ -0,0 +1,377 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+from .utils import rot_mat_to_euler
+
+
+def find_dynamic_lmk_idx_and_bcoords(vertices, pose, dynamic_lmk_faces_idx,
+                                     dynamic_lmk_b_coords,
+                                     neck_kin_chain, dtype=torch.float32):
+    ''' Compute the faces, barycentric coordinates for the dynamic landmarks
+
+
+        To do so, we first compute the rotation of the neck around the y-axis
+        and then use a pre-computed look-up table to find the faces and the
+        barycentric coordinates that will be used.
+
+        Special thanks to Soubhik Sanyal (soubhik.sanyal@tuebingen.mpg.de)
+        for providing the original TensorFlow implementation and for the LUT.
+
+        Parameters
+        ----------
+        vertices: torch.tensor BxVx3, dtype = torch.float32
+            The tensor of input vertices
+        pose: torch.tensor Bx(Jx3), dtype = torch.float32
+            The current pose of the body model
+        dynamic_lmk_faces_idx: torch.tensor L, dtype = torch.long
+            The look-up table from neck rotation to faces
+        dynamic_lmk_b_coords: torch.tensor Lx3, dtype = torch.float32
+            The look-up table from neck rotation to barycentric coordinates
+        neck_kin_chain: list
+            A python list that contains the indices of the joints that form the
+            kinematic chain of the neck.
+        dtype: torch.dtype, optional
+
+        Returns
+        -------
+        dyn_lmk_faces_idx: torch.tensor, dtype = torch.long
+            A tensor of size BxL that contains the indices of the faces that
+            will be used to compute the current dynamic landmarks.
+        dyn_lmk_b_coords: torch.tensor, dtype = torch.float32
+            A tensor of size BxL that contains the indices of the faces that
+            will be used to compute the current dynamic landmarks.
+    '''
+
+    batch_size = vertices.shape[0]
+
+    aa_pose = torch.index_select(pose.view(batch_size, -1, 3), 1,
+                                 neck_kin_chain)
+    rot_mats = batch_rodrigues(
+        aa_pose.view(-1, 3), dtype=dtype).view(batch_size, -1, 3, 3)
+
+    rel_rot_mat = torch.eye(3, device=vertices.device,
+                            dtype=dtype).unsqueeze_(dim=0)
+    for idx in range(len(neck_kin_chain)):
+        rel_rot_mat = torch.bmm(rot_mats[:, idx], rel_rot_mat)
+
+    y_rot_angle = torch.round(
+        torch.clamp(-rot_mat_to_euler(rel_rot_mat) * 180.0 / np.pi,
+                    max=39)).to(dtype=torch.long)
+    neg_mask = y_rot_angle.lt(0).to(dtype=torch.long)
+    mask = y_rot_angle.lt(-39).to(dtype=torch.long)
+    neg_vals = mask * 78 + (1 - mask) * (39 - y_rot_angle)
+    y_rot_angle = (neg_mask * neg_vals +
+                   (1 - neg_mask) * y_rot_angle)
+
+    dyn_lmk_faces_idx = torch.index_select(dynamic_lmk_faces_idx,
+                                           0, y_rot_angle)
+    dyn_lmk_b_coords = torch.index_select(dynamic_lmk_b_coords,
+                                          0, y_rot_angle)
+
+    return dyn_lmk_faces_idx, dyn_lmk_b_coords
+
+
+def vertices2landmarks(vertices, faces, lmk_faces_idx, lmk_bary_coords):
+    ''' Calculates landmarks by barycentric interpolation
+
+        Parameters
+        ----------
+        vertices: torch.tensor BxVx3, dtype = torch.float32
+            The tensor of input vertices
+        faces: torch.tensor Fx3, dtype = torch.long
+            The faces of the mesh
+        lmk_faces_idx: torch.tensor L, dtype = torch.long
+            The tensor with the indices of the faces used to calculate the
+            landmarks.
+        lmk_bary_coords: torch.tensor Lx3, dtype = torch.float32
+            The tensor of barycentric coordinates that are used to interpolate
+            the landmarks
+
+        Returns
+        -------
+        landmarks: torch.tensor BxLx3, dtype = torch.float32
+            The coordinates of the landmarks for each mesh in the batch
+    '''
+    # Extract the indices of the vertices for each face
+    # BxLx3
+    batch_size, num_verts = vertices.shape[:2]
+    device = vertices.device
+
+    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).expand(
+        batch_size, -1, -1).long()
+
+    lmk_faces = lmk_faces +  torch.arange(
+        batch_size, dtype=torch.long, device=device).view(-1, 1, 1) * num_verts
+
+    lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(
+        batch_size, -1, 3, 3)
+
+    landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])
+    return landmarks
+
+
+def lbs(betas, pose, v_template, shapedirs, posedirs, J_regressor, parents,
+        lbs_weights, pose2rot=True, dtype=torch.float32, pose_blend=True):
+    ''' Performs Linear Blend Skinning with the given shape and pose parameters
+
+        Parameters
+        ----------
+        betas : torch.tensor BxNB
+            The tensor of shape parameters
+        pose : torch.tensor Bx(J + 1) * 3
+            The pose parameters in axis-angle format
+        v_template torch.tensor BxVx3
+            The template mesh that will be deformed
+        shapedirs : torch.tensor 1xNB
+            The tensor of PCA shape displacements
+        posedirs : torch.tensor Px(V * 3)
+            The pose PCA coefficients
+        J_regressor : torch.tensor JxV
+            The regressor array that is used to calculate the joints from
+            the position of the vertices
+        parents: torch.tensor J
+            The array that describes the kinematic tree for the model
+        lbs_weights: torch.tensor N x V x (J + 1)
+            The linear blend skinning weights that represent how much the
+            rotation matrix of each part affects each vertex
+        pose2rot: bool, optional
+            Flag on whether to convert the input pose tensor to rotation
+            matrices. The default value is True. If False, then the pose tensor
+            should already contain rotation matrices and have a size of
+            Bx(J + 1)x9
+        dtype: torch.dtype, optional
+
+        Returns
+        -------
+        verts: torch.tensor BxVx3
+            The vertices of the mesh after applying the shape and pose
+            displacements.
+        joints: torch.tensor BxJx3
+            The joints of the model
+    '''
+
+    batch_size = max(betas.shape[0], pose.shape[0])
+    device = betas.device
+
+    # Add shape contribution
+    v_shaped = v_template + blend_shapes(betas, shapedirs)
+
+    # Get the joints
+    # NxJx3 array
+    J = vertices2joints(J_regressor, v_shaped)
+
+    # 3. Add pose blend shapes
+    # N x J x 3 x 3
+    ident = torch.eye(3, dtype=dtype, device=device)
+
+
+    if pose2rot:
+        rot_mats = batch_rodrigues(
+            pose.view(-1, 3), dtype=dtype).view([batch_size, -1, 3, 3])
+
+        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
+        # (N x P) x (P, V * 3) -> N x V x 3
+        pose_offsets = torch.matmul(pose_feature, posedirs) \
+            .view(batch_size, -1, 3)
+    else:
+        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
+        rot_mats = pose.view(batch_size, -1, 3, 3)
+
+        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1),
+                                    posedirs).view(batch_size, -1, 3)
+
+    if pose_blend:
+        v_posed = pose_offsets + v_shaped
+    else:
+        v_posed = v_shaped
+
+    # 4. Get the global joint location
+    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)
+
+    # 5. Do skinning:
+    # W is N x V x (J + 1)
+    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
+    # (N x V x (J + 1)) x (N x (J + 1) x 16)
+    num_joints = J_regressor.shape[0]
+    T = torch.matmul(W, A.view(batch_size, num_joints, 16)) \
+        .view(batch_size, -1, 4, 4)
+
+    homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1],
+                               dtype=dtype, device=device)
+    v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
+    v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
+
+    verts = v_homo[:, :, :3, 0]
+
+    return verts, J_transformed, T, W, A.view(batch_size, num_joints, 4,4)
+
+
+def vertices2joints(J_regressor, vertices):
+    ''' Calculates the 3D joint locations from the vertices
+
+    Parameters
+    ----------
+    J_regressor : torch.tensor JxV
+        The regressor array that is used to calculate the joints from the
+        position of the vertices
+    vertices : torch.tensor BxVx3
+        The tensor of mesh vertices
+
+    Returns
+    -------
+    torch.tensor BxJx3
+        The location of the joints
+    '''
+
+    return torch.einsum('bik,ji->bjk', [vertices, J_regressor])
+
+
+def blend_shapes(betas, shape_disps):
+    ''' Calculates the per vertex displacement due to the blend shapes
+
+
+    Parameters
+    ----------
+    betas : torch.tensor Bx(num_betas)
+        Blend shape coefficients
+    shape_disps: torch.tensor Vx3x(num_betas)
+        Blend shapes
+
+    Returns
+    -------
+    torch.tensor BxVx3
+        The per-vertex displacement due to shape deformation
+    '''
+
+    # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
+    # i.e. Multiply each shape displacement by its corresponding beta and
+    # then sum them.
+    blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
+    return blend_shape
+
+
+def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
+    ''' Calculates the rotation matrices for a batch of rotation vectors
+        Parameters
+        ----------
+        rot_vecs: torch.tensor Nx3
+            array of N axis-angle vectors
+        Returns
+        -------
+        R: torch.tensor Nx3x3
+            The rotation matrices for the given axis-angle parameters
+    '''
+
+    batch_size = rot_vecs.shape[0]
+    device = rot_vecs.device
+
+    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
+    rot_dir = rot_vecs / angle
+
+    cos = torch.unsqueeze(torch.cos(angle), dim=1)
+    sin = torch.unsqueeze(torch.sin(angle), dim=1)
+
+    # Bx1 arrays
+    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
+    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
+
+    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
+    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
+        .view((batch_size, 3, 3))
+
+    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
+    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
+    return rot_mat
+
+
+def transform_mat(R, t):
+    ''' Creates a batch of transformation matrices
+        Args:
+            - R: Bx3x3 array of a batch of rotation matrices
+            - t: Bx3x1 array of a batch of translation vectors
+        Returns:
+            - T: Bx4x4 Transformation matrix
+    '''
+    # No padding left or right, only add an extra row
+    return torch.cat([F.pad(R, [0, 0, 0, 1]),
+                      F.pad(t, [0, 0, 0, 1], value=1)], dim=2)
+
+
+def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
+    """
+    Applies a batch of rigid transformations to the joints
+
+    Parameters
+    ----------
+    rot_mats : torch.tensor BxNx3x3
+        Tensor of rotation matrices
+    joints : torch.tensor BxNx3
+        Locations of joints
+    parents : torch.tensor BxN
+        The kinematic tree of each object
+    dtype : torch.dtype, optional:
+        The data type of the created tensors, the default is torch.float32
+
+    Returns
+    -------
+    posed_joints : torch.tensor BxNx3
+        The locations of the joints after applying the pose rotations
+    rel_transforms : torch.tensor BxNx4x4
+        The relative (with respect to the root joint) rigid transformations
+        for all the joints
+    """
+
+    joints = torch.unsqueeze(joints, dim=-1)
+
+    rel_joints = joints.clone()
+    rel_joints[:, 1:] = rel_joints[:, 1:] - joints[:, parents[1:]]
+
+    transforms_mat = transform_mat(
+        rot_mats.reshape(-1, 3, 3),
+        rel_joints.reshape(-1, 3, 1)).reshape(-1, joints.shape[1], 4, 4)
+
+    transform_chain = [transforms_mat[:, 0]]
+    for i in range(1, parents.shape[0]):
+        # Subtract the joint location at the rest pose
+        # No need for rotation, since it's identity when at rest
+        curr_res = torch.matmul(transform_chain[parents[i]],
+                                transforms_mat[:, i])
+        transform_chain.append(curr_res)
+
+    transforms = torch.stack(transform_chain, dim=1)
+
+    # The last column of the transformations contains the posed joints
+    posed_joints = transforms[:, :, :3, 3]
+
+    # The last column of the transformations contains the posed joints
+    posed_joints = transforms[:, :, :3, 3]
+
+    joints_homogen = F.pad(joints, [0, 0, 0, 1])
+
+    rel_transforms = transforms - F.pad(
+        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0])
+
+    return posed_joints, rel_transforms

code/lib/smpl/smpl_model/SMPL_FEMALE.pkl

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

code/lib/smpl/smpl_model/SMPL_MALE.pkl

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

code/lib/smpl/utils.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import print_function
+from __future__ import absolute_import
+from __future__ import division
+
+import numpy as np
+import torch
+
+
+def to_tensor(array, dtype=torch.float32):
+    if 'torch.tensor' not in str(type(array)):
+        return torch.tensor(array, dtype=dtype)
+
+
+class Struct(object):
+    def __init__(self, **kwargs):
+        for key, val in kwargs.items():
+            setattr(self, key, val)
+
+
+def to_np(array, dtype=np.float32):
+    if 'scipy.sparse' in str(type(array)):
+        array = array.todense()
+    return np.array(array, dtype=dtype)
+
+
+def rot_mat_to_euler(rot_mats):
+    # Calculates rotation matrix to euler angles
+    # Careful for extreme cases of eular angles like [0.0, pi, 0.0]
+
+    sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
+                    rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
+    return torch.atan2(-rot_mats[:, 2, 0], sy)

code/lib/smpl/vertex_ids.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import print_function
+from __future__ import absolute_import
+from __future__ import division
+
+# Joint name to vertex mapping. SMPL/SMPL-H/SMPL-X vertices that correspond to
+# MSCOCO and OpenPose joints
+vertex_ids = {
+    'smplh': {
+        'nose':		    332,
+        'reye':		    6260,
+        'leye':		    2800,
+        'rear':		    4071,
+        'lear':		    583,
+        'rthumb':		6191,
+        'rindex':		5782,
+        'rmiddle':		5905,
+        'rring':		6016,
+        'rpinky':		6133,
+        'lthumb':		2746,
+        'lindex':		2319,
+        'lmiddle':		2445,
+        'lring':		2556,
+        'lpinky':		2673,
+        'LBigToe':		3216,
+        'LSmallToe':	3226,
+        'LHeel':		3387,
+        'RBigToe':		6617,
+        'RSmallToe':    6624,
+        'RHeel':		6787
+    },
+    'smplx': {
+        'nose':		    9120,
+        'reye':		    9929,
+        'leye':		    9448,
+        'rear':		    616,
+        'lear':		    6,
+        'rthumb':		8079,
+        'rindex':		7669,
+        'rmiddle':		7794,
+        'rring':		7905,
+        'rpinky':		8022,
+        'lthumb':		5361,
+        'lindex':		4933,
+        'lmiddle':		5058,
+        'lring':		5169,
+        'lpinky':		5286,
+        'LBigToe':		5770,
+        'LSmallToe':    5780,
+        'LHeel':		8846,
+        'RBigToe':		8463,
+        'RSmallToe': 	8474,
+        'RHeel':  		8635
+    }
+}

code/lib/smpl/vertex_joint_selector.py

@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems and the Max Planck Institute for Biological
+# Cybernetics. All rights reserved.
+#
+# Contact: ps-license@tuebingen.mpg.de
+
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from .utils import to_tensor
+
+
+class VertexJointSelector(nn.Module):
+
+    def __init__(self, vertex_ids=None,
+                 use_hands=True,
+                 use_feet_keypoints=True, **kwargs):
+        super(VertexJointSelector, self).__init__()
+
+        extra_joints_idxs = []
+
+        face_keyp_idxs = np.array([
+            vertex_ids['nose'],
+            vertex_ids['reye'],
+            vertex_ids['leye'],
+            vertex_ids['rear'],
+            vertex_ids['lear']], dtype=np.int64)
+
+        extra_joints_idxs = np.concatenate([extra_joints_idxs,
+                                            face_keyp_idxs])
+
+        if use_feet_keypoints:
+            feet_keyp_idxs = np.array([vertex_ids['LBigToe'],
+                                       vertex_ids['LSmallToe'],
+                                       vertex_ids['LHeel'],
+                                       vertex_ids['RBigToe'],
+                                       vertex_ids['RSmallToe'],
+                                       vertex_ids['RHeel']], dtype=np.int32)
+
+            extra_joints_idxs = np.concatenate(
+                [extra_joints_idxs, feet_keyp_idxs])
+
+        if use_hands:
+            self.tip_names = ['thumb', 'index', 'middle', 'ring', 'pinky']
+
+            tips_idxs = []
+            for hand_id in ['l', 'r']:
+                for tip_name in self.tip_names:
+                    tips_idxs.append(vertex_ids[hand_id + tip_name])
+
+            extra_joints_idxs = np.concatenate(
+                [extra_joints_idxs, tips_idxs])
+
+        self.register_buffer('extra_joints_idxs',
+                             to_tensor(extra_joints_idxs, dtype=torch.long))
+
+    def forward(self, vertices, joints):
+        extra_joints = torch.index_select(vertices, 1, self.extra_joints_idxs)
+        joints = torch.cat([joints, extra_joints], dim=1)
+        return joints

code/lib/utils/meshing.py

@@ -0,0 +1,63 @@
+import numpy as np
+import torch
+from skimage import measure
+from lib.libmise import mise
+import trimesh
+
+def generate_mesh(func, verts, level_set=0, res_init=32, res_up=3, point_batch=5000):
+    
+    scale = 1.1  # Scale of the padded bbox regarding the tight one.
+    verts = verts.data.cpu().numpy()
+    
+    gt_bbox = np.stack([verts.min(axis=0), verts.max(axis=0)], axis=0)
+    gt_center = (gt_bbox[0] + gt_bbox[1]) * 0.5
+    gt_scale = (gt_bbox[1] - gt_bbox[0]).max()
+
+    mesh_extractor = mise.MISE(res_init, res_up, level_set)
+    points = mesh_extractor.query()
+    
+    # query occupancy grid
+    while points.shape[0] != 0:
+        
+        orig_points = points
+        points = points.astype(np.float32)
+        points = (points / mesh_extractor.resolution - 0.5) * scale
+        points = points * gt_scale + gt_center
+        points = torch.tensor(points).float().cuda()
+        
+        values = []
+        for _, pnts in enumerate((torch.split(points,point_batch,dim=0))):
+            out = func(pnts)
+            values.append(out['sdf'].data.cpu().numpy())
+        values = np.concatenate(values, axis=0).astype(np.float64)[:,0]
+        
+        mesh_extractor.update(orig_points, values)
+        
+        points = mesh_extractor.query()
+    
+    value_grid = mesh_extractor.to_dense()
+
+    # marching cube
+    verts, faces, normals, values = measure.marching_cubes_lewiner(
+                                                volume=value_grid,
+                                                gradient_direction='ascent',
+                                                level=level_set)
+
+    verts = (verts / mesh_extractor.resolution - 0.5) * scale
+    verts = verts * gt_scale + gt_center
+    faces = faces[:, [0,2,1]]
+    meshexport = trimesh.Trimesh(verts, faces, normals, vertex_colors=values)
+
+    #remove disconnect part
+    connected_comp = meshexport.split(only_watertight=False)
+    max_area = 0
+    max_comp = None
+    for comp in connected_comp:
+        if comp.area > max_area:
+            max_area = comp.area
+            max_comp = comp
+    meshexport = max_comp
+
+    return meshexport
+
+

code/lib/utils/utils.py

@@ -0,0 +1,232 @@
+import numpy as np
+import cv2
+import torch
+from torch.nn import functional as F
+
+
+def split_input(model_input, total_pixels, n_pixels = 10000):
+    '''
+     Split the input to fit Cuda memory for large resolution.
+     Can decrease the value of n_pixels in case of cuda out of memory error.
+     '''
+
+    split = []
+
+    for i, indx in enumerate(torch.split(torch.arange(total_pixels).cuda(), n_pixels, dim=0)):
+        data = model_input.copy()
+        data['uv'] = torch.index_select(model_input['uv'], 1, indx)
+        split.append(data)
+    return split
+
+
+def merge_output(res, total_pixels, batch_size):
+    ''' Merge the split output. '''
+
+    model_outputs = {}
+    for entry in res[0]:
+        if res[0][entry] is None:
+            continue
+        if len(res[0][entry].shape) == 1:
+            model_outputs[entry] = torch.cat([r[entry].reshape(batch_size, -1, 1) for r in res],
+                                             1).reshape(batch_size * total_pixels)
+        else:
+            model_outputs[entry] = torch.cat([r[entry].reshape(batch_size, -1, r[entry].shape[-1]) for r in res],
+                                             1).reshape(batch_size * total_pixels, -1)
+    return model_outputs
+
+
+def get_psnr(img1, img2, normalize_rgb=False):
+    if normalize_rgb: # [-1,1] --> [0,1]
+        img1 = (img1 + 1.) / 2.
+        img2 = (img2 + 1. ) / 2.
+
+    mse = torch.mean((img1 - img2) ** 2)
+    psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]).cuda())
+
+    return psnr
+
+
+def load_K_Rt_from_P(filename, P=None):
+    if P is None:
+        lines = open(filename).read().splitlines()
+        if len(lines) == 4:
+            lines = lines[1:]
+        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
+        P = np.asarray(lines).astype(np.float32).squeeze()
+
+    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 get_camera_params(uv, pose, intrinsics):
+    if pose.shape[1] == 7: #In case of quaternion vector representation
+        cam_loc = pose[:, 4:]
+        R = quat_to_rot(pose[:,:4])
+        p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float()
+        p[:, :3, :3] = R
+        p[:, :3, 3] = cam_loc
+    else: # In case of pose matrix representation
+        cam_loc = pose[:, :3, 3]
+        p = pose
+
+    batch_size, num_samples, _ = uv.shape
+
+    depth = torch.ones((batch_size, num_samples)).cuda()
+    x_cam = uv[:, :, 0].view(batch_size, -1)
+    y_cam = uv[:, :, 1].view(batch_size, -1)
+    z_cam = depth.view(batch_size, -1)
+
+    pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics)
+
+    # permute for batch matrix product
+    pixel_points_cam = pixel_points_cam.permute(0, 2, 1)
+
+    world_coords = torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3]
+    ray_dirs = world_coords - cam_loc[:, None, :]
+    ray_dirs = F.normalize(ray_dirs, dim=2)
+
+    return ray_dirs, cam_loc
+
+def lift(x, y, z, intrinsics):
+    # parse intrinsics
+    intrinsics = intrinsics.cuda()
+    fx = intrinsics[:, 0, 0]
+    fy = intrinsics[:, 1, 1]
+    cx = intrinsics[:, 0, 2]
+    cy = intrinsics[:, 1, 2]
+    sk = intrinsics[:, 0, 1]
+
+    x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z
+    y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z
+
+    # homogeneous
+    return torch.stack((x_lift, y_lift, z, torch.ones_like(z).cuda()), dim=-1)
+
+
+def quat_to_rot(q):
+    batch_size, _ = q.shape
+    q = F.normalize(q, dim=1)
+    R = torch.ones((batch_size, 3,3)).cuda()
+    qr=q[:,0]
+    qi = q[:, 1]
+    qj = q[:, 2]
+    qk = q[:, 3]
+    R[:, 0, 0]=1-2 * (qj**2 + qk**2)
+    R[:, 0, 1] = 2 * (qj *qi -qk*qr)
+    R[:, 0, 2] = 2 * (qi * qk + qr * qj)
+    R[:, 1, 0] = 2 * (qj * qi + qk * qr)
+    R[:, 1, 1] = 1-2 * (qi**2 + qk**2)
+    R[:, 1, 2] = 2*(qj*qk - qi*qr)
+    R[:, 2, 0] = 2 * (qk * qi-qj * qr)
+    R[:, 2, 1] = 2 * (qj*qk + qi*qr)
+    R[:, 2, 2] = 1-2 * (qi**2 + qj**2)
+    return R
+
+
+def rot_to_quat(R):
+    batch_size, _,_ = R.shape
+    q = torch.ones((batch_size, 4)).cuda()
+
+    R00 = R[:, 0,0]
+    R01 = R[:, 0, 1]
+    R02 = R[:, 0, 2]
+    R10 = R[:, 1, 0]
+    R11 = R[:, 1, 1]
+    R12 = R[:, 1, 2]
+    R20 = R[:, 2, 0]
+    R21 = R[:, 2, 1]
+    R22 = R[:, 2, 2]
+
+    q[:,0]=torch.sqrt(1.0+R00+R11+R22)/2
+    q[:, 1]=(R21-R12)/(4*q[:,0])
+    q[:, 2] = (R02 - R20) / (4 * q[:, 0])
+    q[:, 3] = (R10 - R01) / (4 * q[:, 0])
+    return q
+
+
+def get_sphere_intersections(cam_loc, ray_directions, r = 1.0):
+    # Input: n_rays x 3 ; n_rays x 3
+    # Output: n_rays x 1, n_rays x 1 (close and far)
+
+    ray_cam_dot = torch.bmm(ray_directions.view(-1, 1, 3),
+                            cam_loc.view(-1, 3, 1)).squeeze(-1)
+    under_sqrt = ray_cam_dot ** 2 - (cam_loc.norm(2, 1, keepdim=True) ** 2 - r ** 2)
+
+    # sanity check
+    if (under_sqrt <= 0).sum() > 0:
+        print('BOUNDING SPHERE PROBLEM!')
+        exit()
+
+    sphere_intersections = torch.sqrt(under_sqrt) * torch.Tensor([-1, 1]).cuda().float() - ray_cam_dot
+    sphere_intersections = sphere_intersections.clamp_min(0.0)
+
+    return sphere_intersections
+
+def bilinear_interpolation(xs, ys, dist_map):
+    x1 = np.floor(xs).astype(np.int32)
+    y1 = np.floor(ys).astype(np.int32)
+    x2 = x1 + 1
+    y2 = y1 + 1
+
+    dx = np.expand_dims(np.stack([x2 - xs, xs - x1], axis=1), axis=1)
+    dy = np.expand_dims(np.stack([y2 - ys, ys - y1], axis=1), axis=2)
+    Q = np.stack([
+        dist_map[x1, y1], dist_map[x1, y2], dist_map[x2, y1], dist_map[x2, y2]
+    ], axis=1).reshape(-1, 2, 2)
+    return np.squeeze(dx @ Q @ dy)  # ((x2 - x1) * (y2 - y1)) = 1
+
+def get_index_outside_of_bbox(samples_uniform, bbox_min, bbox_max):
+    samples_uniform_row = samples_uniform[:, 0]
+    samples_uniform_col = samples_uniform[:, 1]
+    index_outside = np.where((samples_uniform_row < bbox_min[0]) | (samples_uniform_row > bbox_max[0]) | (samples_uniform_col < bbox_min[1]) | (samples_uniform_col > bbox_max[1]))[0]
+    return index_outside
+
+
+def weighted_sampling(data, img_size, num_sample, bbox_ratio=0.9):
+    """
+    More sampling within the bounding box
+    """
+    
+    # calculate bounding box
+    mask = data["object_mask"]
+    where = np.asarray(np.where(mask))
+    bbox_min = where.min(axis=1)
+    bbox_max = where.max(axis=1)
+
+    num_sample_bbox = int(num_sample * bbox_ratio)
+    samples_bbox = np.random.rand(num_sample_bbox, 2)
+    samples_bbox = samples_bbox * (bbox_max - bbox_min) + bbox_min
+
+    num_sample_uniform = num_sample - num_sample_bbox
+    samples_uniform = np.random.rand(num_sample_uniform, 2)
+    samples_uniform *= (img_size[0] - 1, img_size[1] - 1)
+
+    # get indices for uniform samples outside of bbox
+    index_outside = get_index_outside_of_bbox(samples_uniform, bbox_min, bbox_max) + num_sample_bbox
+
+    indices = np.concatenate([samples_bbox, samples_uniform], axis=0)
+    output = {}
+    for key, val in data.items():
+        if len(val.shape) == 3:
+            new_val = np.stack([
+                bilinear_interpolation(indices[:, 0], indices[:, 1], val[:, :, i])
+                for i in range(val.shape[2])
+            ], axis=-1)
+        else:
+            new_val = bilinear_interpolation(indices[:, 0], indices[:, 1], val)
+        new_val = new_val.reshape(-1, *val.shape[2:])
+        output[key] = new_val
+    
+    return output, index_outside

code/setup.py

@@ -0,0 +1,34 @@
+# Copyright 2020 The TensorFlow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Set-up script for installing extension modules."""
+from Cython.Build import cythonize
+import numpy
+from setuptools import Extension
+from setuptools import setup
+
+# Get the numpy include directory.
+numpy_include_dir = numpy.get_include()
+
+# mise (efficient mesh extraction)
+mise_module = Extension(
+    "lib.libmise.mise",
+    sources=["lib/libmise/mise.pyx"],
+)
+
+# Gather all extension modules
+ext_modules = [
+    mise_module,
+]
+
+setup(ext_modules=cythonize(ext_modules),)

code/test.py

@@ -0,0 +1,39 @@
+from v2a_model import V2AModel
+from lib.datasets import create_dataset
+import hydra
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import WandbLogger
+import os
+import glob
+
+@hydra.main(config_path="confs", config_name="base")
+def main(opt):
+    pl.seed_everything(42)
+    print("Working dir:", os.getcwd())
+
+    checkpoint_callback = pl.callbacks.ModelCheckpoint(
+        dirpath="checkpoints/",
+        filename="{epoch:04d}-{loss}",
+        save_on_train_epoch_end=True,
+        save_last=True)
+    logger = WandbLogger(project=opt.project_name, name=f"{opt.exp}/{opt.run}")
+
+    trainer = pl.Trainer(
+        gpus=1,
+        accelerator="gpu",
+        callbacks=[checkpoint_callback],
+        max_epochs=8000,
+        check_val_every_n_epoch=50,
+        logger=logger,
+        log_every_n_steps=1,
+        num_sanity_val_steps=0
+    )
+
+    model = V2AModel(opt)
+    checkpoint = sorted(glob.glob("checkpoints/*.ckpt"))[-1]
+    testset = create_dataset(opt.dataset.metainfo, opt.dataset.test)
+
+    trainer.test(model, testset, ckpt_path=checkpoint)
+
+if __name__ == '__main__':
+    main()

code/train.py

@@ -0,0 +1,45 @@
+from v2a_model import V2AModel
+from lib.datasets import create_dataset
+import hydra
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import WandbLogger
+import os
+import glob
+
+@hydra.main(config_path="confs", config_name="base")
+def main(opt):
+    pl.seed_everything(42)
+    print("Working dir:", os.getcwd())
+
+    checkpoint_callback = pl.callbacks.ModelCheckpoint(
+        dirpath="checkpoints/",
+        filename="{epoch:04d}-{loss}",
+        save_on_train_epoch_end=True,
+        save_last=True)
+    logger = WandbLogger(project=opt.project_name, name=f"{opt.exp}/{opt.run}")
+
+    trainer = pl.Trainer(
+        gpus=1,
+        accelerator="gpu",
+        callbacks=[checkpoint_callback],
+        max_epochs=8000,
+        check_val_every_n_epoch=50,
+        logger=logger,
+        log_every_n_steps=1,
+        num_sanity_val_steps=0
+    )
+
+    
+    model = V2AModel(opt)
+    trainset = create_dataset(opt.dataset.metainfo, opt.dataset.train)
+    validset = create_dataset(opt.dataset.metainfo, opt.dataset.valid)
+
+    if opt.model.is_continue == True:
+        checkpoint = sorted(glob.glob("checkpoints/*.ckpt"))[-1]
+        trainer.fit(model, trainset, validset, ckpt_path=checkpoint)
+    else: 
+        trainer.fit(model, trainset, validset)
+
+
+if __name__ == '__main__':
+    main()

code/v2a_model.py

@@ -0,0 +1,311 @@
+import pytorch_lightning as pl
+import torch.optim as optim
+from lib.model.v2a import V2A
+from lib.model.body_model_params import BodyModelParams
+from lib.model.deformer import SMPLDeformer
+import cv2
+import torch
+from lib.model.loss import Loss
+import hydra
+import os
+import numpy as np
+from lib.utils.meshing import generate_mesh
+from kaolin.ops.mesh import index_vertices_by_faces
+import trimesh
+from lib.model.deformer import skinning
+from lib.utils import utils
+class V2AModel(pl.LightningModule):
+    def __init__(self, opt) -> None:
+        super().__init__()
+
+        self.opt = opt
+        num_training_frames = opt.dataset.metainfo.end_frame - opt.dataset.metainfo.start_frame
+        self.betas_path = os.path.join(hydra.utils.to_absolute_path('..'), 'data', opt.dataset.metainfo.data_dir, 'mean_shape.npy')
+        self.gender = opt.dataset.metainfo.gender
+        self.model = V2A(opt.model, self.betas_path, self.gender, num_training_frames)
+        self.start_frame = opt.dataset.metainfo.start_frame
+        self.end_frame = opt.dataset.metainfo.end_frame
+        self.training_modules = ["model"]
+
+        self.training_indices = list(range(self.start_frame, self.end_frame))
+        self.body_model_params = BodyModelParams(num_training_frames, model_type='smpl')
+        self.load_body_model_params()
+        optim_params = self.body_model_params.param_names
+        for param_name in optim_params:
+            self.body_model_params.set_requires_grad(param_name, requires_grad=True)
+        self.training_modules += ['body_model_params']
+        
+        self.loss = Loss(opt.model.loss)
+        
+    def load_body_model_params(self):
+        body_model_params = {param_name: [] for param_name in self.body_model_params.param_names}
+        data_root = os.path.join('../data', self.opt.dataset.metainfo.data_dir)
+        data_root = hydra.utils.to_absolute_path(data_root)
+
+        body_model_params['betas'] = torch.tensor(np.load(os.path.join(data_root, 'mean_shape.npy'))[None], dtype=torch.float32)
+        body_model_params['global_orient'] = torch.tensor(np.load(os.path.join(data_root, 'poses.npy'))[self.training_indices][:, :3], dtype=torch.float32)
+        body_model_params['body_pose'] = torch.tensor(np.load(os.path.join(data_root, 'poses.npy'))[self.training_indices] [:, 3:], dtype=torch.float32)
+        body_model_params['transl'] = torch.tensor(np.load(os.path.join(data_root, 'normalize_trans.npy'))[self.training_indices], dtype=torch.float32)
+
+        for param_name in body_model_params.keys():
+            self.body_model_params.init_parameters(param_name, body_model_params[param_name], requires_grad=False) 
+
+    def configure_optimizers(self):
+        params = [{'params': self.model.parameters(), 'lr':self.opt.model.learning_rate}]
+        params.append({'params': self.body_model_params.parameters(), 'lr':self.opt.model.learning_rate*0.1})
+        self.optimizer = optim.Adam(params, lr=self.opt.model.learning_rate, eps=1e-8)
+        self.scheduler = optim.lr_scheduler.MultiStepLR(
+            self.optimizer, milestones=self.opt.model.sched_milestones, gamma=self.opt.model.sched_factor)
+        return [self.optimizer], [self.scheduler]
+
+    def training_step(self, batch):
+        inputs, targets = batch
+
+        batch_idx = inputs["idx"]
+
+        body_model_params = self.body_model_params(batch_idx)
+        inputs['smpl_pose'] = torch.cat((body_model_params['global_orient'], body_model_params['body_pose']), dim=1)
+        inputs['smpl_shape'] = body_model_params['betas']
+        inputs['smpl_trans'] = body_model_params['transl']
+
+        inputs['current_epoch'] = self.current_epoch
+        model_outputs = self.model(inputs)
+
+        loss_output = self.loss(model_outputs, targets)
+        for k, v in loss_output.items():
+            if k in ["loss"]:
+                self.log(k, v.item(), prog_bar=True, on_step=True)
+            else:
+                self.log(k, v.item(), prog_bar=True, on_step=True)
+        return loss_output["loss"]
+
+    def training_epoch_end(self, outputs) -> None:        
+        # Canonical mesh update every 20 epochs
+        if self.current_epoch != 0 and self.current_epoch % 20 == 0:
+            cond = {'smpl': torch.zeros(1, 69).float().cuda()}
+            mesh_canonical = generate_mesh(lambda x: self.query_oc(x, cond), self.model.smpl_server.verts_c[0], point_batch=10000, res_up=2)
+            self.model.mesh_v_cano = torch.tensor(mesh_canonical.vertices[None], device = self.model.smpl_v_cano.device).float()
+            self.model.mesh_f_cano = torch.tensor(mesh_canonical.faces.astype(np.int64), device=self.model.smpl_v_cano.device)
+            self.model.mesh_face_vertices = index_vertices_by_faces(self.model.mesh_v_cano, self.model.mesh_f_cano)
+        return super().training_epoch_end(outputs)
+
+    def query_oc(self, x, cond):
+        
+        x = x.reshape(-1, 3)
+        mnfld_pred = self.model.implicit_network(x, cond)[:,:,0].reshape(-1,1)
+        return {'sdf':mnfld_pred}
+
+    def query_wc(self, x):
+        
+        x = x.reshape(-1, 3)
+        w = self.model.deformer.query_weights(x)
+    
+        return w
+
+    def query_od(self, x, cond, smpl_tfs, smpl_verts):
+        
+        x = x.reshape(-1, 3)
+        x_c, _ = self.model.deformer.forward(x, smpl_tfs, return_weights=False, inverse=True, smpl_verts=smpl_verts)
+        output = self.model.implicit_network(x_c, cond)[0]
+        sdf = output[:, 0:1]
+        
+        return {'sdf': sdf}
+
+    def get_deformed_mesh_fast_mode(self, verts, smpl_tfs):
+        verts = torch.tensor(verts).cuda().float()
+        weights = self.model.deformer.query_weights(verts)
+        verts_deformed = skinning(verts.unsqueeze(0),  weights, smpl_tfs).data.cpu().numpy()[0]
+        return verts_deformed
+
+    def validation_step(self, batch, *args, **kwargs):
+
+        output = {}
+        inputs, targets = batch
+        inputs['current_epoch'] = self.current_epoch
+        self.model.eval()
+
+        body_model_params = self.body_model_params(inputs['image_id'])
+        inputs['smpl_pose'] = torch.cat((body_model_params['global_orient'], body_model_params['body_pose']), dim=1)
+        inputs['smpl_shape'] = body_model_params['betas']
+        inputs['smpl_trans'] = body_model_params['transl']
+
+        cond = {'smpl': inputs["smpl_pose"][:, 3:]/np.pi}
+        mesh_canonical = generate_mesh(lambda x: self.query_oc(x, cond), self.model.smpl_server.verts_c[0], point_batch=10000, res_up=3)
+        
+        mesh_canonical = trimesh.Trimesh(mesh_canonical.vertices, mesh_canonical.faces)
+        
+        output.update({
+            'canonical_mesh':mesh_canonical
+        })
+
+        split = utils.split_input(inputs, targets["total_pixels"][0], n_pixels=min(targets['pixel_per_batch'], targets["img_size"][0] * targets["img_size"][1]))
+
+        res = []
+        for s in split:
+
+            out = self.model(s)
+
+            for k, v in out.items():
+                try:
+                    out[k] = v.detach()
+                except:
+                    out[k] = v
+
+            res.append({
+                'rgb_values': out['rgb_values'].detach(),
+                'normal_values': out['normal_values'].detach(),
+                'fg_rgb_values': out['fg_rgb_values'].detach(),
+            })
+        batch_size = targets['rgb'].shape[0]
+
+        model_outputs = utils.merge_output(res, targets["total_pixels"][0], batch_size)
+
+        output.update({
+            "rgb_values": model_outputs["rgb_values"].detach().clone(),
+            "normal_values": model_outputs["normal_values"].detach().clone(),
+            "fg_rgb_values": model_outputs["fg_rgb_values"].detach().clone(),
+            **targets,
+        })
+            
+        return output
+
+    def validation_step_end(self, batch_parts):
+        return batch_parts
+
+    def validation_epoch_end(self, outputs) -> None:
+        img_size = outputs[0]["img_size"]
+
+        rgb_pred = torch.cat([output["rgb_values"] for output in outputs], dim=0)
+        rgb_pred = rgb_pred.reshape(*img_size, -1)
+
+        fg_rgb_pred = torch.cat([output["fg_rgb_values"] for output in outputs], dim=0)
+        fg_rgb_pred = fg_rgb_pred.reshape(*img_size, -1)
+
+        normal_pred = torch.cat([output["normal_values"] for output in outputs], dim=0)
+        normal_pred = (normal_pred.reshape(*img_size, -1) + 1) / 2
+
+        rgb_gt = torch.cat([output["rgb"] for output in outputs], dim=1).squeeze(0)
+        rgb_gt = rgb_gt.reshape(*img_size, -1)
+        if 'normal' in outputs[0].keys():
+            normal_gt = torch.cat([output["normal"] for output in outputs], dim=1).squeeze(0)
+            normal_gt = (normal_gt.reshape(*img_size, -1) + 1) / 2
+            normal = torch.cat([normal_gt, normal_pred], dim=0).cpu().numpy()
+        else:
+            normal = torch.cat([normal_pred], dim=0).cpu().numpy()
+
+        rgb = torch.cat([rgb_gt, rgb_pred], dim=0).cpu().numpy()
+        rgb = (rgb * 255).astype(np.uint8)
+
+        fg_rgb = torch.cat([fg_rgb_pred], dim=0).cpu().numpy()
+        fg_rgb = (fg_rgb * 255).astype(np.uint8)
+
+        normal = (normal * 255).astype(np.uint8)
+
+        os.makedirs("rendering", exist_ok=True)
+        os.makedirs("normal", exist_ok=True)
+        os.makedirs('fg_rendering', exist_ok=True)
+
+        canonical_mesh = outputs[0]['canonical_mesh']
+        canonical_mesh.export(f"rendering/{self.current_epoch}.ply")
+
+        cv2.imwrite(f"rendering/{self.current_epoch}.png", rgb[:, :, ::-1])
+        cv2.imwrite(f"normal/{self.current_epoch}.png", normal[:, :, ::-1])
+        cv2.imwrite(f"fg_rendering/{self.current_epoch}.png", fg_rgb[:, :, ::-1])
+    
+    def test_step(self, batch, *args, **kwargs):
+        inputs, targets, pixel_per_batch, total_pixels, idx = batch
+        num_splits = (total_pixels + pixel_per_batch -
+                       1) // pixel_per_batch
+        results = []
+
+        scale, smpl_trans, smpl_pose, smpl_shape = torch.split(inputs["smpl_params"], [1, 3, 72, 10], dim=1)
+
+        body_model_params = self.body_model_params(inputs['idx'])
+        smpl_shape = body_model_params['betas'] if body_model_params['betas'].dim() == 2 else body_model_params['betas'].unsqueeze(0)
+        smpl_trans = body_model_params['transl']
+        smpl_pose = torch.cat((body_model_params['global_orient'], body_model_params['body_pose']), dim=1)
+
+        smpl_outputs = self.model.smpl_server(scale, smpl_trans, smpl_pose, smpl_shape)
+        smpl_tfs = smpl_outputs['smpl_tfs']
+        cond = {'smpl': smpl_pose[:, 3:]/np.pi}
+
+        mesh_canonical = generate_mesh(lambda x: self.query_oc(x, cond), self.model.smpl_server.verts_c[0], point_batch=10000, res_up=4)
+        self.model.deformer = SMPLDeformer(betas=np.load(self.betas_path), gender=self.gender, K=7)
+        verts_deformed = self.get_deformed_mesh_fast_mode(mesh_canonical.vertices, smpl_tfs)
+        mesh_deformed = trimesh.Trimesh(vertices=verts_deformed, faces=mesh_canonical.faces, process=False)
+
+        os.makedirs("test_mask", exist_ok=True)
+        os.makedirs("test_rendering", exist_ok=True)
+        os.makedirs("test_fg_rendering", exist_ok=True)
+        os.makedirs("test_normal", exist_ok=True)
+        os.makedirs("test_mesh", exist_ok=True)
+        
+        mesh_canonical.export(f"test_mesh/{int(idx.cpu().numpy()):04d}_canonical.ply")
+        mesh_deformed.export(f"test_mesh/{int(idx.cpu().numpy()):04d}_deformed.ply")
+        self.model.deformer = SMPLDeformer(betas=np.load(self.betas_path), gender=self.gender)
+        for i in range(num_splits):
+            indices = list(range(i * pixel_per_batch,
+                                min((i + 1) * pixel_per_batch, total_pixels)))
+            batch_inputs = {"uv": inputs["uv"][:, indices],
+                            "intrinsics": inputs['intrinsics'],
+                            "pose": inputs['pose'],
+                            "smpl_params": inputs["smpl_params"],
+                            "smpl_pose": inputs["smpl_params"][:, 4:76],
+                            "smpl_shape": inputs["smpl_params"][:, 76:],
+                            "smpl_trans": inputs["smpl_params"][:, 1:4],
+                            "idx": inputs["idx"] if 'idx' in inputs.keys() else None}
+
+            body_model_params = self.body_model_params(inputs['idx'])
+
+            batch_inputs.update({'smpl_pose': torch.cat((body_model_params['global_orient'], body_model_params['body_pose']), dim=1)})
+            batch_inputs.update({'smpl_shape': body_model_params['betas']})
+            batch_inputs.update({'smpl_trans': body_model_params['transl']})
+
+            batch_targets = {"rgb": targets["rgb"][:, indices].detach().clone() if 'rgb' in targets.keys() else None,
+                             "img_size": targets["img_size"]}
+
+            with torch.no_grad():
+                model_outputs = self.model(batch_inputs)
+            results.append({"rgb_values":model_outputs["rgb_values"].detach().clone(), 
+                            "fg_rgb_values":model_outputs["fg_rgb_values"].detach().clone(),
+                            "normal_values": model_outputs["normal_values"].detach().clone(),
+                            "acc_map": model_outputs["acc_map"].detach().clone(),
+                            **batch_targets})         
+
+        img_size = results[0]["img_size"]
+        rgb_pred = torch.cat([result["rgb_values"] for result in results], dim=0)
+        rgb_pred = rgb_pred.reshape(*img_size, -1)
+
+        fg_rgb_pred = torch.cat([result["fg_rgb_values"] for result in results], dim=0)
+        fg_rgb_pred = fg_rgb_pred.reshape(*img_size, -1)
+
+        normal_pred = torch.cat([result["normal_values"] for result in results], dim=0)
+        normal_pred = (normal_pred.reshape(*img_size, -1) + 1) / 2
+
+        pred_mask = torch.cat([result["acc_map"] for result in results], dim=0)
+        pred_mask = pred_mask.reshape(*img_size, -1)
+
+        if results[0]['rgb'] is not None:
+            rgb_gt = torch.cat([result["rgb"] for result in results], dim=1).squeeze(0)
+            rgb_gt = rgb_gt.reshape(*img_size, -1)
+            rgb = torch.cat([rgb_gt, rgb_pred], dim=0).cpu().numpy()
+        else:
+            rgb = torch.cat([rgb_pred], dim=0).cpu().numpy()
+        if 'normal' in results[0].keys():
+            normal_gt = torch.cat([result["normal"] for result in results], dim=1).squeeze(0)
+            normal_gt = (normal_gt.reshape(*img_size, -1) + 1) / 2
+            normal = torch.cat([normal_gt, normal_pred], dim=0).cpu().numpy()
+        else:
+            normal = torch.cat([normal_pred], dim=0).cpu().numpy()
+        
+        rgb = (rgb * 255).astype(np.uint8)
+
+        fg_rgb = torch.cat([fg_rgb_pred], dim=0).cpu().numpy()
+        fg_rgb = (fg_rgb * 255).astype(np.uint8)
+
+        normal = (normal * 255).astype(np.uint8)
+
+        cv2.imwrite(f"test_mask/{int(idx.cpu().numpy()):04d}.png", pred_mask.cpu().numpy() * 255)
+        cv2.imwrite(f"test_rendering/{int(idx.cpu().numpy()):04d}.png", rgb[:, :, ::-1])
+        cv2.imwrite(f"test_normal/{int(idx.cpu().numpy()):04d}.png", normal[:, :, ::-1])
+        cv2.imwrite(f"test_fg_rendering/{int(idx.cpu().numpy()):04d}.png", fg_rgb[:, :, ::-1])

data/parkinglot/cameras.npz

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

data/parkinglot/cameras_normalize.npz

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

data/parkinglot/checkpoints/epoch=6299-loss=0.01887552998960018.ckpt

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