thanks to One-2-3-45 ❤

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.DS_Store
+*.ipynb
+*.egg-info/

README.md

@@ -0,0 +1,34 @@
+---
+license: mit
+---
+
+
+# One-2-3-45's Inference Model
+
+<div>
+<a style="display:inline-block" href="http://one-2-3-45.com"><img src="https://img.shields.io/badge/Project_Homepage-f9f7f7?logo=data:image/webp;base64,UklGRmIRAABXRUJQVlA4IFYRAABQPwCdASrIAMgAPm00lEekpiolqDvpMIANiWJu3pE7maI+vTDkhN5f7PfmGT7nS6p8nKBr0I+YBzr/ML+2/rG/8j1Sf3/1AP6v/sOsW/bn2GP2W9Zv/zeyf/cf+Z+63tReoB/6uCx2p5E/iUkPwG1FO9t/XgHPL7RH7TzI8EvEA4Mr1T2CP0J6xH+d5J/rz2F/LG9h37l+x9+4xQ3m86D2Te/zeVV/tWyTw7s85XZ0ABD4N2CpzWHt8feKiWkqdTkRjojREWrbUDAKXlYsV7EGU9rWR2gCxVXnstqpNVhwra603swvYRlMyRzKc5nJHEEeLuV8EDee/MpPVIq2DNUcXpCZXvFLHgXBWvZWzSZCFo4iub8df+Yu9q7rw5qemOe2Nt1IIoyjBmXdjCunMherehhPjQIQGiI6PDcriy/zhhwHE7O+0gmpUsYmcDR+ixOfLPY0yjnQosZkIoK1pfttDGtirMbSMqndDVi73JMxcSlNb0MNFtgdAAXNk5Z77wgsPz9RRj9oWO/KRpXn5ra4gUt+mMgSCvFG86zgghSehTRD54z10sNxqnG3/rpKDifOvT4EQU1uA9ZckUZcUt5L5C0+dOdj1I56uLJEAsn0432gHD5wRG7dgSfYusXvhGl2uMaczlXSJ0JfX+Z0e9q7sHywvyEkWC+iJREwvtWi1K+NQAD+/JSRhGP+QTeW9xU73vXKZO+JaR/TAb6vV9dNzIjket6jYZdxK0qCcaf95QeouegLeSQL/9WeH5l2/DE2AKdYjhEYzzIefp7c6cTfM3D3q3kSFxAF/xP/f/3YUFjjOzfzl5xrD3XaWz0TAehn6+ze5pANq6t5GDX8ZOfpIBGUplJj6UZXd76ropLkDdM+d/F2Megl53hry7QvtcUGNlKgjLd7/txvzvkYIPre5sKVvAJzj9DEml706Piekk2NTtBnCMQtQAPO7/Soo3p3QbqLnMIY2PKCq3jFUkeMDAB6uvaHy7e8G/yi+LlFCfYgju+h+ha+jj6NYh6xUx/9TpQoQ1VFrpEw7pCAaQ2NbzVcj/EfBLQUWQBwliZd6FG70L3ATK7AS/cu+Pm/ASndDhIDTx08uveDvY2kW7Mqproq8D4ImWzJ7ZwM8JfrvyN9/wh0Iu00O3UbTDU58dYfWzxI1gDb2Yt6+AyvgjRY/WUM8aikx5MTFi6ZEWeffMc8ruwWeKmfwJtpDxNYhJgSN5gZoOS+XedZmwoYfiuaf9hhPdDtJCM429liA9mZQ2GNfMOPtcLJV/4xTUuWJx4/d43remtcIdsy1GlD79SNNSlfWuSCF8LIopGEcQwZnBVOBmJ7O2bQvBHNAQ6dlz+Pc8zL7MgsN7uff5PDGyXHqV4lCh2Q/xbeZaYVv1agSO4QU5FvEX/6AQpNxQqgtvrXM+XsXTWYJpiO7+ucPvPudldDswT7/ITpp7AdSJ9OjPYr3cMRBVy5sXXkyY8SVv0z//QqGJbxMA3IV81dfN5dUNNlMvl+EBv6Qrnq42ZAEXMEbW/zcmuvIO+539I0BKM+COuGTuEmhulQBlMmdlBNVaII5lFuENjHGpIPNULKARA/PhM9jOvRG2xs8SFCjLZ9ZNLyWznJxq3sZWeGUNXcyJPKDr3WAWCmP8tNyjKQk6aKOD1s/+2MCNQ9b4Zb2XJJLW4kBOS6P10n42Scz8D1K6GTeouELbGApoVNEYjw2jdbNqsNLZiJi6XfSs7Kz5ozupOLJsYOzfLQTjDz7BZOzEYFoB+acySl5Qs3mO84Mw6pIb9MqHhzFqtmrsD4V6uwssP9SUvJTMA4orRWZAMdCA9lMHQi6XBPR9CBidQdudyeYtUj5gWMyzdcSkG8l/2zbzpz8THi23y/8+Ijg5naj6GtYnpJna+6QpI6dlY+E2KF7bOK2ctKwBjClUvAjHXa1162i6DsToHLQE4stmDJdvI1POb9Hj0Mq+azo1wrOfqVFcAS5XNc37IJeYBs/cQYZ08mg2vXWWJYVWz648jTHABHf+LiHsy4WRaVo4oOOSyeampoUYSM9WUJ3iOlTMis5U2DCrGoAiATOAyyuwMcYgTni5FGSpdE5BnoS6ORUiYapPetM/XmsvikTkKNn4z4jhiLFFcU+bH1pZ2DseVK9vCgY5s9ZDjNb9Ky+8fwn9dJtsZ6M7opvXhqde9Ljos6KWQ/8hj3pswa2lLZ7WRc9xaxTjq1sytCxfOd+J+VrsXedNuuYDMwumYIzF1Xsbz1VCURDw6C1twAPizF49s4JfToWsGhgG4wtBE5NAU4KvnGleFGzv54AwBR9qqPVD9bcN7ZmZhphTcAGnR2oOcvT98FmknnVXqJYdHSeP9nWG6A8YUUgDmy7rYLtbxgpid5ysrxvLeyMOeTaQXuNZXn5kQeqDGELMfQ5U2PAk+/DhnbWTsirG5NAP0tgGbSNGM5cv9+trgSk6sXdw1lZOLrfqLGZ8Dt19DWcxmjsbDU30CoSc1alYxX5G+uIHy72tQxjzsot1O4iZeNO34PItzBoeg0Fq+YQZGsHdknwJkAbffRl96xFsDxM6l4g22OaMEHxLMC9uFFE8Ee/xf+krkjv7YCfJrCM3Nw6xfyrhtxN3x2GxSg4YTu2dtjb3zVI/6jYNwgGiaDWoh5I29uQ8ZvES8Ros5jgxDzeKB1tJ3HtDM9SFGNJfQiLiSyYZQLBjCcGbi3+vlythB3k6af+P5rDqah2oPFl29Ngnw/tmpkmRIvri5i55FPeY9J4nXfvWYHTHdoB0oVA2NEk2nropP+T7GXhAxA2NgyGtzHaVU2yxiSju87w8MLIo1eac26wOnbEo/oD6Zcb8vyu0x7ug9iERQ5FlppDnIktT6QC6Kk3qBxovLzOPdEvYQoytaN256n2dmkxAaq78klv6PnU7HiH3e/I9RC27VOP0j6JDW19KvC9/uN9tfOi6WMr0IGKpTsZAUZXm+Ukyk/Rpu9ZPIH5/3CL+yfj3ROts+BWIZNj8lpFHfmYhmN/J0+/lDIGmbRVMbvmif9tqr53fqb8EkFpdMHnK8jc0oIYu2Io5SWOzHc7GMdwt5RB8LR5jUjV6Xv+rR7N4IcTQRphe7WarFsxHmihpNr8sLroqsVxBH+6dOjC5DPhRV6aJB9ZB0NjpLtbjRsEKe1By1huo8rJa+DS73fTUfxWPaJjQsasBOcc6xwuob3OBjTFjUOxfiBbhMDNUFcamlMphrkbmTe2smHz0hrScXZjoHxphV537e8PNenBpI//N58bUOcmV4Lo1H1BLLjNTw1gK+rKFgaU/WOZQ0DZ1kRRqCa86XYnFposmkLgDNooS/yeW/RGfvopRDH40d2TeW8t1+2fDHQcwocSXolq+dxC6JMGsu2rCrhdjzhqd1KPMp5EVGQuCyLc8LfjUhQ8fSs63P9aVDYZBDhO8oWSI2Lbk7cRpKJ38ww9dD0b0OjvucHkJl1zIwyQFqKKEfIN7RPvV8Q1Xxot6Y5f8/UqOCOVZRt+IM1JFcJ4AstPMOXs9hAyZzEs1EY9lv3976/18LNNvL8K7RPNH1uz3qwAajMXLaOTEK7IzCjex3YZQ0LCICPzWVKMNbkSFpmy5ow1A54fK4F45T0apL1FE8dc/Jy6ERymiJ8ZvT+BJHUtbS5oB72w8NeIb0zTuqTzYwMQiKeCI+DlJTd6R3dgbvDETb7XtLT1L5quVxBiyJLxgARoeWU1DY3eWTFJkicFp/UqIFCYgLUhQgGm/1gAxylWf4wZmbQy6RGlY3/pfn5qxqFq8Xmza7Unght3AckydGZ6u6yWcooxZwILsHaklA/Bu2HRlCLzLer57IQWfvHUjJ8pqEoZ/TE0WqZc4SF6CBVC4KGEIqyPnH/+chaIQRfGuKg0rKAAc5tB+7vGl4ck72A+dA9iW0UUwXqD6Y333q9MEdov8hbXuiRkRMv1CEm0h1N8yhxOEe1SLWxlCmvUHcVvhojM6S4ODYr2rxHxOqx63MVVCk6PpQAB2gn4D/9+QHVBBqAxLV8Pggh6aRlEmPuHNEc+b/1Zqh4lvwxUgyMFngOgTAhqZAZqBpRRD41KfU7wEbzruYQhOIxxPMbGIe93LRFgYZLz21sLmS/02JhQ1eY6cSu/iumOWleWujzjCcB9qxDpOBlsugpiveHwunOO9Lu66uGNeKuw6Aqo+zTcdEX+BOlserNtSeyYmhQrwLA+mnEqaAtKv5eAyTC03krdlSEI+++xVMU+kqsGF+6H9yNBQj5aZxmOzd7BejBdBBInEjlj868zR80jlgVKb+yQ7XkdiFIvQl/XvaFpPGqYb4UR70U0jNe/I9UuFggu8M90wyOi7Ihm3t9FZTBPv4zmee4ue5pVKpdsOOZLwBSODTdpUJb6ctU3I9n0KAUBM4+RkNoQzCyb+iXoXl22CL2eQWlOlBi8IG84Y2bMIiLnPs5qeUth9zlniN14oQNTtVJibuIgkylT50ExHyuqz1ra2+wW3QDltErT6yyrKnL8rmkPesI3aPAL880z4U6TWXqcU6hkryL8W5gdI94KYuDTBEim0GM6IAAKf8JZNX3sM/OIB9h3XbFUuNXRocJY9iqQAGdinm3YPLbRBxP5S5EWwlTdIVK5yjUpV+tCN0HXOVf7xj9pnyIMPDz/Znf4zufz+0zonywFQLgAdKiVwC5a6EDC2rmxYC4L82QIO17SKc8NCAJZPTWwwrPGgb0nhQdi3g32QzHUAqE2qhq2jyM7WINI34P28PN5IE50uRx/XFn2a8h2Qgla55PnsIT7KbDBo0Nd4XUkCWINxReQK0/NZEDZrUuZghyZYnnIuIi0pTpecJWliTLvfxyiRkIsb9t2mT6VzM8H2HN8nq0rF7BC27r0JoLl/5YgZQZmw763cQ625wkmPOX0vr1M35fZYv06zKm1ux/L+W6O3ju3VdFudKgEgRIeT+bIOQKoKaT+knRugmPDGt1JAt6bKTT2bvIYnf5OvZs9id5x+qy5UeotL3uxYiBj7SyGxTCHdovbak2BG5hGmuVWxRojEJS9IEqUKwy133zg24keiFy0bXsG125D1XgQ/uI3IM8dijJ4N6jHObWneJl3zHvKb+cX97XFAv5VV5ySEfm0Iglkir9QaTOXP9SuTeCao7Q3fULO6Jcp+sOU6x+jCjlsmHiF7oDCAb/RYITB9oIcAGGQJkT6ccTNJPlyPiNf4/nefe5V4j5RWLTDm7Wb/kt426CIGzE2ekBjrvExlch914MGYjMJcBllUj/LTfXTYLSqPwPzU/xBVSUOR9o7wFnFBZTaI0ApgY11rRsoEgTu9yRgBxw0h71O/RjpN5Ku/U/er87C9/jHzucfXpRDcP1JxOOxoJziSE01YxnPjmyDigmCcur63bY/xXdZeNQNprWvE3mAIP14fFkdJ4+0vwkAP+BXokPPQBkZuEWFAUEz1H/YQf4Q9bCQZXl/WSpUpG/TjBo8EpZLTJ2Jwa1G3H2hVIUlifUnV/SvKDYbpUvl6mKuwdgglJxkJOXjtf84FjvjeHUOzf8ZhLw3PH53rUrDz0INySaGJ/n4a/iuvLMaL146Ii98kND4sM0nElTIxnJe+LJF/8aimynAAiTshwnKc7MHqCtuDaUFfEQCGw0tmys5yKZM5zzawr6LW7MdQs9XUDiyTrX+YcI0uPZZ43oMnO737u5Tmc/sAeKCNGIWt8kw87EMQ+BP5NMrf8b8wDvSD2XVEZu7xqwNCizeSYQJGVJVSAdJ27XwXrFfHtdHHrlojW+3BFzE5rOzDsUsA00zYHxt+e9zo9Yn0sImcxGhbDFBGD892Rgz9G+eor3huRF8h4p1qYpjTe/ykVkhWyvHRjNNevOV7Gk1jhgiwOajzrXwNsIJNUvAQQB017GRYgey7MAEeBoAx5RuxYU+oMH6DNk5eYcrJDxo48XGbO4QhCMRgAA"></a>
+<a style="display:inline-block; margin-left: .5em" href="https://arxiv.org/abs/2306.16928"><img src="https://img.shields.io/badge/2306.16928-f9f7f7?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAADcAAABMCAYAAADJPi9EAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAuIwAALiMBeKU/dgAAABl0RVh0U29mdHdhcmUAd3d3Lmlua3NjYXBlLm9yZ5vuPBoAAAa2SURBVHja3Zt7bBRFGMAXUCDGF4rY7m7bAwuhlggKStFgLBgFEkCIIRJEEoOBYHwRFYKilUgEReVNJEGCJJpehHI3M9vZvd3bUP1DjNhEIRQQsQgSHiJgQZ5dv7krWEvvdmZ7d7vHJN+ft/f99pv5XvOtJMFCqvoCUpTdIEeRLC+L9Ox5i3Q9LACaCeK0kXoSChVcD3C/tQPHpAEsquQ73IkUcEz2kcLCknyGW5MGjkljRFVL8xJOKyi4CwCOuQAeAkfTP1+tNxLkogvgEbDgffkJqKqvuMA5ifOpqg/5qWecRstNg7xoUTI1Fovdxg8oy2s5AP8CGeYHmGngeZaOL4I4LXLcpHg4149/GDz4xqgsb+UAbMKKUpkrqHA43MUyyJpWUK0EHeG2YKRXr7tB+QMcgGewLD+ebTDbtrtbBt7UPlhS4rV4IvcDI7J8P1OeA/AcAI7LHljN7aB8XTowJmZt9EFRD/o0SDMH4HlwMhMyDWZZSAHFf3YDs3RS49WDLuaAY3IJq+qzmQKLxXAZKN7oDoYbdV3v5elPqiSpMyiOuAEVZVqHXb1OhloUH+MA+ztO0cAO/RkrfyBE7OAEbAZvO8vzVtTRWFD6DAfY5biBM3PWiaL0a4lvXICwnV8WjmE6ntYmhqX2jjp5LbMZjCw/wbYeN6CizOa2GMVzQOlmHjB4Ceuyk6LJ8huccEmR5Xddg7OOV/NAtchW+E3XbOag60QA4Qwuarca0bRuEJyr+cFQwzcY98huxhAKdQelt4kAQpj4qJ3gvFXAYn+aJumXk1yPlpQUgtIHhbYoFMUstNRRWgjnpl4A7IKlayNymqFHFaWCpV9CFry3LGxR1CgA5kB5M8OX2goApwpaz6mdOMGxtAgXWJySxb4WuQD4qTDgU+N5AAnzpr7ChSWpCyisiQJqY0Y7FtmSKpbV23b45kC0KHBxcQ9QeI8w4KgnHRPVtIU7rOtbioLVg5Hl/qDwSVFAMqLSMSObroCdZYlzIJtMRFVHCaRo/wFWPgaAXzdbBpkc2A4aKzCNd97+URQuESYGDDhIVfWOQIKZJu4D2+oXlgDTV1865gUQZDts756BArMNMoR1oa46BYqbyPixZz1ZUFV3sgwoGBajuBKATl3btIn8QYYMuezRgrsiRUWyr2BxA40EkPMpA/Hm6gbUu7fjEXA3azP6AsbKD9bxdUuhjM9W7fII52BF+daRpE4+WA3P501+jbfmHvQKyFqMuXf7Ot4mkN2fr50y+bRH61X7AXdUpHSxaPQ4GVbR5AGw3g+434XgQGKfr72I+vQRhfsu92dOx7WicInzt3CBg1RVpMm0NveWo2SqFzgmdNZMbriILD+S+zoueWf2vSdAipzacWN5nMl6XxNlUHa/J8DoJodUDE0HR8Ll5V0lPxcrLEHZPV4AzS83OLis7FowVa3RSku7BSNxJqQAlN3hBTC2apmDSkpaw22wJemGQFUG7J4MlP3JC6A+f96V7vRyX9It3nzT/GrjIU8edM7rMSnIi10f476lzbE1K7yEiEuWro0OJBguLCwDuFOJc1Na6sRWL/cCeMIwUN9ggSVbe3v/5/EgzTKWLvEAiBrYRUkgwNI2ZaFQNT75UDxEUEx97zYnzpmiLEmbaYCbNxYtFAb0/Z4AztgUrhyxuNgxPnhfHFDHz/vTgFWUQZxTRkkJhQ6YNdVUEPAfO6ZV5BRss6LcCVb7VaAma9giy0XJZBt9IQh42NY0NSdgbLIPlLUF6rEdrdt0CUCK1wsCbkcI3ZSLc7ZSwGLbmJXbPsNxnE5xilYKAobZ77LpGZ8TAIun+/iCKQoF71IxQDI3K2CCd+ARNvXg9sykBcnHAoCZG4u66hlDoQLe6QV4CRtFSxZQ+D0BwNO2jgdkzoGoah1nj3FVlSR19taTSYxI8QLut23U8dsgzqHulJNCQpcqBnpTALCuQ6NSYLHpmR5i42gZzuIdcrMMvMJbQlxe3jXxyZnLACl7ARm/FjPIDOY8ODtpM71sxwfcZpvBeUzKWmfNINM5AS+wO0Khh7dMqKccu4+qatarZjYAwDlgetzStHtEt+XedsBOQtU9XMrRgjg4KTnc5nr+dmqadit/4C4uLm8DuA9koJTj1TL7fI5nDL+qqoo/FLGAzL7dYT17PzvAcQONYSUQRxW/QMrHZVIyik0ZuQA2mzp+Ji8BW4YM3Mbzm9inaHkJCGfrUZZjujiYailfFwA8DHIy3acwUj4v9vUVa+SmgNsl5fuyDTKovW9/IAmfLV0Pi2UncA515kjYdrwC9i9rpuHiq3JwtAAAAABJRU5ErkJggg=="></a>
+<a style="display:inline-block; margin-left: .5em" href='https://github.com/One-2-3-45/One-2-3-45'><img src='https://img.shields.io/github/stars/One-2-3-45/One-2-3-45?style=social' /></a>
+</div>
+
+This inference model supports the demo for [One-2-3-45](http://One-2-3-45.com). 
+
+Try out our 🤗 Hugging Face Demo:
+<a target="_blank" href="https://huggingface.co/spaces/One-2-3-45/One-2-3-45">
+  <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/open-in-hf-spaces-sm.svg" alt="Open in HuggingFace"/>
+</a>
+
+Please refer to our [GitHub repo](https://github.com/One-2-3-45/One-2-3-45) for full code release and local deployment.
+
+## Citation
+
+```bibtex
+@misc{liu2023one2345,
+      title={One-2-3-45: Any Single Image to 3D Mesh in 45 Seconds without Per-Shape Optimization}, 
+      author={Minghua Liu and Chao Xu and Haian Jin and Linghao Chen and Mukund Varma T and Zexiang Xu and Hao Su},
+      year={2023},
+      eprint={2306.16928},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

SparseNeuS_demo_v1/confs/one2345_lod0_val_demo.conf

@@ -0,0 +1,135 @@
+# - for the lod1 geometry network, using adaptive cost for sparse cost regularization network
+#- for lod1 rendering network, using depth-adaptive render
+
+general {
+
+  base_exp_dir = exp/lod0 # !!! where you store the results and checkpoints to be used
+  recording = [
+    ./,
+    ./data
+    ./ops
+    ./models
+    ./loss
+  ]
+}
+
+dataset {
+  trainpath = ../
+  valpath = ../ # !!! where you store the validation data
+  testpath = ../
+
+  imgScale_train = 1.0
+  imgScale_test = 1.0
+  nviews = 5
+  clean_image = True
+  importance_sample = True
+  test_ref_views = [23]
+
+  # test dataset
+  test_n_views = 2
+  test_img_wh = [256, 256]
+  test_clip_wh = [0, 0]
+  test_scan_id = scan110
+  train_img_idx = [49, 50, 52, 53, 54, 56, 58] #[21, 22, 23, 24, 25]  #
+  test_img_idx = [51, 55, 57]  #[32, 33, 34]  #
+
+  test_dir_comment = train
+}
+
+train {
+  learning_rate = 2e-4
+  learning_rate_milestone = [100000, 150000, 200000]
+  learning_rate_factor = 0.5
+  end_iter = 200000
+  save_freq = 5000
+  val_freq = 1
+  val_mesh_freq = 1
+  report_freq = 100
+
+  N_rays = 512
+
+  validate_resolution_level = 4
+  anneal_start = 0
+  anneal_end = 25000
+  anneal_start_lod1 = 0
+  anneal_end_lod1 = 15000
+
+  use_white_bkgd = True
+
+  # Loss
+  # ! for training the lod1 network, don't use this regularization in first 10k steps; then use the regularization
+  sdf_igr_weight = 0.1
+  sdf_sparse_weight = 0.02  # 0.002 for lod1 network;  0.02 for lod0 network
+  sdf_decay_param = 100   # cannot be too large, which decide the tsdf range
+  fg_bg_weight = 0.01  # first 0.01
+  bg_ratio = 0.3
+
+  if_fix_lod0_networks = False
+}
+
+model {
+  num_lods = 1
+
+  sdf_network_lod0 {
+    lod = 0,
+    ch_in = 56,  # the channel num of fused pyramid features
+    voxel_size = 0.02105263,  # 0.02083333, should be 2/95
+    vol_dims = [96, 96, 96],
+    hidden_dim = 128,
+    cost_type = variance_mean
+    d_pyramid_feature_compress = 16,
+    regnet_d_out = 16,
+    num_sdf_layers = 4,
+    # position embedding
+    multires = 6
+  }
+
+
+  sdf_network_lod1 {
+    lod = 1,
+    ch_in = 56,  # the channel num of fused pyramid features
+    voxel_size = 0.0104712, #0.01041667, should be 2/191
+    vol_dims = [192, 192, 192],
+    hidden_dim = 128,
+    cost_type = variance_mean
+    d_pyramid_feature_compress = 8,
+    regnet_d_out = 16,
+    num_sdf_layers = 4,
+
+    # position embedding
+    multires = 6
+  }
+
+
+  variance_network {
+    init_val = 0.2
+  }
+
+  variance_network_lod1 {
+    init_val = 0.2
+  }
+
+  rendering_network {
+    in_geometry_feat_ch = 16
+    in_rendering_feat_ch = 56
+    anti_alias_pooling = True
+  }
+
+  rendering_network_lod1 {
+    in_geometry_feat_ch = 16 # default 8
+    in_rendering_feat_ch = 56
+    anti_alias_pooling = True
+
+  }
+
+
+  trainer {
+    n_samples_lod0 = 64
+    n_importance_lod0 = 64
+    n_samples_lod1 = 64
+    n_importance_lod1 = 64
+    n_outside = 0  # 128 if render_outside_uniform_sampling
+    perturb = 1.0
+    alpha_type = div
+  }
+}

SparseNeuS_demo_v1/data/blender_general_narrow_all_eval_new_data.py

@@ -0,0 +1,394 @@
+from torch.utils.data import Dataset
+import os
+import json
+import numpy as np
+import cv2
+from PIL import Image
+import torch
+from torchvision import transforms as T
+from data.scene import get_boundingbox
+
+from models.rays import gen_rays_from_single_image, gen_random_rays_from_single_image
+from kornia import create_meshgrid
+
+def get_ray_directions(H, W, focal, center=None):
+    """
+    Get ray directions for all pixels in camera coordinate.
+    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
+               ray-tracing-generating-camera-rays/standard-coordinate-systems
+    Inputs:
+        H, W, focal: image height, width and focal length
+    Outputs:
+        directions: (H, W, 3), the direction of the rays in camera coordinate
+    """
+    grid = create_meshgrid(H, W, normalized_coordinates=False)[0] + 0.5 # 1xHxWx2
+
+    i, j = grid.unbind(-1)
+    # the direction here is without +0.5 pixel centering as calibration is not so accurate
+    # see https://github.com/bmild/nerf/issues/24
+    cent = center if center is not None else [W / 2, H / 2]
+    directions = torch.stack([(i - cent[0]) / focal[0], (j - cent[1]) / focal[1], torch.ones_like(i)], -1)  # (H, W, 3)
+
+    return directions
+
+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()  # ? why need transpose here
+    pose[:3, 3] = (t[:3] / t[3])[:, 0]
+
+    return intrinsics, pose  # ! return cam2world matrix here
+
+
+# ! load one ref-image with multiple src-images in camera coordinate system
+class BlenderPerView(Dataset):
+    def __init__(self, root_dir, split, n_views=3, img_wh=(256, 256), downSample=1.0,
+                 split_filepath=None, pair_filepath=None,
+                 N_rays=512,
+                 vol_dims=[128, 128, 128], batch_size=1,
+                 clean_image=False, importance_sample=False, test_ref_views=[],
+                 specific_dataset_name = 'GSO'
+                 ):
+
+        # print("root_dir: ", root_dir)
+        self.root_dir = root_dir
+        self.split = split
+
+        self.specific_dataset_name = specific_dataset_name
+        self.n_views = n_views
+        self.N_rays = N_rays
+        self.batch_size = batch_size  # - used for construct new metas for gru fusion training
+
+        self.clean_image = clean_image
+        self.importance_sample = importance_sample
+        self.test_ref_views = test_ref_views  # used for testing
+        self.scale_factor = 1.0
+        self.scale_mat = np.float32(np.diag([1, 1, 1, 1.0]))
+        assert self.split == 'val' or 'export_mesh', 'only support val or export_mesh'
+        # find all subfolders
+        main_folder = os.path.join(root_dir, self.specific_dataset_name)
+        self.shape_list = [""] # os.listdir(main_folder) # MODIFIED
+        self.shape_list.sort()
+
+        # self.shape_list = ['barrel_render']
+        # self.shape_list = ["barrel", "bag", "mailbox", "shoe", "chair", "car", "dog", "teddy"] # TO BE DELETED
+
+
+        self.lvis_paths = []
+        for shape_name in self.shape_list:
+            self.lvis_paths.append(os.path.join(main_folder, shape_name))
+
+        if img_wh is not None:
+            assert img_wh[0] % 32 == 0 and img_wh[1] % 32 == 0, \
+                'img_wh must both be multiples of 32!'
+
+
+        # * bounding box for rendering
+        self.bbox_min = np.array([-1.0, -1.0, -1.0])
+        self.bbox_max = np.array([1.0, 1.0, 1.0])
+
+        # - used for cost volume regularization
+        self.voxel_dims = torch.tensor(vol_dims, dtype=torch.float32)
+        self.partial_vol_origin = torch.tensor([-1., -1., -1.], dtype=torch.float32)
+        
+
+    def define_transforms(self):
+        self.transform = T.Compose([T.ToTensor()])
+
+
+
+    def load_cam_info(self):
+        for vid, img_id in enumerate(self.img_ids):
+            intrinsic, extrinsic, near_far = self.intrinsic, np.linalg.inv(self.c2ws[vid]), self.near_far
+            self.all_intrinsics.append(intrinsic)
+            self.all_extrinsics.append(extrinsic)
+            self.all_near_fars.append(near_far)
+
+    def read_mask(self, filename):
+        mask_h = cv2.imread(filename, 0)
+        mask_h = cv2.resize(mask_h, None, fx=self.downSample, fy=self.downSample,
+                            interpolation=cv2.INTER_NEAREST)
+        mask = cv2.resize(mask_h, None, fx=0.25, fy=0.25,
+                          interpolation=cv2.INTER_NEAREST)
+
+        mask[mask > 0] = 1  # the masks stored in png are not binary
+        mask_h[mask_h > 0] = 1
+
+        return mask, mask_h
+
+    def cal_scale_mat(self, img_hw, intrinsics, extrinsics, near_fars, factor=1.):
+
+        center, radius, bounds = get_boundingbox(img_hw, intrinsics, extrinsics, near_fars)
+
+        radius = radius * factor
+        scale_mat = np.diag([radius, radius, radius, 1.0])
+        scale_mat[:3, 3] = center.cpu().numpy()
+        scale_mat = scale_mat.astype(np.float32)
+
+        return scale_mat, 1. / radius.cpu().numpy()
+
+    def __len__(self):
+        # return 8*len(self.lvis_paths)
+        return len(self.lvis_paths)
+
+    def __getitem__(self, idx):
+        sample = {}
+        idx = idx * 8 # to be deleted
+        origin_idx = idx
+        imgs, depths_h, masks_h = [], [], []  # full size (256, 256)
+        intrinsics, w2cs, c2ws, near_fars = [], [], [], []  # record proj-mats between views
+
+        folder_path = self.lvis_paths[idx//8]
+        idx = idx % 8 # [0, 7]
+
+        # last subdir name
+        shape_name = os.path.split(folder_path)[-1]
+
+        pose_json_path = os.path.join(folder_path, "pose.json")
+        with open(pose_json_path, 'r') as f:
+            meta = json.load(f)
+        
+        self.img_ids = list(meta["c2ws"].keys()) # e.g. "view_0", "view_7", "view_0_2_10"
+        self.img_wh = (256, 256)
+        self.input_poses = np.array(list(meta["c2ws"].values()))
+        intrinsic = np.eye(4)
+        intrinsic[:3, :3] = np.array(meta["intrinsics"])
+        self.intrinsic = intrinsic
+        self.near_far = np.array(meta["near_far"])
+        self.near_far[1] = 1.8
+        self.define_transforms()
+        self.blender2opencv = np.array(
+            [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]]
+        )
+        
+        self.c2ws = []
+        self.w2cs = []
+        self.near_fars = []
+        for image_idx, img_id in enumerate(self.img_ids):
+            pose = self.input_poses[image_idx]
+            c2w = pose @ self.blender2opencv
+            self.c2ws.append(c2w)
+            self.w2cs.append(np.linalg.inv(c2w))
+            self.near_fars.append(self.near_far)
+        self.c2ws = np.stack(self.c2ws, axis=0)
+        self.w2cs = np.stack(self.w2cs, axis=0)
+
+
+        self.all_intrinsics = []  # the cam info of the whole scene
+        self.all_extrinsics = []
+        self.all_near_fars = []
+        self.load_cam_info() 
+
+
+        # target view
+        c2w = self.c2ws[idx]
+        w2c = np.linalg.inv(c2w)
+        w2c_ref = w2c
+        w2c_ref_inv = np.linalg.inv(w2c_ref)
+
+        w2cs.append(w2c @ w2c_ref_inv)
+        c2ws.append(np.linalg.inv(w2c @ w2c_ref_inv))
+
+        img_filename = os.path.join(folder_path, 'stage1_8', f'{self.img_ids[idx]}')
+
+        img = Image.open(img_filename)
+        img = self.transform(img)  # (4, h, w)
+        
+
+        if img.shape[0] == 4:
+            img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+        imgs += [img]
+
+
+        depth_h = torch.ones((img.shape[1], img.shape[2]), dtype=torch.float32)
+        depth_h = depth_h.fill_(-1.0)
+        mask_h = torch.ones((img.shape[1], img.shape[2]), dtype=torch.int32)
+
+
+        depths_h.append(depth_h)
+        masks_h.append(mask_h)
+        
+        intrinsic = self.intrinsic
+        intrinsics.append(intrinsic)
+    
+
+        near_fars.append(self.near_fars[idx])
+        image_perm = 0  # only supervised on reference view
+
+        mask_dilated = None
+
+
+        src_views = range(8, 8 + 8 * 4)
+
+        for vid in src_views:
+
+            img_filename = os.path.join(folder_path, 'stage2_8', f'{self.img_ids[vid]}')
+            img = Image.open(img_filename)
+            img_wh = self.img_wh
+
+            img = self.transform(img)
+            if img.shape[0] == 4:
+                img = img[:3] * img[-1:] + (1 - img[-1:])  # blend A to RGB
+
+            imgs += [img]
+            depth_h = np.ones(img.shape[1:], dtype=np.float32)
+            depths_h.append(depth_h)
+            masks_h.append(np.ones(img.shape[1:], dtype=np.int32))
+
+            near_fars.append(self.all_near_fars[vid])
+            intrinsics.append(self.all_intrinsics[vid])
+
+            w2cs.append(self.all_extrinsics[vid] @ w2c_ref_inv)
+
+        
+        # ! estimate scale_mat
+        scale_mat, scale_factor = self.cal_scale_mat(
+                                                     img_hw=[img_wh[1], img_wh[0]],
+                                                     intrinsics=intrinsics, extrinsics=w2cs,
+                                                     near_fars=near_fars, factor=1.1
+                                                     )
+
+
+        new_near_fars = []
+        new_w2cs = []
+        new_c2ws = []
+        new_affine_mats = []
+        new_depths_h = []
+        for intrinsic, extrinsic, near_far, depth in zip(intrinsics, w2cs, near_fars, depths_h):
+
+            P = intrinsic @ extrinsic @ scale_mat
+            P = P[:3, :4]
+            # - should use load_K_Rt_from_P() to obtain c2w
+            c2w = load_K_Rt_from_P(None, P)[1]
+            w2c = np.linalg.inv(c2w)
+            new_w2cs.append(w2c)
+            new_c2ws.append(c2w)
+            affine_mat = np.eye(4)
+            affine_mat[:3, :4] = intrinsic[:3, :3] @ w2c[:3, :4]
+            new_affine_mats.append(affine_mat)
+
+            camera_o = c2w[:3, 3]
+            dist = np.sqrt(np.sum(camera_o ** 2))
+            near = dist - 1
+            far = dist + 1
+
+            new_near_fars.append([0.95 * near, 1.05 * far])
+            new_depths_h.append(depth * scale_factor)
+
+        imgs = torch.stack(imgs).float()
+        depths_h = np.stack(new_depths_h)
+        masks_h = np.stack(masks_h)
+
+        affine_mats = np.stack(new_affine_mats)
+        intrinsics, w2cs, c2ws, near_fars = np.stack(intrinsics), np.stack(new_w2cs), np.stack(new_c2ws), np.stack(
+            new_near_fars)
+        
+        if self.split == 'train':
+            start_idx = 0
+        else:
+            start_idx = 1
+
+
+
+        target_w2cs = []
+        target_intrinsics = []
+        new_target_w2cs = []
+        for i_idx in range(8):
+            target_w2cs.append(self.all_extrinsics[i_idx] @ w2c_ref_inv)
+            target_intrinsics.append(self.all_intrinsics[i_idx])
+
+        for intrinsic, extrinsic in zip(target_intrinsics, target_w2cs):
+
+            P = intrinsic @ extrinsic @ scale_mat
+            P = P[:3, :4]
+            # - should use load_K_Rt_from_P() to obtain c2w
+            c2w = load_K_Rt_from_P(None, P)[1]
+            w2c = np.linalg.inv(c2w)
+            new_target_w2cs.append(w2c)
+        target_w2cs = np.stack(new_target_w2cs)
+
+
+
+        view_ids = [idx] + list(src_views)
+        sample['origin_idx'] = origin_idx
+        sample['images'] = imgs  # (V, 3, H, W)
+        sample['depths_h'] = torch.from_numpy(depths_h.astype(np.float32))  # (V, H, W)
+        sample['masks_h'] = torch.from_numpy(masks_h.astype(np.float32))  # (V, H, W)
+        sample['w2cs'] = torch.from_numpy(w2cs.astype(np.float32))  # (V, 4, 4)
+        sample['c2ws'] = torch.from_numpy(c2ws.astype(np.float32))  # (V, 4, 4)
+        sample['target_candidate_w2cs'] = torch.from_numpy(target_w2cs.astype(np.float32))  # (8, 4, 4)
+        sample['near_fars'] = torch.from_numpy(near_fars.astype(np.float32))  # (V, 2)
+        sample['intrinsics'] = torch.from_numpy(intrinsics.astype(np.float32))[:, :3, :3]  # (V, 3, 3)
+        sample['view_ids'] = torch.from_numpy(np.array(view_ids))
+        sample['affine_mats'] = torch.from_numpy(affine_mats.astype(np.float32))  # ! in world space
+
+        sample['scan'] = shape_name
+
+        sample['scale_factor'] = torch.tensor(scale_factor)
+        sample['img_wh'] = torch.from_numpy(np.array(img_wh))
+        sample['render_img_idx'] = torch.tensor(image_perm)
+        sample['partial_vol_origin'] = self.partial_vol_origin
+        sample['meta'] = str(self.specific_dataset_name) + '_' + str(shape_name) + "_refview" + str(view_ids[0])
+        # print("meta: ", sample['meta'])
+
+        # - image to render
+        sample['query_image'] = sample['images'][0]
+        sample['query_c2w'] = sample['c2ws'][0]
+        sample['query_w2c'] = sample['w2cs'][0]
+        sample['query_intrinsic'] = sample['intrinsics'][0]
+        sample['query_depth'] = sample['depths_h'][0]
+        sample['query_mask'] = sample['masks_h'][0]
+        sample['query_near_far'] = sample['near_fars'][0]
+
+        sample['images'] = sample['images'][start_idx:]  # (V, 3, H, W)
+        sample['depths_h'] = sample['depths_h'][start_idx:]  # (V, H, W)
+        sample['masks_h'] = sample['masks_h'][start_idx:]  # (V, H, W)
+        sample['w2cs'] = sample['w2cs'][start_idx:]  # (V, 4, 4)
+        sample['c2ws'] = sample['c2ws'][start_idx:]  # (V, 4, 4)
+        sample['intrinsics'] = sample['intrinsics'][start_idx:]  # (V, 3, 3)
+        sample['view_ids'] = sample['view_ids'][start_idx:]
+        sample['affine_mats'] = sample['affine_mats'][start_idx:]  # ! in world space
+
+        sample['scale_mat'] = torch.from_numpy(scale_mat)
+        sample['trans_mat'] = torch.from_numpy(w2c_ref_inv)
+
+        # - generate rays
+        if ('val' in self.split) or ('test' in self.split):
+            sample_rays = gen_rays_from_single_image(
+                img_wh[1], img_wh[0],
+                sample['query_image'],
+                sample['query_intrinsic'],
+                sample['query_c2w'],
+                depth=sample['query_depth'],
+                mask=sample['query_mask'] if self.clean_image else None)
+        else:
+            sample_rays = gen_random_rays_from_single_image(
+                img_wh[1], img_wh[0],
+                self.N_rays,
+                sample['query_image'],
+                sample['query_intrinsic'],
+                sample['query_c2w'],
+                depth=sample['query_depth'],
+                mask=sample['query_mask'] if self.clean_image else None,
+                dilated_mask=mask_dilated,
+                importance_sample=self.importance_sample)
+
+
+        sample['rays'] = sample_rays
+
+        return sample

SparseNeuS_demo_v1/data/scene.py

@@ -0,0 +1,101 @@
+import numpy as np
+import torch
+
+
+def rigid_transform(xyz, transform):
+    """Applies a rigid transform (c2w) to an (N, 3) pointcloud.
+    """
+    device = xyz.device
+    xyz_h = torch.cat([xyz, torch.ones((len(xyz), 1)).to(device)], dim=1)  # (N, 4)
+    xyz_t_h = (transform @ xyz_h.T).T  # * checked: the same with the below
+
+    return xyz_t_h[:, :3]
+
+
+def get_view_frustum(min_depth, max_depth, size, cam_intr, c2w):
+    """Get corners of 3D camera view frustum of depth image
+    """
+    device = cam_intr.device
+    im_h, im_w = size
+    im_h = int(im_h)
+    im_w = int(im_w)
+    view_frust_pts = torch.stack([
+        (torch.tensor([0, 0, im_w, im_w, 0, 0, im_w, im_w]).to(device) - cam_intr[0, 2]) * torch.tensor(
+            [min_depth, min_depth, min_depth, min_depth, max_depth, max_depth, max_depth, max_depth]).to(device) /
+        cam_intr[0, 0],
+        (torch.tensor([0, im_h, 0, im_h, 0, im_h, 0, im_h]).to(device) - cam_intr[1, 2]) * torch.tensor(
+            [min_depth, min_depth, min_depth, min_depth, max_depth, max_depth, max_depth, max_depth]).to(device) /
+        cam_intr[1, 1],
+        torch.tensor([min_depth, min_depth, min_depth, min_depth, max_depth, max_depth, max_depth, max_depth]).to(
+            device)
+    ])
+    view_frust_pts = view_frust_pts.type(torch.float32)
+    c2w = c2w.type(torch.float32)
+    view_frust_pts = rigid_transform(view_frust_pts.T, c2w).T
+    return view_frust_pts
+
+
+def set_pixel_coords(h, w):
+    i_range = torch.arange(0, h).view(1, h, 1).expand(1, h, w).type(torch.float32)  # [1, H, W]
+    j_range = torch.arange(0, w).view(1, 1, w).expand(1, h, w).type(torch.float32)  # [1, H, W]
+    ones = torch.ones(1, h, w).type(torch.float32)
+
+    pixel_coords = torch.stack((j_range, i_range, ones), dim=1)  # [1, 3, H, W]
+
+    return pixel_coords
+
+
+def get_boundingbox(img_hw, intrinsics, extrinsics, near_fars):
+    """
+    # get the minimum bounding box of all visual hulls
+    :param img_hw:
+    :param intrinsics:
+    :param extrinsics:
+    :param near_fars:
+    :return:
+    """
+
+    bnds = torch.zeros((3, 2))
+    bnds[:, 0] = np.inf
+    bnds[:, 1] = -np.inf
+
+    if isinstance(intrinsics, list):
+        num = len(intrinsics)
+    else:
+        num = intrinsics.shape[0]
+    # print("num: ", num)
+    view_frust_pts_list = []
+    for i in range(num):
+        if not isinstance(intrinsics[i], torch.Tensor):
+            cam_intr = torch.tensor(intrinsics[i])
+            w2c = torch.tensor(extrinsics[i])
+            c2w = torch.inverse(w2c)
+        else:
+            cam_intr = intrinsics[i]
+            w2c = extrinsics[i]
+            c2w = torch.inverse(w2c)
+        min_depth, max_depth = near_fars[i][0], near_fars[i][1]
+        # todo: check the coresponding points are matched
+
+        view_frust_pts = get_view_frustum(min_depth, max_depth, img_hw, cam_intr, c2w)
+        bnds[:, 0] = torch.min(bnds[:, 0], torch.min(view_frust_pts, dim=1)[0])
+        bnds[:, 1] = torch.max(bnds[:, 1], torch.max(view_frust_pts, dim=1)[0])
+        view_frust_pts_list.append(view_frust_pts)
+    all_view_frust_pts = torch.cat(view_frust_pts_list, dim=1)
+
+    # print("all_view_frust_pts: ", all_view_frust_pts.shape)
+    # distance = torch.norm(all_view_frust_pts, dim=0)
+    # print("distance: ", distance)
+
+    # print("all_view_frust_pts_z: ", all_view_frust_pts[2, :])
+
+    center = torch.tensor(((bnds[0, 1] + bnds[0, 0]) / 2, (bnds[1, 1] + bnds[1, 0]) / 2,
+                           (bnds[2, 1] + bnds[2, 0]) / 2))
+
+    lengths = bnds[:, 1] - bnds[:, 0]
+
+    max_length, _ = torch.max(lengths, dim=0)
+    radius = max_length / 2
+
+    # print("radius: ", radius)
+    return center, radius, bnds

SparseNeuS_demo_v1/exp/lod0/.gitignore

@@ -0,0 +1 @@
+checkpoints_*/

SparseNeuS_demo_v1/exp/lod0/checkpoints/ckpt_215000.pth

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

SparseNeuS_demo_v1/exp_runner_generic_blender_val.py

@@ -0,0 +1,629 @@
+import os
+import logging
+import argparse
+import numpy as np
+from shutil import copyfile
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from rich import print
+from tqdm import tqdm
+from pyhocon import ConfigFactory
+
+import sys
+sys.path.append(os.path.dirname(__file__))
+
+from models.fields import SingleVarianceNetwork
+from models.featurenet import FeatureNet
+from models.trainer_generic import GenericTrainer
+from models.sparse_sdf_network import SparseSdfNetwork
+from models.rendering_network import GeneralRenderingNetwork
+from data.blender_general_narrow_all_eval_new_data import BlenderPerView
+
+
+from datetime import datetime
+
+class Runner:
+    def __init__(self, conf_path, mode='train', is_continue=False,
+                 is_restore=False, restore_lod0=False, local_rank=0):
+
+        # Initial setting
+        self.device = torch.device('cuda:%d' % local_rank)
+        # self.device = torch.device('cuda')
+        self.num_devices = torch.cuda.device_count()
+        self.is_continue = is_continue or (mode == "export_mesh")
+        self.is_restore = is_restore
+        self.restore_lod0 = restore_lod0
+        self.mode = mode
+        self.model_list = []
+        self.logger = logging.getLogger('exp_logger')
+
+        print("detected %d GPUs" % self.num_devices)
+
+        self.conf_path = conf_path
+        self.conf = ConfigFactory.parse_file(conf_path)
+        self.timestamp = None
+        if not self.is_continue:
+            self.timestamp = '_{:%Y_%m_%d_%H_%M_%S}'.format(datetime.now())
+            self.base_exp_dir = self.conf['general.base_exp_dir'] + self.timestamp # jha comment this when testing and use this when training
+        else:
+            self.base_exp_dir = self.conf['general.base_exp_dir']
+        self.conf['general.base_exp_dir'] = self.base_exp_dir # jha use this when testing
+        print("base_exp_dir: " + self.base_exp_dir)
+        os.makedirs(self.base_exp_dir, exist_ok=True)
+        self.iter_step = 0
+        self.val_step = 0
+
+        # trainning parameters
+        self.end_iter = self.conf.get_int('train.end_iter')
+        self.save_freq = self.conf.get_int('train.save_freq')
+        self.report_freq = self.conf.get_int('train.report_freq')
+        self.val_freq = self.conf.get_int('train.val_freq')
+        self.val_mesh_freq = self.conf.get_int('train.val_mesh_freq')
+        self.batch_size = self.num_devices  # use DataParallel to warp
+        self.validate_resolution_level = self.conf.get_int('train.validate_resolution_level')
+        self.learning_rate = self.conf.get_float('train.learning_rate')
+        self.learning_rate_milestone = self.conf.get_list('train.learning_rate_milestone')
+        self.learning_rate_factor = self.conf.get_float('train.learning_rate_factor')
+        self.use_white_bkgd = self.conf.get_bool('train.use_white_bkgd')
+        self.N_rays = self.conf.get_int('train.N_rays')
+
+        # warmup params for sdf gradient
+        self.anneal_start_lod0 = self.conf.get_float('train.anneal_start', default=0)
+        self.anneal_end_lod0 = self.conf.get_float('train.anneal_end', default=0)
+        self.anneal_start_lod1 = self.conf.get_float('train.anneal_start_lod1', default=0)
+        self.anneal_end_lod1 = self.conf.get_float('train.anneal_end_lod1', default=0)
+
+        self.writer = None
+
+        # Networks
+        self.num_lods = self.conf.get_int('model.num_lods')
+
+        self.rendering_network_outside = None
+        self.sdf_network_lod0 = None
+        self.sdf_network_lod1 = None
+        self.variance_network_lod0 = None
+        self.variance_network_lod1 = None
+        self.rendering_network_lod0 = None
+        self.rendering_network_lod1 = None
+        self.pyramid_feature_network = None  # extract 2d pyramid feature maps from images, used for geometry
+        self.pyramid_feature_network_lod1 = None  # may use different feature network for different lod
+
+        # * pyramid_feature_network
+        self.pyramid_feature_network = FeatureNet().to(self.device)
+        self.sdf_network_lod0 = SparseSdfNetwork(**self.conf['model.sdf_network_lod0']).to(self.device)
+        self.variance_network_lod0 = SingleVarianceNetwork(**self.conf['model.variance_network']).to(self.device)
+
+        if self.num_lods > 1:
+            self.sdf_network_lod1 = SparseSdfNetwork(**self.conf['model.sdf_network_lod1']).to(self.device)
+            self.variance_network_lod1 = SingleVarianceNetwork(**self.conf['model.variance_network']).to(self.device)
+
+        self.rendering_network_lod0 = GeneralRenderingNetwork(**self.conf['model.rendering_network']).to(
+            self.device)
+
+        if self.num_lods > 1:
+            self.pyramid_feature_network_lod1 = FeatureNet().to(self.device)
+            self.rendering_network_lod1 = GeneralRenderingNetwork(
+                **self.conf['model.rendering_network_lod1']).to(self.device)
+        if self.mode == 'export_mesh' or self.mode == 'val':
+            # base_exp_dir_to_store = os.path.join(self.base_exp_dir, '{:%Y_%m_%d_%H_%M_%S}'.format(datetime.now()))
+            base_exp_dir_to_store = os.path.join("../", args.specific_dataset_name) #"../gradio_tmp" # MODIFIED
+        else:
+            base_exp_dir_to_store = self.base_exp_dir
+
+        print(f"Store in: {base_exp_dir_to_store}")
+        # Renderer model
+        self.trainer = GenericTrainer(
+            self.rendering_network_outside,
+            self.pyramid_feature_network,
+            self.pyramid_feature_network_lod1,
+            self.sdf_network_lod0,
+            self.sdf_network_lod1,
+            self.variance_network_lod0,
+            self.variance_network_lod1,
+            self.rendering_network_lod0,
+            self.rendering_network_lod1,
+            **self.conf['model.trainer'],
+            timestamp=self.timestamp,
+            base_exp_dir=base_exp_dir_to_store,
+            conf=self.conf)
+
+        self.data_setup()  # * data setup
+
+        self.optimizer_setup()
+
+        # Load checkpoint
+        latest_model_name = None
+        if self.is_continue:
+            model_list_raw = os.listdir(os.path.join(self.base_exp_dir, 'checkpoints'))
+            model_list = []
+            for model_name in model_list_raw:
+                if model_name.startswith('ckpt'):
+                    if model_name[-3:] == 'pth':  # and int(model_name[5:-4]) <= self.end_iter:
+                        model_list.append(model_name)
+            model_list.sort()
+            latest_model_name = model_list[-1]
+
+        if latest_model_name is not None:
+            self.logger.info('Find checkpoint: {}'.format(latest_model_name))
+            self.load_checkpoint(latest_model_name)
+
+        self.trainer = torch.nn.DataParallel(self.trainer).to(self.device)
+
+        if self.mode[:5] == 'train':
+            self.file_backup()
+
+    def optimizer_setup(self):
+        self.params_to_train = self.trainer.get_trainable_params()
+        self.optimizer = torch.optim.Adam(self.params_to_train, lr=self.learning_rate)
+
+    def data_setup(self):
+        """
+        if use ddp, use setup() not prepare_data(),
+        prepare_data() only called on 1 GPU/TPU in distributed
+        :return:
+        """
+
+        self.train_dataset = BlenderPerView(
+            root_dir=self.conf['dataset.trainpath'],
+            split=self.conf.get_string('dataset.train_split', default='train'),
+            split_filepath=self.conf.get_string('dataset.train_split_filepath', default=None),
+            n_views=self.conf['dataset.nviews'],
+            downSample=self.conf['dataset.imgScale_train'],
+            N_rays=self.N_rays,
+            batch_size=self.batch_size,
+            clean_image=True,  # True for training
+            importance_sample=self.conf.get_bool('dataset.importance_sample', default=False),
+            specific_dataset_name = args.specific_dataset_name
+        )
+
+        self.val_dataset = BlenderPerView(
+            root_dir=self.conf['dataset.valpath'],
+            split=self.conf.get_string('dataset.test_split', default='test'),
+            split_filepath=self.conf.get_string('dataset.val_split_filepath', default=None),
+            n_views=3,
+            downSample=self.conf['dataset.imgScale_test'],
+            N_rays=self.N_rays,
+            batch_size=self.batch_size,
+            clean_image=self.conf.get_bool('dataset.mask_out_image',
+                                           default=False) if self.mode != 'train' else False,
+            importance_sample=self.conf.get_bool('dataset.importance_sample', default=False),
+            test_ref_views=self.conf.get_list('dataset.test_ref_views', default=[]),
+            specific_dataset_name = args.specific_dataset_name
+        )
+
+        # item = self.train_dataset.__getitem__(0)
+        self.train_dataloader = DataLoader(self.train_dataset,
+                                           shuffle=True,
+                                           num_workers=4 * self.batch_size,
+                                        #    num_workers=1,
+                                           batch_size=self.batch_size,
+                                           pin_memory=True,
+                                           drop_last=True
+                                           )
+        
+        self.val_dataloader = DataLoader(self.val_dataset,
+                                        #  shuffle=False if self.mode == 'train' else True,
+                                         shuffle=False,
+                                         num_workers=4 * self.batch_size,
+                                        #  num_workers=1,
+                                         batch_size=self.batch_size,
+                                         pin_memory=True,
+                                         drop_last=False
+                                         )
+
+        self.val_dataloader_iterator = iter(self.val_dataloader)  # - should be after "reconstruct_metas_for_gru_fusion"
+
+    def train(self):
+        self.writer = SummaryWriter(log_dir=os.path.join(self.base_exp_dir, 'logs'))
+        res_step = self.end_iter - self.iter_step
+
+        dataloader = self.train_dataloader
+
+        epochs = int(1 + res_step // len(dataloader))
+
+        self.adjust_learning_rate()
+        print("starting training learning rate: {:.5f}".format(self.optimizer.param_groups[0]['lr']))
+
+        background_rgb = None
+        if self.use_white_bkgd:
+            # background_rgb = torch.ones([1, 3]).to(self.device)
+            background_rgb = 1.0
+
+        for epoch_i in range(epochs):
+
+            print("current epoch %d" % epoch_i)
+            dataloader = tqdm(dataloader)
+
+            for batch in dataloader:
+                # print("Checker1:, fetch data")
+                batch['batch_idx'] = torch.tensor([x for x in range(self.batch_size)])  # used to get meta
+
+                # - warmup params
+                if self.num_lods == 1:
+                    alpha_inter_ratio_lod0 = self.get_alpha_inter_ratio(self.anneal_start_lod0, self.anneal_end_lod0)
+                else:
+                    alpha_inter_ratio_lod0 = 1.
+                alpha_inter_ratio_lod1 = self.get_alpha_inter_ratio(self.anneal_start_lod1, self.anneal_end_lod1)
+
+                losses = self.trainer(
+                    batch,
+                    background_rgb=background_rgb,
+                    alpha_inter_ratio_lod0=alpha_inter_ratio_lod0,
+                    alpha_inter_ratio_lod1=alpha_inter_ratio_lod1,
+                    iter_step=self.iter_step,
+                    mode='train',
+                )
+
+                loss_types = ['loss_lod0', 'loss_lod1']
+                # print("[TEST]: weights_sum in trainer return", losses['losses_lod0']['weights_sum'].mean())
+
+                losses_lod0 = losses['losses_lod0']
+                losses_lod1 = losses['losses_lod1']
+                # import ipdb; ipdb.set_trace()
+                loss = 0
+                for loss_type in loss_types:
+                    if losses[loss_type] is not None:
+                        loss = loss + losses[loss_type].mean()
+                # print("Checker4:, begin BP")
+                self.optimizer.zero_grad()
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.params_to_train, 1.0)
+                self.optimizer.step()
+                # print("Checker5:, end BP")
+                self.iter_step += 1
+
+                if self.iter_step % self.report_freq == 0:
+                    self.writer.add_scalar('Loss/loss', loss, self.iter_step)
+
+                    if losses_lod0 is not None:
+                        self.writer.add_scalar('Loss/d_loss_lod0',
+                                               losses_lod0['depth_loss'].mean() if losses_lod0 is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('Loss/sparse_loss_lod0',
+                                               losses_lod0[
+                                                   'sparse_loss'].mean() if losses_lod0 is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('Loss/color_loss_lod0',
+                                               losses_lod0['color_fine_loss'].mean()
+                                               if losses_lod0['color_fine_loss'] is not None else 0,
+                                               self.iter_step)
+
+                        self.writer.add_scalar('statis/psnr_lod0',
+                                               losses_lod0['psnr'].mean()
+                                               if losses_lod0['psnr'] is not None else 0,
+                                               self.iter_step)
+
+                        self.writer.add_scalar('param/variance_lod0',
+                                               1. / torch.exp(self.variance_network_lod0.variance * 10),
+                                               self.iter_step)
+                        self.writer.add_scalar('param/eikonal_loss', losses_lod0['gradient_error_loss'].mean() if losses_lod0 is not None else 0,
+                                               self.iter_step)
+
+                    ######## - lod 1
+                    if self.num_lods > 1:
+                        self.writer.add_scalar('Loss/d_loss_lod1',
+                                               losses_lod1['depth_loss'].mean() if losses_lod1 is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('Loss/sparse_loss_lod1',
+                                               losses_lod1[
+                                                   'sparse_loss'].mean() if losses_lod1 is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('Loss/color_loss_lod1',
+                                               losses_lod1['color_fine_loss'].mean()
+                                               if losses_lod1['color_fine_loss'] is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('statis/sdf_mean_lod1',
+                                               losses_lod1['sdf_mean'].mean() if losses_lod1 is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('statis/psnr_lod1',
+                                               losses_lod1['psnr'].mean()
+                                               if losses_lod1['psnr'] is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('statis/sparseness_0.01_lod1',
+                                               losses_lod1['sparseness_1'].mean()
+                                               if losses_lod1['sparseness_1'] is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('statis/sparseness_0.02_lod1',
+                                               losses_lod1['sparseness_2'].mean()
+                                               if losses_lod1['sparseness_2'] is not None else 0,
+                                               self.iter_step)
+                        self.writer.add_scalar('param/variance_lod1',
+                                               1. / torch.exp(self.variance_network_lod1.variance * 10),
+                                               self.iter_step)
+
+                    print(self.base_exp_dir)
+                    print(
+                        'iter:{:8>d} '
+                        'loss = {:.4f} '
+                        'd_loss_lod0 = {:.4f} '
+                        'color_loss_lod0 = {:.4f} '
+                        'sparse_loss_lod0= {:.4f} '
+                        'd_loss_lod1 = {:.4f} '
+                        'color_loss_lod1 = {:.4f} '
+                        '  lr = {:.5f}'.format(
+                            self.iter_step, loss,
+                            losses_lod0['depth_loss'].mean() if losses_lod0 is not None else 0,
+                            losses_lod0['color_fine_loss'].mean() if losses_lod0 is not None else 0,
+                            losses_lod0['sparse_loss'].mean() if losses_lod0 is not None else 0,
+                            losses_lod1['depth_loss'].mean() if losses_lod1 is not None else 0,
+                            losses_lod1['color_fine_loss'].mean() if losses_lod1 is not None else 0,
+                            self.optimizer.param_groups[0]['lr']))
+
+                    print('alpha_inter_ratio_lod0 = {:.4f} alpha_inter_ratio_lod1 = {:.4f}\n'.format(
+                        alpha_inter_ratio_lod0, alpha_inter_ratio_lod1))
+
+                    if losses_lod0 is not None:
+                        # print("[TEST]: weights_sum in print", losses_lod0['weights_sum'].mean())
+                        # import ipdb; ipdb.set_trace()
+                        print(
+                            'iter:{:8>d} '
+                            'variance = {:.5f} '
+                            'weights_sum = {:.4f} '
+                            'weights_sum_fg = {:.4f} '
+                            'alpha_sum = {:.4f} '
+                            'sparse_weight= {:.4f} '
+                            'background_loss = {:.4f} '
+                            'background_weight = {:.4f} '
+                                .format(
+                                self.iter_step,
+                                losses_lod0['variance'].mean(),
+                                losses_lod0['weights_sum'].mean(),
+                                losses_lod0['weights_sum_fg'].mean(),
+                                losses_lod0['alpha_sum'].mean(),
+                                losses_lod0['sparse_weight'].mean(),
+                                losses_lod0['fg_bg_loss'].mean(),
+                                losses_lod0['fg_bg_weight'].mean(),
+                            ))
+
+                    if losses_lod1 is not None:
+                        print(
+                            'iter:{:8>d} '
+                            'variance = {:.5f} '
+                            ' weights_sum = {:.4f} '
+                            'alpha_sum = {:.4f} '
+                            'fg_bg_loss = {:.4f} '
+                            'fg_bg_weight = {:.4f} '
+                            'sparse_weight= {:.4f} '
+                            'fg_bg_loss = {:.4f} '
+                            'fg_bg_weight = {:.4f} '
+                                .format(
+                                self.iter_step,
+                                losses_lod1['variance'].mean(),
+                                losses_lod1['weights_sum'].mean(),
+                                losses_lod1['alpha_sum'].mean(),
+                                losses_lod1['fg_bg_loss'].mean(),
+                                losses_lod1['fg_bg_weight'].mean(),
+                                losses_lod1['sparse_weight'].mean(),
+                                losses_lod1['fg_bg_loss'].mean(),
+                                losses_lod1['fg_bg_weight'].mean(),
+                            ))
+
+                if self.iter_step % self.save_freq == 0:
+                    self.save_checkpoint()
+
+                if self.iter_step % self.val_freq == 0:
+                    self.validate()
+
+                # - ajust learning rate
+                self.adjust_learning_rate()
+
+    def adjust_learning_rate(self):
+        # - ajust learning rate, cosine learning schedule
+        learning_rate = (np.cos(np.pi * self.iter_step / self.end_iter) + 1.0) * 0.5 * 0.9 + 0.1
+        learning_rate = self.learning_rate * learning_rate
+        for g in self.optimizer.param_groups:
+            g['lr'] = learning_rate
+
+    def get_alpha_inter_ratio(self, start, end):
+        if end == 0.0:
+            return 1.0
+        elif self.iter_step < start:
+            return 0.0
+        else:
+            return np.min([1.0, (self.iter_step - start) / (end - start)])
+
+    def file_backup(self):
+        # copy python file
+        dir_lis = self.conf['general.recording']
+        os.makedirs(os.path.join(self.base_exp_dir, 'recording'), exist_ok=True)
+        for dir_name in dir_lis:
+            cur_dir = os.path.join(self.base_exp_dir, 'recording', dir_name)
+            os.makedirs(cur_dir, exist_ok=True)
+            files = os.listdir(dir_name)
+            for f_name in files:
+                if f_name[-3:] == '.py':
+                    copyfile(os.path.join(dir_name, f_name), os.path.join(cur_dir, f_name))
+
+        # copy configs
+        copyfile(self.conf_path, os.path.join(self.base_exp_dir, 'recording', 'config.conf'))
+
+    def load_checkpoint(self, checkpoint_name):
+
+        def load_state_dict(network, checkpoint, comment):
+            if network is not None:
+                try:
+                    pretrained_dict = checkpoint[comment]
+
+                    model_dict = network.state_dict()
+
+                    # 1. filter out unnecessary keys
+                    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+                    # 2. overwrite entries in the existing state dict
+                    model_dict.update(pretrained_dict)
+                    # 3. load the new state dict
+                    network.load_state_dict(pretrained_dict)
+                except:
+                    print(comment + " load fails")
+
+        checkpoint = torch.load(os.path.join(self.base_exp_dir, 'checkpoints', checkpoint_name),
+                                map_location=self.device)
+
+        load_state_dict(self.rendering_network_outside, checkpoint, 'rendering_network_outside')
+
+        load_state_dict(self.sdf_network_lod0, checkpoint, 'sdf_network_lod0')
+        load_state_dict(self.sdf_network_lod1, checkpoint, 'sdf_network_lod1')
+
+        load_state_dict(self.pyramid_feature_network, checkpoint, 'pyramid_feature_network')
+        load_state_dict(self.pyramid_feature_network_lod1, checkpoint, 'pyramid_feature_network_lod1')
+
+        load_state_dict(self.variance_network_lod0, checkpoint, 'variance_network_lod0')
+        load_state_dict(self.variance_network_lod1, checkpoint, 'variance_network_lod1')
+
+        load_state_dict(self.rendering_network_lod0, checkpoint, 'rendering_network_lod0')
+        load_state_dict(self.rendering_network_lod1, checkpoint, 'rendering_network_lod1')
+
+        if self.restore_lod0:  # use the trained lod0 networks to initialize lod1 networks
+            load_state_dict(self.sdf_network_lod1, checkpoint, 'sdf_network_lod0')
+            load_state_dict(self.pyramid_feature_network_lod1, checkpoint, 'pyramid_feature_network')
+            load_state_dict(self.rendering_network_lod1, checkpoint, 'rendering_network_lod0')
+
+        if self.is_continue and (not self.restore_lod0):
+            try:
+                self.optimizer.load_state_dict(checkpoint['optimizer'])
+            except:
+                print("load optimizer fails")
+            self.iter_step = checkpoint['iter_step']
+            self.val_step = checkpoint['val_step'] if 'val_step' in checkpoint.keys() else 0
+
+        self.logger.info('End')
+
+    def save_checkpoint(self):
+
+        def save_state_dict(network, checkpoint, comment):
+            if network is not None:
+                checkpoint[comment] = network.state_dict()
+
+        checkpoint = {
+            'optimizer': self.optimizer.state_dict(),
+            'iter_step': self.iter_step,
+            'val_step': self.val_step,
+        }
+
+        save_state_dict(self.sdf_network_lod0, checkpoint, "sdf_network_lod0")
+        save_state_dict(self.sdf_network_lod1, checkpoint, "sdf_network_lod1")
+
+        save_state_dict(self.rendering_network_outside, checkpoint, 'rendering_network_outside')
+        save_state_dict(self.rendering_network_lod0, checkpoint, "rendering_network_lod0")
+        save_state_dict(self.rendering_network_lod1, checkpoint, "rendering_network_lod1")
+
+        save_state_dict(self.variance_network_lod0, checkpoint, 'variance_network_lod0')
+        save_state_dict(self.variance_network_lod1, checkpoint, 'variance_network_lod1')
+
+        save_state_dict(self.pyramid_feature_network, checkpoint, 'pyramid_feature_network')
+        save_state_dict(self.pyramid_feature_network_lod1, checkpoint, 'pyramid_feature_network_lod1')
+
+        os.makedirs(os.path.join(self.base_exp_dir, 'checkpoints'), exist_ok=True)
+        torch.save(checkpoint,
+                   os.path.join(self.base_exp_dir, 'checkpoints', 'ckpt_{:0>6d}.pth'.format(self.iter_step)))
+
+    def validate(self, resolution_level=-1):
+        # validate image
+        print("iter_step: ", self.iter_step)
+        self.logger.info('Validate begin')
+        self.val_step += 1
+
+        try:
+            batch = next(self.val_dataloader_iterator)
+        except:
+            self.val_dataloader_iterator = iter(self.val_dataloader)  # reset
+            
+            batch = next(self.val_dataloader_iterator)
+
+
+        background_rgb = None
+        if self.use_white_bkgd:
+            # background_rgb = torch.ones([1, 3]).to(self.device)
+            background_rgb = 1.0
+
+        batch['batch_idx'] = torch.tensor([x for x in range(self.batch_size)])
+
+        # - warmup params
+        if self.num_lods == 1:
+            alpha_inter_ratio_lod0 = self.get_alpha_inter_ratio(self.anneal_start_lod0, self.anneal_end_lod0)
+        else:
+            alpha_inter_ratio_lod0 = 1.
+        alpha_inter_ratio_lod1 = self.get_alpha_inter_ratio(self.anneal_start_lod1, self.anneal_end_lod1)
+
+        self.trainer(
+            batch,
+            background_rgb=background_rgb,
+            alpha_inter_ratio_lod0=alpha_inter_ratio_lod0,
+            alpha_inter_ratio_lod1=alpha_inter_ratio_lod1,
+            iter_step=self.iter_step,
+            save_vis=True,
+            mode='val',
+        )
+
+
+    def export_mesh(self, resolution_level=-1):
+        print("iter_step: ", self.iter_step)
+        self.logger.info('Validate begin')
+        self.val_step += 1
+
+        try:
+            batch = next(self.val_dataloader_iterator)
+        except:
+            self.val_dataloader_iterator = iter(self.val_dataloader)  # reset
+            
+            batch = next(self.val_dataloader_iterator)
+
+
+        background_rgb = None
+        if self.use_white_bkgd:
+            background_rgb = 1.0
+
+        batch['batch_idx'] = torch.tensor([x for x in range(self.batch_size)])
+
+        # - warmup params
+        if self.num_lods == 1:
+            alpha_inter_ratio_lod0 = self.get_alpha_inter_ratio(self.anneal_start_lod0, self.anneal_end_lod0)
+        else:
+            alpha_inter_ratio_lod0 = 1.
+        alpha_inter_ratio_lod1 = self.get_alpha_inter_ratio(self.anneal_start_lod1, self.anneal_end_lod1)
+        self.trainer(
+            batch,
+            background_rgb=background_rgb,
+            alpha_inter_ratio_lod0=alpha_inter_ratio_lod0,
+            alpha_inter_ratio_lod1=alpha_inter_ratio_lod1,
+            iter_step=self.iter_step,
+            save_vis=True,
+            mode='export_mesh',
+        )
+
+
+if __name__ == '__main__':
+    # torch.set_default_tensor_type('torch.cuda.FloatTensor')
+    torch.set_default_dtype(torch.float32)
+    FORMAT = "[%(filename)s:%(lineno)s - %(funcName)20s() ] %(message)s"
+    logging.basicConfig(level=logging.INFO, format=FORMAT)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--conf', type=str, default='./confs/base.conf')
+    parser.add_argument('--mode', type=str, default='train')
+    parser.add_argument('--threshold', type=float, default=0.0)
+    parser.add_argument('--is_continue', default=False, action="store_true")
+    parser.add_argument('--is_restore', default=False, action="store_true")
+    parser.add_argument('--is_finetune', default=False, action="store_true")
+    parser.add_argument('--train_from_scratch', default=False, action="store_true")
+    parser.add_argument('--restore_lod0', default=False, action="store_true")
+    parser.add_argument('--local_rank', type=int, default=0)
+    parser.add_argument('--specific_dataset_name', type=str, default='GSO')
+
+
+    args = parser.parse_args()
+
+    torch.cuda.set_device(args.local_rank)
+    torch.backends.cudnn.benchmark = True  # ! make training 2x faster
+
+    runner = Runner(args.conf, args.mode, args.is_continue, args.is_restore, args.restore_lod0,
+                    args.local_rank)
+
+    if args.mode == 'train':
+        runner.train()
+    elif args.mode == 'val':
+        for i in range(len(runner.val_dataset)):
+            runner.validate()
+    elif args.mode == 'export_mesh':
+        for i in range(len(runner.val_dataset)):
+            runner.export_mesh()

SparseNeuS_demo_v1/loss/color_loss.py

@@ -0,0 +1,152 @@
+import torch
+import torch.nn as nn
+from loss.ncc import NCC
+
+
+class Normalize(nn.Module):
+    def __init__(self):
+        super(Normalize, self).__init__()
+
+    def forward(self, bottom):
+        qn = torch.norm(bottom, p=2, dim=1).unsqueeze(dim=1) + 1e-12
+        top = bottom.div(qn)
+
+        return top
+
+
+class OcclusionColorLoss(nn.Module):
+    def __init__(self, alpha=1, beta=0.025, gama=0.01, occlusion_aware=True, weight_thred=[0.6]):
+        super(OcclusionColorLoss, self).__init__()
+        self.alpha = alpha
+        self.beta = beta
+        self.gama = gama
+        self.occlusion_aware = occlusion_aware
+        self.eps = 1e-4
+
+        self.weight_thred = weight_thred
+        self.adjuster = ParamAdjuster(self.weight_thred, self.beta)
+
+    def forward(self, pred, gt, weight, mask, detach=False, occlusion_aware=True):
+        """
+
+        :param pred: [N_pts, 3]
+        :param gt: [N_pts, 3]
+        :param weight: [N_pts]
+        :param mask: [N_pts]
+        :return:
+        """
+        if detach:
+            weight = weight.detach()
+
+        error = torch.abs(pred - gt).sum(dim=-1, keepdim=False)  # [N_pts]
+        error = error[mask]
+
+        if not (self.occlusion_aware and occlusion_aware):
+            return torch.mean(error), torch.mean(error)
+
+        beta = self.adjuster(weight.mean())
+
+        # weight = weight[mask]
+        weight = weight.clamp(0.0, 1.0)
+        term1 = self.alpha * torch.mean(weight[mask] * error)
+        term2 = beta * torch.log(1 - weight + self.eps).mean()
+        term3 = self.gama * torch.log(weight + self.eps).mean()
+
+        return term1 + term2 + term3, term1
+
+
+class OcclusionColorPatchLoss(nn.Module):
+    def __init__(self, alpha=1, beta=0.025, gama=0.015,
+                 occlusion_aware=True, type='l1', h_patch_size=3, weight_thred=[0.6]):
+        super(OcclusionColorPatchLoss, self).__init__()
+        self.alpha = alpha
+        self.beta = beta
+        self.gama = gama
+        self.occlusion_aware = occlusion_aware
+        self.type = type  # 'l1' or 'ncc' loss
+        self.ncc = NCC(h_patch_size=h_patch_size)
+        self.eps = 1e-4
+        self.weight_thred = weight_thred
+
+        self.adjuster = ParamAdjuster(self.weight_thred, self.beta)
+
+        print("type {} patch_size {}  beta {}  gama {}  weight_thred {}".format(type, h_patch_size, beta, gama,
+                                                                                weight_thred))
+
+    def forward(self, pred, gt, weight, mask, penalize_ratio=0.9, detach=False, occlusion_aware=True):
+        """
+
+        :param pred: [N_pts, Npx, 3]
+        :param gt: [N_pts, Npx, 3]
+        :param weight: [N_pts]
+        :param mask: [N_pts]
+        :return:
+        """
+
+        if detach:
+            weight = weight.detach()
+
+        if self.type == 'l1':
+            error = torch.abs(pred - gt).mean(dim=-1, keepdim=False).sum(dim=-1, keepdim=False)  # [N_pts]
+        elif self.type == 'ncc':
+            error = 1 - self.ncc(pred[:, None, :, :], gt)[:, 0]  # ncc 1 positive, -1 negative
+            error, indices = torch.sort(error)
+            mask = torch.index_select(mask, 0, index=indices)
+            mask[int(penalize_ratio * mask.shape[0]):] = False  # can help boundaries
+        elif self.type == 'ssd':
+            error = ((pred - gt) ** 2).mean(dim=-1, keepdim=False).sum(dim=-1, keepdims=False)
+
+        error = error[mask]
+        if not (self.occlusion_aware and occlusion_aware):
+            return torch.mean(error), torch.mean(error), 0.
+
+        # * weight adjuster
+        beta = self.adjuster(weight.mean())
+
+        # weight = weight[mask]
+        weight = weight.clamp(0.0, 1.0)
+
+        term1 = self.alpha * torch.mean(weight[mask] * error)
+        term2 = beta * torch.log(1 - weight + self.eps).mean()
+        term3 = self.gama * torch.log(weight + self.eps).mean()
+
+        return term1 + term2 + term3, term1, beta
+
+
+class ParamAdjuster(nn.Module):
+    def __init__(self, weight_thred, param):
+        super(ParamAdjuster, self).__init__()
+        self.weight_thred = weight_thred
+        self.thred_num = len(weight_thred)
+        self.param = param
+        self.global_step = 0
+        self.statis_window = 100
+        self.counter = 0
+        self.adjusted = False
+        self.adjusted_step = 0
+        self.thred_idx = 0
+
+    def reset(self):
+        self.counter = 0
+        self.adjusted = False
+
+    def adjust(self):
+        if (self.counter / self.statis_window) > 0.3:
+            self.param = self.param + 0.005
+            self.adjusted = True
+            self.adjusted_step = self.global_step
+            self.thred_idx += 1
+            print("adjusted param, now {}".format(self.param))
+
+    def forward(self, weight_mean):
+        self.global_step += 1
+
+        if (self.global_step % self.statis_window == 0) and self.adjusted is False:
+            self.adjust()
+            self.reset()
+
+        if self.thred_idx < self.thred_num:
+            if weight_mean < self.weight_thred[self.thred_idx] and (not self.adjusted):
+                self.counter += 1
+
+        return self.param

SparseNeuS_demo_v1/loss/depth_loss.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DepthLoss(nn.Module):
+    def __init__(self, type='l1'):
+        super(DepthLoss, self).__init__()
+        self.type = type
+
+
+    def forward(self, depth_pred, depth_gt, mask=None):
+            if (depth_gt < 0).sum() > 0:
+                # print("no depth loss")
+                return torch.tensor(0.0).to(depth_pred.device)
+            if mask is not None:
+                mask_d = (depth_gt > 0).float()
+
+                mask = mask * mask_d
+
+                mask_sum = mask.sum() + 1e-5
+                depth_error = (depth_pred - depth_gt) * mask
+                depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device),
+                                    reduction='sum') / mask_sum
+            else:
+                depth_error = depth_pred - depth_gt
+                depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device),
+                                    reduction='mean')
+            return depth_loss
+
+def forward(self, depth_pred, depth_gt, mask=None):
+        if mask is not None:
+            mask_d = (depth_gt > 0).float()
+
+            mask = mask * mask_d
+
+            mask_sum = mask.sum() + 1e-5
+            depth_error = (depth_pred - depth_gt) * mask
+            depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device),
+                                   reduction='sum') / mask_sum
+        else:
+            depth_error = depth_pred - depth_gt
+            depth_loss = F.l1_loss(depth_error, torch.zeros_like(depth_error).to(depth_error.device),
+                                   reduction='mean')
+        return depth_loss
+
+class DepthSmoothLoss(nn.Module):
+    def __init__(self):
+        super(DepthSmoothLoss, self).__init__()
+
+    def forward(self, disp, img, mask):
+        """
+        Computes the smoothness loss for a disparity image
+        The color image is used for edge-aware smoothness
+        :param disp: [B, 1, H, W]
+        :param img: [B, 1, H, W]
+        :param mask: [B, 1, H, W]
+        :return:
+        """
+        grad_disp_x = torch.abs(disp[:, :, :, :-1] - disp[:, :, :, 1:])
+        grad_disp_y = torch.abs(disp[:, :, :-1, :] - disp[:, :, 1:, :])
+
+        grad_img_x = torch.mean(torch.abs(img[:, :, :, :-1] - img[:, :, :, 1:]), 1, keepdim=True)
+        grad_img_y = torch.mean(torch.abs(img[:, :, :-1, :] - img[:, :, 1:, :]), 1, keepdim=True)
+
+        grad_disp_x *= torch.exp(-grad_img_x)
+        grad_disp_y *= torch.exp(-grad_img_y)
+
+        grad_disp = (grad_disp_x * mask[:, :, :, :-1]).mean() + (grad_disp_y * mask[:, :, :-1, :]).mean()
+
+        return grad_disp

SparseNeuS_demo_v1/loss/depth_metric.py

@@ -0,0 +1,240 @@
+import numpy as np
+
+
+def l1(depth1, depth2):
+    """
+    Computes the l1 errors between the two depth maps.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth1:  one depth map
+    depth2:  another depth map
+
+    Returns:
+        L1(log)
+
+    """
+    assert (np.all(np.isfinite(depth1) & np.isfinite(depth2) & (depth1 >= 0) & (depth2 >= 0)))
+    diff = depth1 - depth2
+    num_pixels = float(diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sum(np.absolute(diff)) / num_pixels
+
+
+def l1_inverse(depth1, depth2):
+    """
+    Computes the l1 errors between inverses of two depth maps.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth1:  one depth map
+    depth2:  another depth map
+
+    Returns:
+        L1(log)
+
+    """
+    assert (np.all(np.isfinite(depth1) & np.isfinite(depth2) & (depth1 >= 0) & (depth2 >= 0)))
+    diff = np.reciprocal(depth1) - np.reciprocal(depth2)
+    num_pixels = float(diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sum(np.absolute(diff)) / num_pixels
+
+
+def rmse_log(depth1, depth2):
+    """
+    Computes the root min square errors between the logs of two depth maps.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth1:  one depth map
+    depth2:  another depth map
+
+    Returns:
+        RMSE(log)
+
+    """
+    assert (np.all(np.isfinite(depth1) & np.isfinite(depth2) & (depth1 >= 0) & (depth2 >= 0)))
+    log_diff = np.log(depth1) - np.log(depth2)
+    num_pixels = float(log_diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sqrt(np.sum(np.square(log_diff)) / num_pixels)
+
+
+def rmse(depth1, depth2):
+    """
+    Computes the root min square errors between the two depth maps.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth1:  one depth map
+    depth2:  another depth map
+
+    Returns:
+        RMSE(log)
+
+    """
+    assert (np.all(np.isfinite(depth1) & np.isfinite(depth2) & (depth1 >= 0) & (depth2 >= 0)))
+    diff = depth1 - depth2
+    num_pixels = float(diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sqrt(np.sum(np.square(diff)) / num_pixels)
+
+
+def scale_invariant(depth1, depth2):
+    """
+    Computes the scale invariant loss based on differences of logs of depth maps.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth1:  one depth map
+    depth2:  another depth map
+
+    Returns:
+        scale_invariant_distance
+
+    """
+    # sqrt(Eq. 3)
+    assert (np.all(np.isfinite(depth1) & np.isfinite(depth2) & (depth1 >= 0) & (depth2 >= 0)))
+    log_diff = np.log(depth1) - np.log(depth2)
+    num_pixels = float(log_diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sqrt(np.sum(np.square(log_diff)) / num_pixels - np.square(np.sum(log_diff)) / np.square(num_pixels))
+
+
+def abs_relative(depth_pred, depth_gt):
+    """
+    Computes relative absolute distance.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth_pred:  depth map prediction
+    depth_gt:    depth map ground truth
+
+    Returns:
+        abs_relative_distance
+
+    """
+    assert (np.all(np.isfinite(depth_pred) & np.isfinite(depth_gt) & (depth_pred >= 0) & (depth_gt >= 0)))
+    diff = depth_pred - depth_gt
+    num_pixels = float(diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sum(np.absolute(diff) / depth_gt) / num_pixels
+
+
+def avg_log10(depth1, depth2):
+    """
+    Computes average log_10 error (Liu, Neural Fields, 2015).
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth1:  one depth map
+    depth2:  another depth map
+
+    Returns:
+        abs_relative_distance
+
+    """
+    assert (np.all(np.isfinite(depth1) & np.isfinite(depth2) & (depth1 >= 0) & (depth2 >= 0)))
+    log_diff = np.log10(depth1) - np.log10(depth2)
+    num_pixels = float(log_diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sum(np.absolute(log_diff)) / num_pixels
+
+
+def sq_relative(depth_pred, depth_gt):
+    """
+    Computes relative squared distance.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth_pred:  depth map prediction
+    depth_gt:    depth map ground truth
+
+    Returns:
+        squared_relative_distance
+
+    """
+    assert (np.all(np.isfinite(depth_pred) & np.isfinite(depth_gt) & (depth_pred >= 0) & (depth_gt >= 0)))
+    diff = depth_pred - depth_gt
+    num_pixels = float(diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return np.sum(np.square(diff) / depth_gt) / num_pixels
+
+
+def ratio_threshold(depth1, depth2, threshold):
+    """
+    Computes the percentage of pixels for which the ratio of the two depth maps is less than a given threshold.
+    Takes preprocessed depths (no nans, infs and non-positive values)
+
+    depth1:  one depth map
+    depth2:  another depth map
+
+    Returns:
+        percentage of pixels with ratio less than the threshold
+
+    """
+    assert (threshold > 0.)
+    assert (np.all(np.isfinite(depth1) & np.isfinite(depth2) & (depth1 >= 0) & (depth2 >= 0)))
+    log_diff = np.log(depth1) - np.log(depth2)
+    num_pixels = float(log_diff.size)
+
+    if num_pixels == 0:
+        return np.nan
+    else:
+        return float(np.sum(np.absolute(log_diff) < np.log(threshold))) / num_pixels
+
+
+def compute_depth_errors(depth_pred, depth_gt, valid_mask):
+    """
+    Computes different distance measures between two depth maps.
+
+    depth_pred:           depth map prediction
+    depth_gt:             depth map ground truth
+    distances_to_compute: which distances to compute
+
+    Returns:
+        a dictionary with computed distances, and the number of valid pixels
+
+    """
+    depth_pred = depth_pred[valid_mask]
+    depth_gt = depth_gt[valid_mask]
+    num_valid = np.sum(valid_mask)
+
+    distances_to_compute = ['l1',
+                            'l1_inverse',
+                            'scale_invariant',
+                            'abs_relative',
+                            'sq_relative',
+                            'avg_log10',
+                            'rmse_log',
+                            'rmse',
+                            'ratio_threshold_1.25',
+                            'ratio_threshold_1.5625',
+                            'ratio_threshold_1.953125']
+
+    results = {'num_valid': num_valid}
+    for dist in distances_to_compute:
+        if dist.startswith('ratio_threshold'):
+            threshold = float(dist.split('_')[-1])
+            results[dist] = ratio_threshold(depth_pred, depth_gt, threshold)
+        else:
+            results[dist] = globals()[dist](depth_pred, depth_gt)
+
+    return results

SparseNeuS_demo_v1/loss/ncc.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from math import exp, sqrt
+
+
+class NCC(torch.nn.Module):
+    def __init__(self, h_patch_size, mode='rgb'):
+        super(NCC, self).__init__()
+        self.window_size = 2 * h_patch_size + 1
+        self.mode = mode  # 'rgb' or 'gray'
+        self.channel = 3
+        self.register_buffer("window", create_window(self.window_size, self.channel))
+
+    def forward(self, img_pred, img_gt):
+        """
+        :param img_pred: [Npx, nviews, npatch, c]
+        :param img_gt: [Npx, npatch, c]
+        :return:
+        """
+        ntotpx, nviews, npatch, channels = img_pred.shape
+
+        patch_size = int(sqrt(npatch))
+        patch_img_pred = img_pred.reshape(ntotpx, nviews, patch_size, patch_size, channels).permute(0, 1, 4, 2,
+                                                                                                    3).contiguous()
+        patch_img_gt = img_gt.reshape(ntotpx, patch_size, patch_size, channels).permute(0, 3, 1, 2)
+
+        return _ncc(patch_img_pred, patch_img_gt, self.window, self.channel)
+
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
+    return gauss / gauss.sum()
+
+
+def create_window(window_size, channel, std=1.5):
+    _1D_window = gaussian(window_size, std).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).unsqueeze(0).unsqueeze(0)
+    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+    return window
+
+
+def _ncc(pred, gt, window, channel):
+    ntotpx, nviews, nc, h, w = pred.shape
+    flat_pred = pred.view(-1, nc, h, w)
+    mu1 = F.conv2d(flat_pred, window, padding=0, groups=channel).view(ntotpx, nviews, nc)
+    mu2 = F.conv2d(gt, window, padding=0, groups=channel).view(ntotpx, nc)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2).unsqueeze(1)  # (ntotpx, 1, nc)
+
+    sigma1_sq = F.conv2d(flat_pred * flat_pred, window, padding=0, groups=channel).view(ntotpx, nviews, nc) - mu1_sq
+    sigma2_sq = F.conv2d(gt * gt, window, padding=0, groups=channel).view(ntotpx, 1, 3) - mu2_sq
+
+    sigma1 = torch.sqrt(sigma1_sq + 1e-4)
+    sigma2 = torch.sqrt(sigma2_sq + 1e-4)
+
+    pred_norm = (pred - mu1[:, :, :, None, None]) / (sigma1[:, :, :, None, None] + 1e-8)  # [ntotpx, nviews, nc, h, w]
+    gt_norm = (gt[:, None, :, :, :] - mu2[:, None, :, None, None]) / (
+            sigma2[:, :, :, None, None] + 1e-8)  # ntotpx, nc, h, w
+
+    ncc = F.conv2d((pred_norm * gt_norm).view(-1, nc, h, w), window, padding=0, groups=channel).view(
+        ntotpx, nviews, nc)
+
+    return torch.mean(ncc, dim=2)

SparseNeuS_demo_v1/models/embedder.py

@@ -0,0 +1,101 @@
+import torch
+import torch.nn as nn
+
+""" Positional encoding embedding. Code was taken from https://github.com/bmild/nerf. """
+
+
+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']:
+                if self.kwargs['normalize']:
+                    embed_fns.append(lambda x, p_fn=p_fn,
+                                            freq=freq: p_fn(x * freq) / freq)
+                else:
+                    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, normalize=False, input_dims=3):
+    embed_kwargs = {
+        'include_input': True,
+        'input_dims': input_dims,
+        'max_freq_log2': multires - 1,
+        'num_freqs': multires,
+        'normalize': normalize,
+        'log_sampling': True,
+        'periodic_fns': [torch.sin, torch.cos],
+    }
+
+    embedder_obj = Embedder(**embed_kwargs)
+
+    def embed(x, eo=embedder_obj): return eo.embed(x)
+
+    return embed, embedder_obj.out_dim
+
+
+class Embedding(nn.Module):
+    def __init__(self, in_channels, N_freqs, logscale=True, normalize=False):
+        """
+        Defines a function that embeds x to (x, sin(2^k x), cos(2^k x), ...)
+        in_channels: number of input channels (3 for both xyz and direction)
+        """
+        super(Embedding, self).__init__()
+        self.N_freqs = N_freqs
+        self.in_channels = in_channels
+        self.funcs = [torch.sin, torch.cos]
+        self.out_channels = in_channels * (len(self.funcs) * N_freqs + 1)
+        self.normalize = normalize
+
+        if logscale:
+            self.freq_bands = 2 ** torch.linspace(0, N_freqs - 1, N_freqs)
+        else:
+            self.freq_bands = torch.linspace(1, 2 ** (N_freqs - 1), N_freqs)
+
+    def forward(self, x):
+        """
+        Embeds x to (x, sin(2^k x), cos(2^k x), ...)
+        Different from the paper, "x" is also in the output
+        See https://github.com/bmild/nerf/issues/12
+
+        Inputs:
+            x: (B, self.in_channels)
+
+        Outputs:
+            out: (B, self.out_channels)
+        """
+        out = [x]
+        for freq in self.freq_bands:
+            for func in self.funcs:
+                if self.normalize:
+                    out += [func(freq * x) / freq]
+                else:
+                    out += [func(freq * x)]
+
+        return torch.cat(out, -1)

SparseNeuS_demo_v1/models/fast_renderer.py

@@ -0,0 +1,316 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from icecream import ic
+
+
+# - neus: use sphere-tracing to speed up depth maps extraction
+# This code snippet is heavily borrowed from IDR.
+class FastRenderer(nn.Module):
+    def __init__(self):
+        super(FastRenderer, self).__init__()
+
+        self.sdf_threshold = 5e-5
+        self.line_search_step = 0.5
+        self.line_step_iters = 1
+        self.sphere_tracing_iters = 10
+        self.n_steps = 100
+        self.n_secant_steps = 8
+
+        # - use sdf_network to inference sdf value or directly interpolate sdf value from precomputed sdf_volume
+        self.network_inference = False
+
+    def extract_depth_maps(self, rays_o, rays_d, near, far, sdf_network, conditional_volume):
+        with torch.no_grad():
+            curr_start_points, network_object_mask, acc_start_dis = self.get_intersection(
+                rays_o, rays_d, near, far,
+                sdf_network, conditional_volume)
+
+        network_object_mask = network_object_mask.reshape(-1)
+
+        return network_object_mask, acc_start_dis
+
+    def get_intersection(self, rays_o, rays_d, near, far, sdf_network, conditional_volume):
+        device = rays_o.device
+        num_pixels, _ = rays_d.shape
+
+        curr_start_points, unfinished_mask_start, acc_start_dis, acc_end_dis, min_dis, max_dis = \
+            self.sphere_tracing(rays_o, rays_d, near, far, sdf_network, conditional_volume)
+
+        network_object_mask = (acc_start_dis < acc_end_dis)
+
+        # The non convergent rays should be handled by the sampler
+        sampler_mask = unfinished_mask_start
+        sampler_net_obj_mask = torch.zeros_like(sampler_mask).bool().to(device)
+        if sampler_mask.sum() > 0:
+            # sampler_min_max = torch.zeros((num_pixels, 2)).to(device)
+            # sampler_min_max[sampler_mask, 0] = acc_start_dis[sampler_mask]
+            # sampler_min_max[sampler_mask, 1] = acc_end_dis[sampler_mask]
+
+            # ray_sampler(self, rays_o, rays_d, near, far, sampler_mask):
+            sampler_pts, sampler_net_obj_mask, sampler_dists = self.ray_sampler(rays_o,
+                                                                                rays_d,
+                                                                                acc_start_dis,
+                                                                                acc_end_dis,
+                                                                                sampler_mask,
+                                                                                sdf_network,
+                                                                                conditional_volume
+                                                                                )
+
+            curr_start_points[sampler_mask] = sampler_pts[sampler_mask]
+            acc_start_dis[sampler_mask] = sampler_dists[sampler_mask][:, None]
+            network_object_mask[sampler_mask] = sampler_net_obj_mask[sampler_mask][:, None]
+
+        # print('----------------------------------------------------------------')
+        # print('RayTracing: object = {0}/{1}, secant on {2}/{3}.'
+        #       .format(network_object_mask.sum(), len(network_object_mask), sampler_net_obj_mask.sum(),
+        #               sampler_mask.sum()))
+        # print('----------------------------------------------------------------')
+
+        return curr_start_points, network_object_mask, acc_start_dis
+
+    def sphere_tracing(self, rays_o, rays_d, near, far, sdf_network, conditional_volume):
+        ''' Run sphere tracing algorithm for max iterations from both sides of unit sphere intersection '''
+
+        device = rays_o.device
+
+        unfinished_mask_start = (near < far).reshape(-1).clone()
+        unfinished_mask_end = (near < far).reshape(-1).clone()
+
+        # Initialize start current points
+        curr_start_points = rays_o + rays_d * near
+        acc_start_dis = near.clone()
+
+        # Initialize end current points
+        curr_end_points = rays_o + rays_d * far
+        acc_end_dis = far.clone()
+
+        # Initizlize min and max depth
+        min_dis = acc_start_dis.clone()
+        max_dis = acc_end_dis.clone()
+
+        # Iterate on the rays (from both sides) till finding a surface
+        iters = 0
+
+        next_sdf_start = torch.zeros_like(acc_start_dis).to(device)
+
+        if self.network_inference:
+            sdf_func = sdf_network.sdf
+        else:
+            sdf_func = sdf_network.sdf_from_sdfvolume
+
+        next_sdf_start[unfinished_mask_start] = sdf_func(
+            curr_start_points[unfinished_mask_start],
+            conditional_volume, lod=0, gru_fusion=False)['sdf_pts_scale%d' % 0]
+
+        next_sdf_end = torch.zeros_like(acc_end_dis).to(device)
+        next_sdf_end[unfinished_mask_end] = sdf_func(curr_end_points[unfinished_mask_end],
+                                                     conditional_volume, lod=0, gru_fusion=False)[
+            'sdf_pts_scale%d' % 0]
+
+        while True:
+            # Update sdf
+            curr_sdf_start = torch.zeros_like(acc_start_dis).to(device)
+            curr_sdf_start[unfinished_mask_start] = next_sdf_start[unfinished_mask_start]
+            curr_sdf_start[curr_sdf_start <= self.sdf_threshold] = 0
+
+            curr_sdf_end = torch.zeros_like(acc_end_dis).to(device)
+            curr_sdf_end[unfinished_mask_end] = next_sdf_end[unfinished_mask_end]
+            curr_sdf_end[curr_sdf_end <= self.sdf_threshold] = 0
+
+            # Update masks
+            unfinished_mask_start = unfinished_mask_start & (curr_sdf_start > self.sdf_threshold).reshape(-1)
+            unfinished_mask_end = unfinished_mask_end & (curr_sdf_end > self.sdf_threshold).reshape(-1)
+
+            if (
+                    unfinished_mask_start.sum() == 0 and unfinished_mask_end.sum() == 0) or iters == self.sphere_tracing_iters:
+                break
+            iters += 1
+
+            # Make step
+            # Update distance
+            acc_start_dis = acc_start_dis + curr_sdf_start
+            acc_end_dis = acc_end_dis - curr_sdf_end
+
+            # Update points
+            curr_start_points = rays_o + acc_start_dis * rays_d
+            curr_end_points = rays_o + acc_end_dis * rays_d
+
+            # Fix points which wrongly crossed the surface
+            next_sdf_start = torch.zeros_like(acc_start_dis).to(device)
+            if unfinished_mask_start.sum() > 0:
+                next_sdf_start[unfinished_mask_start] = sdf_func(curr_start_points[unfinished_mask_start],
+                                                                 conditional_volume, lod=0, gru_fusion=False)[
+                    'sdf_pts_scale%d' % 0]
+
+            next_sdf_end = torch.zeros_like(acc_end_dis).to(device)
+            if unfinished_mask_end.sum() > 0:
+                next_sdf_end[unfinished_mask_end] = sdf_func(curr_end_points[unfinished_mask_end],
+                                                             conditional_volume, lod=0, gru_fusion=False)[
+                    'sdf_pts_scale%d' % 0]
+
+            not_projected_start = (next_sdf_start < 0).reshape(-1)
+            not_projected_end = (next_sdf_end < 0).reshape(-1)
+            not_proj_iters = 0
+
+            while (
+                    not_projected_start.sum() > 0 or not_projected_end.sum() > 0) and not_proj_iters < self.line_step_iters:
+                # Step backwards
+                if not_projected_start.sum() > 0:
+                    acc_start_dis[not_projected_start] -= ((1 - self.line_search_step) / (2 ** not_proj_iters)) * \
+                                                          curr_sdf_start[not_projected_start]
+                    curr_start_points[not_projected_start] = (rays_o + acc_start_dis * rays_d)[not_projected_start]
+
+                    next_sdf_start[not_projected_start] = sdf_func(
+                        curr_start_points[not_projected_start],
+                        conditional_volume, lod=0, gru_fusion=False)['sdf_pts_scale%d' % 0]
+
+                if not_projected_end.sum() > 0:
+                    acc_end_dis[not_projected_end] += ((1 - self.line_search_step) / (2 ** not_proj_iters)) * \
+                                                      curr_sdf_end[
+                                                          not_projected_end]
+                    curr_end_points[not_projected_end] = (rays_o + acc_end_dis * rays_d)[not_projected_end]
+
+                    # Calc sdf
+
+                    next_sdf_end[not_projected_end] = sdf_func(
+                        curr_end_points[not_projected_end],
+                        conditional_volume, lod=0, gru_fusion=False)['sdf_pts_scale%d' % 0]
+
+                # Update mask
+                not_projected_start = (next_sdf_start < 0).reshape(-1)
+                not_projected_end = (next_sdf_end < 0).reshape(-1)
+                not_proj_iters += 1
+
+            unfinished_mask_start = unfinished_mask_start & (acc_start_dis < acc_end_dis).reshape(-1)
+            unfinished_mask_end = unfinished_mask_end & (acc_start_dis < acc_end_dis).reshape(-1)
+
+        return curr_start_points, unfinished_mask_start, acc_start_dis, acc_end_dis, min_dis, max_dis
+
+    def ray_sampler(self, rays_o, rays_d, near, far, sampler_mask, sdf_network, conditional_volume):
+        ''' Sample the ray in a given range and run secant on rays which have sign transition '''
+        device = rays_o.device
+        num_pixels, _ = rays_d.shape
+        sampler_pts = torch.zeros(num_pixels, 3).to(device).float()
+        sampler_dists = torch.zeros(num_pixels).to(device).float()
+
+        intervals_dist = torch.linspace(0, 1, steps=self.n_steps).to(device).view(1, -1)
+
+        pts_intervals = near + intervals_dist * (far - near)
+        points = rays_o[:, None, :] + pts_intervals[:, :, None] * rays_d[:, None, :]
+
+        # Get the non convergent rays
+        mask_intersect_idx = torch.nonzero(sampler_mask).flatten()
+        points = points.reshape((-1, self.n_steps, 3))[sampler_mask, :, :]
+        pts_intervals = pts_intervals.reshape((-1, self.n_steps))[sampler_mask]
+
+        if self.network_inference:
+            sdf_func = sdf_network.sdf
+        else:
+            sdf_func = sdf_network.sdf_from_sdfvolume
+
+        sdf_val_all = []
+        for pnts in torch.split(points.reshape(-1, 3), 100000, dim=0):
+            sdf_val_all.append(sdf_func(pnts,
+                                        conditional_volume, lod=0, gru_fusion=False)['sdf_pts_scale%d' % 0])
+        sdf_val = torch.cat(sdf_val_all).reshape(-1, self.n_steps)
+
+        tmp = torch.sign(sdf_val) * torch.arange(self.n_steps, 0, -1).to(device).float().reshape(
+            (1, self.n_steps))  # Force argmin to return the first min value
+        sampler_pts_ind = torch.argmin(tmp, -1)
+        sampler_pts[mask_intersect_idx] = points[torch.arange(points.shape[0]), sampler_pts_ind, :]
+        sampler_dists[mask_intersect_idx] = pts_intervals[torch.arange(pts_intervals.shape[0]), sampler_pts_ind]
+
+        net_surface_pts = (sdf_val[torch.arange(sdf_val.shape[0]), sampler_pts_ind] < 0)
+
+        # take points with minimal SDF value for P_out pixels
+        p_out_mask = ~net_surface_pts
+        n_p_out = p_out_mask.sum()
+        if n_p_out > 0:
+            out_pts_idx = torch.argmin(sdf_val[p_out_mask, :], -1)
+            sampler_pts[mask_intersect_idx[p_out_mask]] = points[p_out_mask, :, :][torch.arange(n_p_out), out_pts_idx,
+                                                          :]
+            sampler_dists[mask_intersect_idx[p_out_mask]] = pts_intervals[p_out_mask, :][
+                torch.arange(n_p_out), out_pts_idx]
+
+        # Get Network object mask
+        sampler_net_obj_mask = sampler_mask.clone()
+        sampler_net_obj_mask[mask_intersect_idx[~net_surface_pts]] = False
+
+        # Run Secant method
+        secant_pts = net_surface_pts
+        n_secant_pts = secant_pts.sum()
+        if n_secant_pts > 0:
+            # Get secant z predictions
+            z_high = pts_intervals[torch.arange(pts_intervals.shape[0]), sampler_pts_ind][secant_pts]
+            sdf_high = sdf_val[torch.arange(sdf_val.shape[0]), sampler_pts_ind][secant_pts]
+            z_low = pts_intervals[secant_pts][torch.arange(n_secant_pts), sampler_pts_ind[secant_pts] - 1]
+            sdf_low = sdf_val[secant_pts][torch.arange(n_secant_pts), sampler_pts_ind[secant_pts] - 1]
+
+            cam_loc_secant = rays_o[mask_intersect_idx[secant_pts]]
+            ray_directions_secant = rays_d[mask_intersect_idx[secant_pts]]
+            z_pred_secant = self.secant(sdf_low, sdf_high, z_low, z_high, cam_loc_secant, ray_directions_secant,
+                                        sdf_network, conditional_volume)
+
+            # Get points
+            sampler_pts[mask_intersect_idx[secant_pts]] = cam_loc_secant + z_pred_secant[:,
+                                                                           None] * ray_directions_secant
+            sampler_dists[mask_intersect_idx[secant_pts]] = z_pred_secant
+
+        return sampler_pts, sampler_net_obj_mask, sampler_dists
+
+    def secant(self, sdf_low, sdf_high, z_low, z_high, rays_o, rays_d, sdf_network, conditional_volume):
+        ''' Runs the secant method for interval [z_low, z_high] for n_secant_steps '''
+
+        if self.network_inference:
+            sdf_func = sdf_network.sdf
+        else:
+            sdf_func = sdf_network.sdf_from_sdfvolume
+
+        z_pred = -sdf_low * (z_high - z_low) / (sdf_high - sdf_low) + z_low
+        for i in range(self.n_secant_steps):
+            p_mid = rays_o + z_pred[:, None] * rays_d
+            sdf_mid = sdf_func(p_mid,
+                               conditional_volume, lod=0, gru_fusion=False)['sdf_pts_scale%d' % 0].reshape(-1)
+            ind_low = (sdf_mid > 0).reshape(-1)
+            if ind_low.sum() > 0:
+                z_low[ind_low] = z_pred[ind_low]
+                sdf_low[ind_low] = sdf_mid[ind_low]
+            ind_high = sdf_mid < 0
+            if ind_high.sum() > 0:
+                z_high[ind_high] = z_pred[ind_high]
+                sdf_high[ind_high] = sdf_mid[ind_high]
+
+            z_pred = - sdf_low * (z_high - z_low) / (sdf_high - sdf_low) + z_low
+
+        return z_pred  # 1D tensor
+
+    def minimal_sdf_points(self, num_pixels, sdf, cam_loc, ray_directions, mask, min_dis, max_dis):
+        ''' Find points with minimal SDF value on rays for P_out pixels '''
+        device = sdf.device
+        n_mask_points = mask.sum()
+
+        n = self.n_steps
+        # steps = torch.linspace(0.0, 1.0,n).to(device)
+        steps = torch.empty(n).uniform_(0.0, 1.0).to(device)
+        mask_max_dis = max_dis[mask].unsqueeze(-1)
+        mask_min_dis = min_dis[mask].unsqueeze(-1)
+        steps = steps.unsqueeze(0).repeat(n_mask_points, 1) * (mask_max_dis - mask_min_dis) + mask_min_dis
+
+        mask_points = cam_loc.unsqueeze(1).repeat(1, num_pixels, 1).reshape(-1, 3)[mask]
+        mask_rays = ray_directions[mask, :]
+
+        mask_points_all = mask_points.unsqueeze(1).repeat(1, n, 1) + steps.unsqueeze(-1) * mask_rays.unsqueeze(
+            1).repeat(1, n, 1)
+        points = mask_points_all.reshape(-1, 3)
+
+        mask_sdf_all = []
+        for pnts in torch.split(points, 100000, dim=0):
+            mask_sdf_all.append(sdf(pnts))
+
+        mask_sdf_all = torch.cat(mask_sdf_all).reshape(-1, n)
+        min_vals, min_idx = mask_sdf_all.min(-1)
+        min_mask_points = mask_points_all.reshape(-1, n, 3)[torch.arange(0, n_mask_points), min_idx]
+        min_mask_dist = steps.reshape(-1, n)[torch.arange(0, n_mask_points), min_idx]
+
+        return min_mask_points, min_mask_dist

SparseNeuS_demo_v1/models/featurenet.py

@@ -0,0 +1,91 @@
+import torch
+
+# ! amazing!!!! autograd.grad with set_detect_anomaly(True) will cause memory leak
+# ! https://github.com/pytorch/pytorch/issues/51349
+# torch.autograd.set_detect_anomaly(True)
+import torch.nn as nn
+import torch.nn.functional as F
+from inplace_abn import InPlaceABN
+
+
+#############################################     MVS Net models        ################################################
+class ConvBnReLU(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=3, stride=1, pad=1,
+                 norm_act=InPlaceABN):
+        super(ConvBnReLU, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels,
+                              kernel_size, stride=stride, padding=pad, bias=False)
+        self.bn = norm_act(out_channels)
+
+    def forward(self, x):
+        return self.bn(self.conv(x))
+
+
+class ConvBnReLU3D(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=3, stride=1, pad=1,
+                 norm_act=InPlaceABN):
+        super(ConvBnReLU3D, self).__init__()
+        self.conv = nn.Conv3d(in_channels, out_channels,
+                              kernel_size, stride=stride, padding=pad, bias=False)
+        self.bn = norm_act(out_channels)
+        # self.bn = nn.ReLU()
+
+    def forward(self, x):
+        return self.bn(self.conv(x))
+
+
+###################################  feature net  ######################################
+class FeatureNet(nn.Module):
+    """
+    output 3 levels of features using a FPN structure
+    """
+
+    def __init__(self, norm_act=InPlaceABN):
+        super(FeatureNet, self).__init__()
+
+        self.conv0 = nn.Sequential(
+            ConvBnReLU(3, 8, 3, 1, 1, norm_act=norm_act),
+            ConvBnReLU(8, 8, 3, 1, 1, norm_act=norm_act))
+
+        self.conv1 = nn.Sequential(
+            ConvBnReLU(8, 16, 5, 2, 2, norm_act=norm_act),
+            ConvBnReLU(16, 16, 3, 1, 1, norm_act=norm_act),
+            ConvBnReLU(16, 16, 3, 1, 1, norm_act=norm_act))
+
+        self.conv2 = nn.Sequential(
+            ConvBnReLU(16, 32, 5, 2, 2, norm_act=norm_act),
+            ConvBnReLU(32, 32, 3, 1, 1, norm_act=norm_act),
+            ConvBnReLU(32, 32, 3, 1, 1, norm_act=norm_act))
+
+        self.toplayer = nn.Conv2d(32, 32, 1)
+        self.lat1 = nn.Conv2d(16, 32, 1)
+        self.lat0 = nn.Conv2d(8, 32, 1)
+
+        # to reduce channel size of the outputs from FPN
+        self.smooth1 = nn.Conv2d(32, 16, 3, padding=1)
+        self.smooth0 = nn.Conv2d(32, 8, 3, padding=1)
+
+    def _upsample_add(self, x, y):
+        return F.interpolate(x, scale_factor=2,
+                             mode="bilinear", align_corners=True) + y
+
+    def forward(self, x):
+        # x: (B, 3, H, W)
+        conv0 = self.conv0(x)  # (B, 8, H, W)
+        conv1 = self.conv1(conv0)  # (B, 16, H//2, W//2)
+        conv2 = self.conv2(conv1)  # (B, 32, H//4, W//4)
+        feat2 = self.toplayer(conv2)  # (B, 32, H//4, W//4)
+        feat1 = self._upsample_add(feat2, self.lat1(conv1))  # (B, 32, H//2, W//2)
+        feat0 = self._upsample_add(feat1, self.lat0(conv0))  # (B, 32, H, W)
+
+        # reduce output channels
+        feat1 = self.smooth1(feat1)  # (B, 16, H//2, W//2)
+        feat0 = self.smooth0(feat0)  # (B, 8, H, W)
+
+        # feats = {"level_0": feat0,
+        #          "level_1": feat1,
+        #          "level_2": feat2}
+
+        return [feat2, feat1, feat0]  # coarser to finer features

SparseNeuS_demo_v1/models/fields.py

@@ -0,0 +1,333 @@
+# The codes are from NeuS
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from models.embedder import get_embedder
+
+
+class SDFNetwork(nn.Module):
+    def __init__(self,
+                 d_in,
+                 d_out,
+                 d_hidden,
+                 n_layers,
+                 skip_in=(4,),
+                 multires=0,
+                 bias=0.5,
+                 scale=1,
+                 geometric_init=True,
+                 weight_norm=True,
+                 activation='softplus',
+                 conditional_type='multiply'):
+        super(SDFNetwork, self).__init__()
+
+        dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out]
+
+        self.embed_fn_fine = None
+
+        if multires > 0:
+            embed_fn, input_ch = get_embedder(multires, input_dims=d_in, normalize=False)
+            self.embed_fn_fine = embed_fn
+            dims[0] = input_ch
+
+        self.num_layers = len(dims)
+        self.skip_in = skip_in
+        self.scale = scale
+
+        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]
+
+            lin = nn.Linear(dims[l], out_dim)
+
+            if geometric_init:
+                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, -bias)
+                elif 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 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)  # ? why dims[0] - 3
+                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 weight_norm:
+                lin = nn.utils.weight_norm(lin)
+
+            setattr(self, "lin" + str(l), lin)
+
+        if activation == 'softplus':
+            self.activation = nn.Softplus(beta=100)
+        else:
+            assert activation == 'relu'
+            self.activation = nn.ReLU()
+
+    def forward(self, inputs):
+        inputs = inputs * self.scale
+        if self.embed_fn_fine is not None:
+            inputs = self.embed_fn_fine(inputs)
+
+        x = inputs
+        for l in range(0, self.num_layers - 1):
+            lin = getattr(self, "lin" + str(l))
+
+            if l in self.skip_in:
+                x = torch.cat([x, inputs], 1) / np.sqrt(2)
+
+            x = lin(x)
+
+            if l < self.num_layers - 2:
+                x = self.activation(x)
+        return torch.cat([x[:, :1] / self.scale, x[:, 1:]], dim=-1)
+
+    def sdf(self, x):
+        return self.forward(x)[:, :1]
+
+    def sdf_hidden_appearance(self, x):
+        return self.forward(x)
+
+    def gradient(self, x):
+        x.requires_grad_(True)
+        y = self.sdf(x)
+        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 VarianceNetwork(nn.Module):
+    def __init__(self, d_in, d_out, d_hidden, n_layers, skip_in=(4,), multires=0):
+        super(VarianceNetwork, self).__init__()
+
+        dims = [d_in] + [d_hidden for _ in range(n_layers)] + [d_out]
+
+        self.embed_fn_fine = None
+
+        if multires > 0:
+            embed_fn, input_ch = get_embedder(multires, normalize=False)
+            self.embed_fn_fine = embed_fn
+            dims[0] = input_ch
+
+        self.num_layers = len(dims)
+        self.skip_in = skip_in
+
+        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]
+
+            lin = nn.Linear(dims[l], out_dim)
+            setattr(self, "lin" + str(l), lin)
+
+        self.relu = nn.ReLU()
+        self.softplus = nn.Softplus(beta=100)
+
+    def forward(self, inputs):
+        if self.embed_fn_fine is not None:
+            inputs = self.embed_fn_fine(inputs)
+
+        x = inputs
+        for l in range(0, self.num_layers - 1):
+            lin = getattr(self, "lin" + str(l))
+
+            if l in self.skip_in:
+                x = torch.cat([x, inputs], 1) / np.sqrt(2)
+
+            x = lin(x)
+
+            if l < self.num_layers - 2:
+                x = self.relu(x)
+
+        # return torch.exp(x)
+        return 1.0 / (self.softplus(x + 0.5) + 1e-3)
+
+    def coarse(self, inputs):
+        return self.forward(inputs)[:, :1]
+
+    def fine(self, inputs):
+        return self.forward(inputs)[:, 1:]
+
+
+class FixVarianceNetwork(nn.Module):
+    def __init__(self, base):
+        super(FixVarianceNetwork, self).__init__()
+        self.base = base
+        self.iter_step = 0
+
+    def set_iter_step(self, iter_step):
+        self.iter_step = iter_step
+
+    def forward(self, x):
+        return torch.ones([len(x), 1]) * np.exp(-self.iter_step / self.base)
+
+
+class SingleVarianceNetwork(nn.Module):
+    def __init__(self, init_val=1.0):
+        super(SingleVarianceNetwork, self).__init__()
+        self.register_parameter('variance', nn.Parameter(torch.tensor(init_val)))
+
+    def forward(self, x):
+        return torch.ones([len(x), 1]).to(x.device) * torch.exp(self.variance * 10.0)
+
+
+
+class RenderingNetwork(nn.Module):
+    def __init__(
+            self,
+            d_feature,
+            mode,
+            d_in,
+            d_out,
+            d_hidden,
+            n_layers,
+            weight_norm=True,
+            multires_view=0,
+            squeeze_out=True,
+            d_conditional_colors=0
+    ):
+        super().__init__()
+
+        self.mode = mode
+        self.squeeze_out = squeeze_out
+        dims = [d_in + d_feature] + [d_hidden for _ in range(n_layers)] + [d_out]
+
+        self.embedview_fn = None
+        if multires_view > 0:
+            embedview_fn, input_ch = get_embedder(multires_view)
+            self.embedview_fn = embedview_fn
+            dims[0] += (input_ch - 3)
+
+        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 weight_norm:
+                lin = nn.utils.weight_norm(lin)
+
+            setattr(self, "lin" + str(l), lin)
+
+        self.relu = nn.ReLU()
+
+    def forward(self, points, normals, view_dirs, feature_vectors):
+        if self.embedview_fn is not None:
+            view_dirs = self.embedview_fn(view_dirs)
+
+        rendering_input = None
+
+        if self.mode == 'idr':
+            rendering_input = torch.cat([points, view_dirs, normals, feature_vectors], dim=-1)
+        elif self.mode == 'no_view_dir':
+            rendering_input = torch.cat([points, normals, feature_vectors], dim=-1)
+        elif self.mode == 'no_normal':
+            rendering_input = torch.cat([points, view_dirs, feature_vectors], dim=-1)
+        elif self.mode == 'no_points':
+            rendering_input = torch.cat([view_dirs, normals, feature_vectors], dim=-1)
+        elif self.mode == 'no_points_no_view_dir':
+            rendering_input = torch.cat([normals, feature_vectors], dim=-1)
+
+        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)
+
+        if self.squeeze_out:
+            x = torch.sigmoid(x)
+        return x
+
+
+# Code from nerf-pytorch
+class NeRF(nn.Module):
+    def __init__(self, D=8, W=256, d_in=3, d_in_view=3, multires=0, multires_view=0, output_ch=4, skips=[4],
+                 use_viewdirs=False):
+        """
+        """
+        super(NeRF, self).__init__()
+        self.D = D
+        self.W = W
+        self.d_in = d_in
+        self.d_in_view = d_in_view
+        self.input_ch = 3
+        self.input_ch_view = 3
+        self.embed_fn = None
+        self.embed_fn_view = None
+
+        if multires > 0:
+            embed_fn, input_ch = get_embedder(multires, input_dims=d_in, normalize=False)
+            self.embed_fn = embed_fn
+            self.input_ch = input_ch
+
+        if multires_view > 0:
+            embed_fn_view, input_ch_view = get_embedder(multires_view, input_dims=d_in_view, normalize=False)
+            self.embed_fn_view = embed_fn_view
+            self.input_ch_view = input_ch_view
+
+        self.skips = skips
+        self.use_viewdirs = use_viewdirs
+
+        self.pts_linears = nn.ModuleList(
+            [nn.Linear(self.input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W)
+                                             for i in
+                                             range(D - 1)])
+
+        ### Implementation according to the official code release (https://github.com/bmild/nerf/blob/master/run_nerf_helpers.py#L104-L105)
+        self.views_linears = nn.ModuleList([nn.Linear(self.input_ch_view + W, W // 2)])
+
+        ### Implementation according to the paper
+        # self.views_linears = nn.ModuleList(
+        #     [nn.Linear(input_ch_views + W, W//2)] + [nn.Linear(W//2, W//2) for i in range(D//2)])
+
+        if use_viewdirs:
+            self.feature_linear = nn.Linear(W, W)
+            self.alpha_linear = nn.Linear(W, 1)
+            self.rgb_linear = nn.Linear(W // 2, 3)
+        else:
+            self.output_linear = nn.Linear(W, output_ch)
+
+    def forward(self, input_pts, input_views):
+        if self.embed_fn is not None:
+            input_pts = self.embed_fn(input_pts)
+        if self.embed_fn_view is not None:
+            input_views = self.embed_fn_view(input_views)
+
+        h = input_pts
+        for i, l in enumerate(self.pts_linears):
+            h = self.pts_linears[i](h)
+            h = F.relu(h)
+            if i in self.skips:
+                h = torch.cat([input_pts, h], -1)
+
+        if self.use_viewdirs:
+            alpha = self.alpha_linear(h)
+            feature = self.feature_linear(h)
+            h = torch.cat([feature, input_views], -1)
+
+            for i, l in enumerate(self.views_linears):
+                h = self.views_linears[i](h)
+                h = F.relu(h)
+
+            rgb = self.rgb_linear(h)
+            return alpha + 1.0, rgb
+        else:
+            assert False

SparseNeuS_demo_v1/models/patch_projector.py

@@ -0,0 +1,211 @@
+"""
+Patch Projector
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from models.render_utils import sample_ptsFeatures_from_featureMaps
+
+
+class PatchProjector():
+    def __init__(self, patch_size):
+        self.h_patch_size = patch_size
+        self.offsets = build_patch_offset(patch_size)  # the warping patch offsets index
+
+        self.z_axis = torch.tensor([0, 0, 1]).float()
+
+        self.plane_dist_thresh = 0.001
+
+    # * correctness checked
+    def pixel_warp(self, pts, imgs, intrinsics,
+                   w2cs, img_wh=None):
+        """
+
+        :param pts: [N_rays, n_samples, 3]
+        :param imgs: [N_views, 3, H, W]
+        :param intrinsics: [N_views, 4, 4]
+        :param c2ws: [N_views, 4, 4]
+        :param img_wh:
+        :return:
+        """
+        if img_wh is None:
+            N_views, _, sizeH, sizeW = imgs.shape
+            img_wh = [sizeW, sizeH]
+
+        pts_color, valid_mask = sample_ptsFeatures_from_featureMaps(
+            pts, imgs, w2cs, intrinsics, img_wh,
+            proj_matrix=None, return_mask=True)  # [N_views, c, N_rays, n_samples], [N_views, N_rays, n_samples]
+
+        pts_color = pts_color.permute(2, 3, 0, 1)
+        valid_mask = valid_mask.permute(1, 2, 0)
+
+        return pts_color, valid_mask  # [N_rays, n_samples, N_views,  3] , [N_rays, n_samples, N_views]
+
+    def patch_warp(self, pts, uv, normals, src_imgs,
+                   ref_intrinsic, src_intrinsics,
+                   ref_c2w, src_c2ws, img_wh=None
+                   ):
+        """
+
+        :param pts: [N_rays, n_samples, 3]
+        :param uv : [N_rays, 2]  normalized in (-1, 1)
+        :param normals: [N_rays, n_samples, 3]  The normal of pt in world space
+        :param src_imgs: [N_src, 3, h, w]
+        :param ref_intrinsic: [4,4]
+        :param src_intrinsics: [N_src, 4, 4]
+        :param ref_c2w: [4,4]
+        :param src_c2ws: [N_src, 4, 4]
+        :return:
+        """
+        device = pts.device
+
+        N_rays, n_samples, _ = pts.shape
+        N_pts = N_rays * n_samples
+
+        N_src, _, sizeH, sizeW = src_imgs.shape
+
+        if img_wh is not None:
+            sizeW, sizeH = img_wh[0], img_wh[1]
+
+        # scale uv from (-1, 1) to (0, W/H)
+        uv[:, 0] = (uv[:, 0] + 1) / 2. * (sizeW - 1)
+        uv[:, 1] = (uv[:, 1] + 1) / 2. * (sizeH - 1)
+
+        ref_intr = ref_intrinsic[:3, :3]
+        inv_ref_intr = torch.inverse(ref_intr)
+        src_intrs = src_intrinsics[:, :3, :3]
+        inv_src_intrs = torch.inverse(src_intrs)
+
+        ref_pose = ref_c2w
+        inv_ref_pose = torch.inverse(ref_pose)
+        src_poses = src_c2ws
+        inv_src_poses = torch.inverse(src_poses)
+
+        ref_cam_loc = ref_pose[:3, 3].unsqueeze(0)  # [1, 3]
+        sampled_dists = torch.norm(pts - ref_cam_loc, dim=-1)  # [N_pts, 1]
+
+        relative_proj = inv_src_poses @ ref_pose
+        R_rel = relative_proj[:, :3, :3]
+        t_rel = relative_proj[:, :3, 3:]
+        R_ref = inv_ref_pose[:3, :3]
+        t_ref = inv_ref_pose[:3, 3:]
+
+        pts = pts.view(-1, 3)
+        normals = normals.view(-1, 3)
+
+        with torch.no_grad():
+            rot_normals = R_ref @ normals.unsqueeze(-1)  # [N_pts, 3, 1]
+            points_in_ref = R_ref @ pts.unsqueeze(
+                -1) + t_ref  # [N_pts, 3, 1]  points in the reference frame coordiantes system
+            d1 = torch.sum(rot_normals * points_in_ref, dim=1).unsqueeze(
+                1)  # distance from the plane to ref camera center
+
+            d2 = torch.sum(rot_normals.unsqueeze(1) * (-R_rel.transpose(1, 2) @ t_rel).unsqueeze(0),
+                           dim=2)  # distance from the plane to src camera center
+            valid_hom = (torch.abs(d1) > self.plane_dist_thresh) & (
+                    torch.abs(d1 - d2) > self.plane_dist_thresh) & ((d2 / d1) < 1)
+
+            d1 = d1.squeeze()
+            sign = torch.sign(d1)
+            sign[sign == 0] = 1
+            d = torch.clamp(torch.abs(d1), 1e-8) * sign
+
+            H = src_intrs.unsqueeze(1) @ (
+                    R_rel.unsqueeze(1) + t_rel.unsqueeze(1) @ rot_normals.view(1, N_pts, 1, 3) / d.view(1,
+                                                                                                        N_pts,
+                                                                                                        1, 1)
+            ) @ inv_ref_intr.view(1, 1, 3, 3)
+
+            # replace invalid homs with fronto-parallel homographies
+            H_invalid = src_intrs.unsqueeze(1) @ (
+                    R_rel.unsqueeze(1) + t_rel.unsqueeze(1) @ self.z_axis.to(device).view(1, 1, 1, 3).expand(-1, N_pts,
+                                                                                                             -1,
+                                                                                                             -1) / sampled_dists.view(
+                1, N_pts, 1, 1)
+            ) @ inv_ref_intr.view(1, 1, 3, 3)
+            tmp_m = ~valid_hom.view(-1, N_src).t()
+            H[tmp_m] = H_invalid[tmp_m]
+
+        pixels = uv.view(N_rays, 1, 2) + self.offsets.float().to(device)
+        Npx = pixels.shape[1]
+        grid, warp_mask_full = self.patch_homography(H, pixels)
+
+        warp_mask_full = warp_mask_full & (grid[..., 0] < (sizeW - self.h_patch_size)) & (
+                grid[..., 1] < (sizeH - self.h_patch_size)) & (grid >= self.h_patch_size).all(dim=-1)
+        warp_mask_full = warp_mask_full.view(N_src, N_rays, n_samples, Npx)
+
+        grid = torch.clamp(normalize(grid, sizeH, sizeW), -10, 10)
+
+        sampled_rgb_val = F.grid_sample(src_imgs, grid.view(N_src, -1, 1, 2), align_corners=True).squeeze(
+            -1).transpose(1, 2)
+        sampled_rgb_val = sampled_rgb_val.view(N_src, N_rays, n_samples, Npx, 3)
+
+        warp_mask_full = warp_mask_full.permute(1, 2, 0, 3).contiguous()  # (N_rays, n_samples, N_src, Npx)
+        sampled_rgb_val = sampled_rgb_val.permute(1, 2, 0, 3, 4).contiguous()  # (N_rays, n_samples, N_src, Npx, 3)
+
+        return sampled_rgb_val, warp_mask_full
+
+    def patch_homography(self, H, uv):
+        N, Npx = uv.shape[:2]
+        Nsrc = H.shape[0]
+        H = H.view(Nsrc, N, -1, 3, 3)
+        hom_uv = add_hom(uv)
+
+        # einsum is 30 times faster
+        # tmp = (H.view(Nsrc, N, -1, 1, 3, 3) @ hom_uv.view(1, N, 1, -1, 3, 1)).squeeze(-1).view(Nsrc, -1, 3)
+        tmp = torch.einsum("vprik,pok->vproi", H, hom_uv).reshape(Nsrc, -1, 3)
+
+        grid = tmp[..., :2] / torch.clamp(tmp[..., 2:], 1e-8)
+        mask = tmp[..., 2] > 0
+        return grid, mask
+
+
+def add_hom(pts):
+    try:
+        dev = pts.device
+        ones = torch.ones(pts.shape[:-1], device=dev).unsqueeze(-1)
+        return torch.cat((pts, ones), dim=-1)
+
+    except AttributeError:
+        ones = np.ones((pts.shape[0], 1))
+        return np.concatenate((pts, ones), axis=1)
+
+
+def normalize(flow, h, w, clamp=None):
+    # either h and w are simple float or N torch.tensor where N batch size
+    try:
+        h.device
+
+    except AttributeError:
+        h = torch.tensor(h, device=flow.device).float().unsqueeze(0)
+        w = torch.tensor(w, device=flow.device).float().unsqueeze(0)
+
+    if len(flow.shape) == 4:
+        w = w.unsqueeze(1).unsqueeze(2)
+        h = h.unsqueeze(1).unsqueeze(2)
+    elif len(flow.shape) == 3:
+        w = w.unsqueeze(1)
+        h = h.unsqueeze(1)
+    elif len(flow.shape) == 5:
+        w = w.unsqueeze(0).unsqueeze(2).unsqueeze(2)
+        h = h.unsqueeze(0).unsqueeze(2).unsqueeze(2)
+
+    res = torch.empty_like(flow)
+    if res.shape[-1] == 3:
+        res[..., 2] = 1
+
+    # for grid_sample with align_corners=True
+    # https://github.com/pytorch/pytorch/blob/c371542efc31b1abfe6f388042aa3ab0cef935f2/aten/src/ATen/native/GridSampler.h#L33
+    res[..., 0] = 2 * flow[..., 0] / (w - 1) - 1
+    res[..., 1] = 2 * flow[..., 1] / (h - 1) - 1
+
+    if clamp:
+        return torch.clamp(res, -clamp, clamp)
+    else:
+        return res
+
+
+def build_patch_offset(h_patch_size):
+    offsets = torch.arange(-h_patch_size, h_patch_size + 1)
+    return torch.stack(torch.meshgrid(offsets, offsets, indexing="ij")[::-1], dim=-1).view(1, -1, 2)  # nb_pixels_patch * 2

SparseNeuS_demo_v1/models/projector.py

@@ -0,0 +1,425 @@
+# The codes are partly from IBRNet
+
+import torch
+import torch.nn.functional as F
+from models.render_utils import sample_ptsFeatures_from_featureMaps, sample_ptsFeatures_from_featureVolume
+
+def safe_l2_normalize(x, dim=None, eps=1e-6):
+    return F.normalize(x, p=2, dim=dim, eps=eps)
+
+class Projector():
+    """
+    Obtain features from geometryVolume and rendering_feature_maps for generalized rendering
+    """
+
+    def compute_angle(self, xyz, query_c2w, supporting_c2ws):
+        """
+
+        :param xyz: [N_rays, n_samples,3 ]
+        :param query_c2w: [1,4,4]
+        :param supporting_c2ws: [n,4,4]
+        :return:
+        """
+        N_rays, n_samples, _ = xyz.shape
+        num_views = supporting_c2ws.shape[0]
+        xyz = xyz.reshape(-1, 3)
+
+        ray2tar_pose = (query_c2w[:, :3, 3].unsqueeze(1) - xyz.unsqueeze(0))
+        ray2tar_pose /= (torch.norm(ray2tar_pose, dim=-1, keepdim=True) + 1e-6)
+        ray2support_pose = (supporting_c2ws[:, :3, 3].unsqueeze(1) - xyz.unsqueeze(0))
+        ray2support_pose /= (torch.norm(ray2support_pose, dim=-1, keepdim=True) + 1e-6)
+        ray_diff = ray2tar_pose - ray2support_pose
+        ray_diff_norm = torch.norm(ray_diff, dim=-1, keepdim=True)
+        ray_diff_dot = torch.sum(ray2tar_pose * ray2support_pose, dim=-1, keepdim=True)
+        ray_diff_direction = ray_diff / torch.clamp(ray_diff_norm, min=1e-6)
+        ray_diff = torch.cat([ray_diff_direction, ray_diff_dot], dim=-1)
+        ray_diff = ray_diff.reshape((num_views, N_rays, n_samples, 4))  # the last dimension (4) is dot-product
+        return ray_diff.detach()
+
+
+    def compute_angle_view_independent(self, xyz, surface_normals, supporting_c2ws):
+        """
+
+        :param xyz: [N_rays, n_samples,3 ]
+        :param surface_normals: [N_rays, n_samples,3 ]
+        :param supporting_c2ws: [n,4,4]
+        :return:
+        """
+        N_rays, n_samples, _ = xyz.shape
+        num_views = supporting_c2ws.shape[0]
+        xyz = xyz.reshape(-1, 3)
+
+        ray2tar_pose = surface_normals
+        ray2support_pose = (supporting_c2ws[:, :3, 3].unsqueeze(1) - xyz.unsqueeze(0))
+        ray2support_pose /= (torch.norm(ray2support_pose, dim=-1, keepdim=True) + 1e-6)
+        ray_diff = ray2tar_pose - ray2support_pose
+        ray_diff_norm = torch.norm(ray_diff, dim=-1, keepdim=True)
+        ray_diff_dot = torch.sum(ray2tar_pose * ray2support_pose, dim=-1, keepdim=True)
+        ray_diff_direction = ray_diff / torch.clamp(ray_diff_norm, min=1e-6)
+        ray_diff = torch.cat([ray_diff_direction, ray_diff_dot], dim=-1)
+        ray_diff = ray_diff.reshape((num_views, N_rays, n_samples, 4))  # the last dimension (4) is dot-product, 
+                                                                        # and the first three dimensions is the normalized ray diff vector
+        return ray_diff.detach()
+
+    @torch.no_grad()
+    def compute_z_diff(self, xyz, w2cs, intrinsics, pred_depth_values):
+        """
+        compute the depth difference of query pts projected on the image and the predicted depth values of the image
+        :param xyz:  [N_rays, n_samples,3 ]
+        :param w2cs:  [N_views, 4, 4]
+        :param intrinsics: [N_views, 3, 3]
+        :param pred_depth_values: [N_views, N_rays, n_samples,1 ]
+        :param pred_depth_masks: [N_views, N_rays, n_samples]
+        :return:
+        """
+        device = xyz.device
+        N_views = w2cs.shape[0]
+        N_rays, n_samples, _ = xyz.shape
+        proj_matrix = torch.matmul(intrinsics, w2cs[:, :3, :])
+
+        proj_rot = proj_matrix[:, :3, :3]
+        proj_trans = proj_matrix[:, :3, 3:]
+
+        batch_xyz = xyz.permute(2, 0, 1).contiguous().view(1, 3, N_rays * n_samples).repeat(N_views, 1, 1)
+
+        proj_xyz = proj_rot.bmm(batch_xyz) + proj_trans
+
+        # X = proj_xyz[:, 0]
+        # Y = proj_xyz[:, 1]
+        Z = proj_xyz[:, 2].clamp(min=1e-3)  # [N_views, N_rays*n_samples]
+        proj_z = Z.view(N_views, N_rays, n_samples, 1)
+
+        z_diff = proj_z - pred_depth_values  # [N_views, N_rays, n_samples,1 ]
+
+        return z_diff
+
+    def compute(self,
+                pts,
+                # * 3d geometry feature volumes
+                geometryVolume=None,
+                geometryVolumeMask=None,
+                vol_dims=None,
+                partial_vol_origin=None,
+                vol_size=None,
+                # * 2d rendering feature maps
+                rendering_feature_maps=None,
+                color_maps=None,
+                w2cs=None,
+                intrinsics=None,
+                img_wh=None,
+                query_img_idx=0,  # the index of the N_views dim for rendering
+                query_c2w=None,
+                pred_depth_maps=None,   # no use here
+                pred_depth_masks=None   # no use here
+                ):
+        """
+        extract features of pts for rendering
+        :param pts:
+        :param geometryVolume:
+        :param vol_dims:
+        :param partial_vol_origin:
+        :param vol_size:
+        :param rendering_feature_maps:
+        :param color_maps:
+        :param w2cs:
+        :param intrinsics:
+        :param img_wh:
+        :param rendering_img_idx: by default, we render the first view of w2cs
+        :return:
+        """
+        device = pts.device
+        c2ws = torch.inverse(w2cs)
+
+        if len(pts.shape) == 2:
+            pts = pts[None, :, :]
+
+        N_rays, n_samples, _ = pts.shape
+        N_views = rendering_feature_maps.shape[0]  # shape (N_views, C, H, W)
+
+        supporting_img_idxs = torch.LongTensor([x for x in range(N_views) if x != query_img_idx]).to(device)
+        query_img_idx = torch.LongTensor([query_img_idx]).to(device)
+
+        if query_c2w is None and query_img_idx > -1:
+            query_c2w = torch.index_select(c2ws, 0, query_img_idx)
+            supporting_c2ws = torch.index_select(c2ws, 0, supporting_img_idxs)
+            supporting_w2cs = torch.index_select(w2cs, 0, supporting_img_idxs)
+            supporting_rendering_feature_maps = torch.index_select(rendering_feature_maps, 0, supporting_img_idxs)
+            supporting_color_maps = torch.index_select(color_maps, 0, supporting_img_idxs)
+            supporting_intrinsics = torch.index_select(intrinsics, 0, supporting_img_idxs)
+
+            if pred_depth_maps is not None:
+                supporting_depth_maps = torch.index_select(pred_depth_maps, 0, supporting_img_idxs)
+                supporting_depth_masks = torch.index_select(pred_depth_masks, 0, supporting_img_idxs)
+            # print("N_supporting_views: ", N_views - 1)
+            N_supporting_views = N_views - 1
+        else:
+            supporting_c2ws = c2ws
+            supporting_w2cs = w2cs
+            supporting_rendering_feature_maps = rendering_feature_maps
+            supporting_color_maps = color_maps
+            supporting_intrinsics = intrinsics
+            supporting_depth_maps = pred_depth_masks
+            supporting_depth_masks = pred_depth_masks
+            # print("N_supporting_views: ", N_views)
+            N_supporting_views = N_views
+        # import ipdb; ipdb.set_trace()
+        if geometryVolume is not None:
+            # * sample feature of pts from 3D feature volume
+            pts_geometry_feature, pts_geometry_masks_0 = sample_ptsFeatures_from_featureVolume(
+                pts, geometryVolume, vol_dims,
+                partial_vol_origin, vol_size)  # [N_rays, n_samples, C], [N_rays, n_samples]
+
+            if len(geometryVolumeMask.shape) == 3:
+                geometryVolumeMask = geometryVolumeMask[None, :, :, :]
+
+            pts_geometry_masks_1, _ = sample_ptsFeatures_from_featureVolume(
+                pts, geometryVolumeMask.to(geometryVolume.dtype), vol_dims,
+                partial_vol_origin, vol_size)  # [N_rays, n_samples, C]
+
+            pts_geometry_masks = pts_geometry_masks_0 & (pts_geometry_masks_1[..., 0] > 0)
+        else:
+            pts_geometry_feature = None
+            pts_geometry_masks = None
+
+        # * sample feature of pts from 2D feature maps
+        pts_rendering_feats, pts_rendering_mask = sample_ptsFeatures_from_featureMaps(
+            pts, supporting_rendering_feature_maps, supporting_w2cs,
+            supporting_intrinsics, img_wh,
+            return_mask=True)  # [N_views, C, N_rays, n_samples], # [N_views, N_rays, n_samples]
+        # import ipdb; ipdb.set_trace()
+        # * size (N_views, N_rays*n_samples, c)
+        pts_rendering_feats = pts_rendering_feats.permute(0, 2, 3, 1).contiguous()
+
+        pts_rendering_colors = sample_ptsFeatures_from_featureMaps(pts, supporting_color_maps, supporting_w2cs,
+                                                                   supporting_intrinsics, img_wh)
+        # * size (N_views, N_rays*n_samples, c)
+        pts_rendering_colors = pts_rendering_colors.permute(0, 2, 3, 1).contiguous()
+
+        rgb_feats = torch.cat([pts_rendering_colors, pts_rendering_feats], dim=-1)  # [N_views, N_rays, n_samples, 3+c]
+
+
+        ray_diff = self.compute_angle(pts, query_c2w, supporting_c2ws)  # [N_views, N_rays, n_samples, 4]
+        # import ipdb; ipdb.set_trace()
+        if pts_geometry_masks is not None:
+            final_mask = pts_geometry_masks[None, :, :].repeat(N_supporting_views, 1, 1) & \
+                         pts_rendering_mask  # [N_views, N_rays, n_samples]
+        else:
+            final_mask = pts_rendering_mask
+        # import ipdb; ipdb.set_trace()
+        z_diff, pts_pred_depth_masks = None, None
+        
+        if pred_depth_maps is not None:
+            pts_pred_depth_values = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_maps, supporting_w2cs,
+                                                                        supporting_intrinsics, img_wh)
+            pts_pred_depth_values = pts_pred_depth_values.permute(0, 2, 3,
+                                                                  1).contiguous()  # (N_views, N_rays*n_samples, 1)
+
+            # - pts_pred_depth_masks are critical than final_mask,
+            # - the ray containing few invalid pts will be treated invalid
+            pts_pred_depth_masks = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_masks.float(),
+                                                                       supporting_w2cs,
+                                                                       supporting_intrinsics, img_wh)
+            
+            pts_pred_depth_masks = pts_pred_depth_masks.permute(0, 2, 3, 1).contiguous()[:, :, :,
+                                   0]  # (N_views, N_rays*n_samples)
+
+            z_diff = self.compute_z_diff(pts, supporting_w2cs, supporting_intrinsics, pts_pred_depth_values)
+        # import ipdb; ipdb.set_trace()
+        return pts_geometry_feature, rgb_feats, ray_diff, final_mask, z_diff, pts_pred_depth_masks
+
+
+    def compute_view_independent(   
+                                    self,
+                                    pts,
+                                    # * 3d geometry feature volumes
+                                    geometryVolume=None,
+                                    geometryVolumeMask=None,
+                                    sdf_network=None,
+                                    lod=0,
+                                    vol_dims=None,
+                                    partial_vol_origin=None,
+                                    vol_size=None,
+                                    # * 2d rendering feature maps
+                                    rendering_feature_maps=None,
+                                    color_maps=None,
+                                    w2cs=None,
+                                    target_candidate_w2cs=None,
+                                    intrinsics=None,
+                                    img_wh=None,
+                                    query_img_idx=0,  # the index of the N_views dim for rendering
+                                    query_c2w=None,
+                                    pred_depth_maps=None,   # no use here
+                                    pred_depth_masks=None   # no use here
+                                ):
+        """
+        extract features of pts for rendering
+        :param pts:
+        :param geometryVolume:
+        :param vol_dims:
+        :param partial_vol_origin:
+        :param vol_size:
+        :param rendering_feature_maps:
+        :param color_maps:
+        :param w2cs:
+        :param intrinsics:
+        :param img_wh:
+        :param rendering_img_idx: by default, we render the first view of w2cs
+        :return:
+        """
+        device = pts.device
+        c2ws = torch.inverse(w2cs)
+
+        if len(pts.shape) == 2:
+            pts = pts[None, :, :]
+
+        N_rays, n_samples, _ = pts.shape
+        N_views = rendering_feature_maps.shape[0]  # shape (N_views, C, H, W)
+
+        supporting_img_idxs = torch.LongTensor([x for x in range(N_views) if x != query_img_idx]).to(device)
+        query_img_idx = torch.LongTensor([query_img_idx]).to(device)
+
+        if query_c2w is None and query_img_idx > -1:
+            query_c2w = torch.index_select(c2ws, 0, query_img_idx)
+            supporting_c2ws = torch.index_select(c2ws, 0, supporting_img_idxs)
+            supporting_w2cs = torch.index_select(w2cs, 0, supporting_img_idxs)
+            supporting_rendering_feature_maps = torch.index_select(rendering_feature_maps, 0, supporting_img_idxs)
+            supporting_color_maps = torch.index_select(color_maps, 0, supporting_img_idxs)
+            supporting_intrinsics = torch.index_select(intrinsics, 0, supporting_img_idxs)
+
+            if pred_depth_maps is not None:
+                supporting_depth_maps = torch.index_select(pred_depth_maps, 0, supporting_img_idxs)
+                supporting_depth_masks = torch.index_select(pred_depth_masks, 0, supporting_img_idxs)
+            # print("N_supporting_views: ", N_views - 1)
+            N_supporting_views = N_views - 1
+        else:
+            supporting_c2ws = c2ws
+            supporting_w2cs = w2cs
+            supporting_rendering_feature_maps = rendering_feature_maps
+            supporting_color_maps = color_maps
+            supporting_intrinsics = intrinsics
+            supporting_depth_maps = pred_depth_masks
+            supporting_depth_masks = pred_depth_masks
+            # print("N_supporting_views: ", N_views)
+            N_supporting_views = N_views
+        # import ipdb; ipdb.set_trace()
+        if geometryVolume is not None:
+            # * sample feature of pts from 3D feature volume
+            pts_geometry_feature, pts_geometry_masks_0 = sample_ptsFeatures_from_featureVolume(
+                pts, geometryVolume, vol_dims,
+                partial_vol_origin, vol_size)  # [N_rays, n_samples, C], [N_rays, n_samples]
+
+            if len(geometryVolumeMask.shape) == 3:
+                geometryVolumeMask = geometryVolumeMask[None, :, :, :]
+
+            pts_geometry_masks_1, _ = sample_ptsFeatures_from_featureVolume(
+                pts, geometryVolumeMask.to(geometryVolume.dtype), vol_dims,
+                partial_vol_origin, vol_size)  # [N_rays, n_samples, C]
+
+            pts_geometry_masks = pts_geometry_masks_0 & (pts_geometry_masks_1[..., 0] > 0)
+        else:
+            pts_geometry_feature = None
+            pts_geometry_masks = None
+
+        # * sample feature of pts from 2D feature maps
+        pts_rendering_feats, pts_rendering_mask = sample_ptsFeatures_from_featureMaps(
+            pts, supporting_rendering_feature_maps, supporting_w2cs,
+            supporting_intrinsics, img_wh,
+            return_mask=True)  # [N_views, C, N_rays, n_samples], # [N_views, N_rays, n_samples]
+
+        # * size (N_views, N_rays*n_samples, c)
+        pts_rendering_feats = pts_rendering_feats.permute(0, 2, 3, 1).contiguous()
+
+        pts_rendering_colors = sample_ptsFeatures_from_featureMaps(pts, supporting_color_maps, supporting_w2cs,
+                                                                   supporting_intrinsics, img_wh)
+        # * size (N_views, N_rays*n_samples, c)
+        pts_rendering_colors = pts_rendering_colors.permute(0, 2, 3, 1).contiguous()
+
+        rgb_feats = torch.cat([pts_rendering_colors, pts_rendering_feats], dim=-1)  # [N_views, N_rays, n_samples, 3+c]
+        
+        # import ipdb; ipdb.set_trace()
+        
+        gradients = sdf_network.gradient(
+                                        pts.reshape(-1, 3), # pts.squeeze(0),
+                                        geometryVolume.unsqueeze(0), 
+                                        lod=lod
+                                        ).squeeze()
+        
+        surface_normals = safe_l2_normalize(gradients, dim=-1) # [npts, 3]
+        # input normals
+        ren_ray_diff = self.compute_angle_view_independent(
+                         xyz=pts,
+                         surface_normals=surface_normals,
+                         supporting_c2ws=supporting_c2ws   
+        )
+
+        # # choose closest target view direction from 32 candidate views
+        # # choose the closest source view as view direction instead of the normals vectors
+        # pts2src_centers = safe_l2_normalize((supporting_c2ws[:, :3, 3].unsqueeze(1) - pts)) # [N_views, npts, 3]
+
+        # cosine_distance = torch.sum(pts2src_centers * surface_normals, dim=-1, keepdim=True) # [N_views, npts, 1]
+        # # choose the largest cosine distance as the view direction
+        # max_idx = torch.argmax(cosine_distance, dim=0) # [npts, 1]
+
+        # chosen_view_direction = pts2src_centers[max_idx.squeeze(), torch.arange(pts.shape[1]), :] # [npts, 3]
+        # ren_ray_diff = self.compute_angle_view_independent(
+        #                  xyz=pts,
+        #                  surface_normals=chosen_view_direction,
+        #                  supporting_c2ws=supporting_c2ws   
+        # )
+
+
+
+        # # choose closest target view direction from 8 candidate views
+        # # choose the closest source view as view direction instead of the normals vectors
+        # target_candidate_c2ws = torch.inverse(target_candidate_w2cs)
+        # pts2src_centers = safe_l2_normalize((target_candidate_c2ws[:, :3, 3].unsqueeze(1) - pts)) # [N_views, npts, 3]
+
+        # cosine_distance = torch.sum(pts2src_centers * surface_normals, dim=-1, keepdim=True) # [N_views, npts, 1]
+        # # choose the largest cosine distance as the view direction
+        # max_idx = torch.argmax(cosine_distance, dim=0) # [npts, 1]
+
+        # chosen_view_direction = pts2src_centers[max_idx.squeeze(), torch.arange(pts.shape[1]), :] # [npts, 3]
+        # ren_ray_diff = self.compute_angle_view_independent(
+        #                  xyz=pts,
+        #                  surface_normals=chosen_view_direction,
+        #                  supporting_c2ws=supporting_c2ws   
+        # )
+
+
+        # ray_diff = self.compute_angle(pts, query_c2w, supporting_c2ws)  # [N_views, N_rays, n_samples, 4]
+        # import ipdb; ipdb.set_trace()
+
+
+        # input_directions = safe_l2_normalize(pts)
+        # ren_ray_diff = self.compute_angle_view_independent(
+        #                  xyz=pts,
+        #                  surface_normals=input_directions,
+        #                  supporting_c2ws=supporting_c2ws   
+        # )
+
+        if pts_geometry_masks is not None:
+            final_mask = pts_geometry_masks[None, :, :].repeat(N_supporting_views, 1, 1) & \
+                         pts_rendering_mask  # [N_views, N_rays, n_samples]
+        else:
+            final_mask = pts_rendering_mask
+        # import ipdb; ipdb.set_trace()
+        z_diff, pts_pred_depth_masks = None, None
+        
+        if pred_depth_maps is not None:
+            pts_pred_depth_values = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_maps, supporting_w2cs,
+                                                                        supporting_intrinsics, img_wh)
+            pts_pred_depth_values = pts_pred_depth_values.permute(0, 2, 3,
+                                                                  1).contiguous()  # (N_views, N_rays*n_samples, 1)
+
+            # - pts_pred_depth_masks are critical than final_mask,
+            # - the ray containing few invalid pts will be treated invalid
+            pts_pred_depth_masks = sample_ptsFeatures_from_featureMaps(pts, supporting_depth_masks.float(),
+                                                                       supporting_w2cs,
+                                                                       supporting_intrinsics, img_wh)
+            
+            pts_pred_depth_masks = pts_pred_depth_masks.permute(0, 2, 3, 1).contiguous()[:, :, :,
+                                   0]  # (N_views, N_rays*n_samples)
+
+            z_diff = self.compute_z_diff(pts, supporting_w2cs, supporting_intrinsics, pts_pred_depth_values)
+        # import ipdb; ipdb.set_trace()
+        return pts_geometry_feature, rgb_feats, ren_ray_diff, final_mask, z_diff, pts_pred_depth_masks

SparseNeuS_demo_v1/models/rays.py

@@ -0,0 +1,320 @@
+import os, torch
+import numpy as np
+
+import torch.nn.functional as F
+
+def build_patch_offset(h_patch_size):
+    offsets = torch.arange(-h_patch_size, h_patch_size + 1)
+    return torch.stack(torch.meshgrid(offsets, offsets)[::-1], dim=-1).view(1, -1, 2)  # nb_pixels_patch * 2
+
+
+def gen_rays_from_single_image(H, W, image, intrinsic, c2w, depth=None, mask=None):
+    """
+    generate rays in world space, for image image
+    :param H:
+    :param W:
+    :param intrinsics: [3,3]
+    :param c2ws: [4,4]
+    :return:
+    """
+    device = image.device
+    ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
+                            torch.linspace(0, W - 1, W), indexing="ij")  # pytorch's meshgrid has indexing='ij'
+    p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3
+
+    # normalized ndc uv coordinates, (-1, 1)
+    ndc_u = 2 * xs / (W - 1) - 1
+    ndc_v = 2 * ys / (H - 1) - 1
+    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)
+
+    intrinsic_inv = torch.inverse(intrinsic)
+
+    p = p.view(-1, 3).float().to(device)  # N_rays, 3
+    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays, 3
+    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_rays, 3
+    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays, 3
+    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays, 3
+
+    image = image.permute(1, 2, 0)
+    color = image.view(-1, 3)
+    depth = depth.view(-1, 1) if depth is not None else None
+    mask = mask.view(-1, 1) if mask is not None else torch.ones([H * W, 1]).to(device)
+    sample = {
+        'rays_o': rays_o,
+        'rays_v': rays_v,
+        'rays_ndc_uv': rays_ndc_uv,
+        'rays_color': color,
+        # 'rays_depth': depth,
+        'rays_mask': mask,
+        'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
+    }
+    if depth is not None:
+        sample['rays_depth'] = depth
+
+    return sample
+
+
+def gen_random_rays_from_single_image(H, W, N_rays, image, intrinsic, c2w, depth=None, mask=None, dilated_mask=None,
+                                      importance_sample=False, h_patch_size=3):
+    """
+    generate random rays in world space, for a single image
+    :param H:
+    :param W:
+    :param N_rays:
+    :param image: [3, H, W]
+    :param intrinsic: [3,3]
+    :param c2w: [4,4]
+    :param depth: [H, W]
+    :param mask: [H, W]
+    :return:
+    """
+    device = image.device
+
+    if dilated_mask is None:
+        dilated_mask = mask
+
+    if not importance_sample:
+        pixels_x = torch.randint(low=0, high=W, size=[N_rays])
+        pixels_y = torch.randint(low=0, high=H, size=[N_rays])
+    elif importance_sample and dilated_mask is not None:  # sample more pts in the valid mask regions
+        pixels_x_1 = torch.randint(low=0, high=W, size=[N_rays // 4])
+        pixels_y_1 = torch.randint(low=0, high=H, size=[N_rays // 4])
+
+        ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
+                                torch.linspace(0, W - 1, W), indexing="ij")  # pytorch's meshgrid has indexing='ij'
+        p = torch.stack([xs, ys], dim=-1)  # H, W, 2
+
+        try:
+            p_valid = p[dilated_mask > 0]  # [num, 2]
+            random_idx = torch.randint(low=0, high=p_valid.shape[0], size=[N_rays // 4 * 3])
+        except:
+            print("dilated_mask.shape: ", dilated_mask.shape)
+            print("dilated_mask valid number", dilated_mask.sum())
+
+            raise ValueError("hhhh")
+        p_select = p_valid[random_idx]  # [N_rays//2, 2]
+        pixels_x_2 = p_select[:, 0]
+        pixels_y_2 = p_select[:, 1]
+
+        pixels_x = torch.cat([pixels_x_1, pixels_x_2], dim=0).to(torch.int64)
+        pixels_y = torch.cat([pixels_y_1, pixels_y_2], dim=0).to(torch.int64)
+
+    # - crop patch from images
+    offsets = build_patch_offset(h_patch_size).to(device)
+    grid_patch = torch.stack([pixels_x, pixels_y], dim=-1).view(-1, 1, 2) + offsets.float()  # [N_pts, Npx, 2]
+    patch_mask = (pixels_x > h_patch_size) * (pixels_x < (W - h_patch_size)) * (pixels_y > h_patch_size) * (
+            pixels_y < H - h_patch_size)  # [N_pts]
+    grid_patch_u = 2 * grid_patch[:, :, 0] / (W - 1) - 1
+    grid_patch_v = 2 * grid_patch[:, :, 1] / (H - 1) - 1
+    grid_patch_uv = torch.stack([grid_patch_u, grid_patch_v], dim=-1)  # [N_pts, Npx, 2]
+    patch_color = F.grid_sample(image[None, :, :, :], grid_patch_uv[None, :, :, :], mode='bilinear',
+                                padding_mode='zeros',align_corners=True)[0]  # [3, N_pts, Npx]
+    patch_color = patch_color.permute(1, 2, 0).contiguous()
+
+    # normalized ndc uv coordinates, (-1, 1)
+    ndc_u = 2 * pixels_x / (W - 1) - 1
+    ndc_v = 2 * pixels_y / (H - 1) - 1
+    rays_ndc_uv = torch.stack([ndc_u, ndc_v], dim=-1).view(-1, 2).float().to(device)
+
+    image = image.permute(1, 2, 0)  # H ,W, C
+    color = image[(pixels_y, pixels_x)]  # N_rays, 3
+
+    if mask is not None:
+        mask = mask[(pixels_y, pixels_x)]  # N_rays
+        patch_mask = patch_mask * mask  # N_rays
+        mask = mask.view(-1, 1)
+    else:
+        mask = torch.ones([N_rays, 1])
+
+    if depth is not None:
+        depth = depth[(pixels_y, pixels_x)]  # N_rays
+        depth = depth.view(-1, 1)
+
+    intrinsic_inv = torch.inverse(intrinsic)
+
+    p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float().to(device)  # N_rays, 3
+    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays, 3
+    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_rays, 3
+    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays, 3
+    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays, 3
+
+    sample = {
+        'rays_o': rays_o,
+        'rays_v': rays_v,
+        'rays_ndc_uv': rays_ndc_uv,
+        'rays_color': color,
+        # 'rays_depth': depth,
+        'rays_mask': mask,
+        'rays_norm_XYZ_cam': p,  # - XYZ_cam, before multiply depth,
+        'rays_patch_color': patch_color,
+        'rays_patch_mask': patch_mask.view(-1, 1)
+    }
+
+    if depth is not None:
+        sample['rays_depth'] = depth
+
+    return sample
+
+
+def gen_random_rays_of_patch_from_single_image(H, W, N_rays, num_neighboring_pts, patch_size,
+                                               image, intrinsic, c2w, depth=None, mask=None):
+    """
+    generate random rays in world space, for a single image
+    sample rays from local patches
+    :param H:
+    :param W:
+    :param N_rays: the number of center rays of patches
+    :param image: [3, H, W]
+    :param intrinsic: [3,3]
+    :param c2w: [4,4]
+    :param depth: [H, W]
+    :param mask: [H, W]
+    :return:
+    """
+    device = image.device
+    patch_radius_max = patch_size // 2
+
+    unit_u = 2 / (W - 1)
+    unit_v = 2 / (H - 1)
+
+    pixels_x_center = torch.randint(low=patch_size, high=W - patch_size, size=[N_rays])
+    pixels_y_center = torch.randint(low=patch_size, high=H - patch_size, size=[N_rays])
+
+    # normalized ndc uv coordinates, (-1, 1)
+    ndc_u_center = 2 * pixels_x_center / (W - 1) - 1
+    ndc_v_center = 2 * pixels_y_center / (H - 1) - 1
+    ndc_uv_center = torch.stack([ndc_u_center, ndc_v_center], dim=-1).view(-1, 2).float().to(device)[:, None,
+                    :]  # [N_rays, 1, 2]
+
+    shift_u, shift_v = torch.rand([N_rays, num_neighboring_pts, 1]), torch.rand(
+        [N_rays, num_neighboring_pts, 1])  # uniform distribution of [0,1)
+    shift_u = 2 * (shift_u - 0.5)  # mapping to [-1, 1)
+    shift_v = 2 * (shift_v - 0.5)
+
+    # - avoid sample points which are too close to center point
+    shift_uv = torch.cat([(shift_u * patch_radius_max) * unit_u, (shift_v * patch_radius_max) * unit_v],
+                         dim=-1)  # [N_rays, num_npts, 2]
+    neighboring_pts_uv = ndc_uv_center + shift_uv  # [N_rays, num_npts, 2]
+
+    sampled_pts_uv = torch.cat([ndc_uv_center, neighboring_pts_uv], dim=1)  # concat the center point
+
+    # sample the gts
+    color = F.grid_sample(image[None, :, :, :], sampled_pts_uv[None, :, :, :], mode='bilinear',
+                          align_corners=True)[0]  # [3, N_rays, num_npts]
+    depth = F.grid_sample(depth[None, None, :, :], sampled_pts_uv[None, :, :, :], mode='bilinear',
+                          align_corners=True)[0]  # [1, N_rays, num_npts]
+
+    mask = F.grid_sample(mask[None, None, :, :].to(torch.float32), sampled_pts_uv[None, :, :, :], mode='nearest',
+                         align_corners=True).to(torch.int64)[0]  # [1, N_rays, num_npts]
+
+    intrinsic_inv = torch.inverse(intrinsic)
+
+    sampled_pts_uv = sampled_pts_uv.view(N_rays * (1 + num_neighboring_pts), 2)
+    color = color.permute(1, 2, 0).contiguous().view(N_rays * (1 + num_neighboring_pts), 3)
+    depth = depth.permute(1, 2, 0).contiguous().view(N_rays * (1 + num_neighboring_pts), 1)
+    mask = mask.permute(1, 2, 0).contiguous().view(N_rays * (1 + num_neighboring_pts), 1)
+
+    pixels_x = (sampled_pts_uv[:, 0] + 1) * (W - 1) / 2
+    pixels_y = (sampled_pts_uv[:, 1] + 1) * (H - 1) / 2
+    p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float().to(device)  # N_rays*num_pts, 3
+    p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze()  # N_rays*num_pts, 3
+    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # N_rays*num_pts, 3
+    rays_v = torch.matmul(c2w[None, :3, :3], rays_v[:, :, None]).squeeze()  # N_rays*num_pts, 3
+    rays_o = c2w[None, :3, 3].expand(rays_v.shape)  # N_rays*num_pts, 3
+
+    sample = {
+        'rays_o': rays_o,
+        'rays_v': rays_v,
+        'rays_ndc_uv': sampled_pts_uv,
+        'rays_color': color,
+        'rays_depth': depth,
+        'rays_mask': mask,
+        # 'rays_norm_XYZ_cam': p  # - XYZ_cam, before multiply depth
+    }
+
+    return sample
+
+
+def gen_random_rays_from_batch_images(H, W, N_rays, images, intrinsics, c2ws, depths=None, masks=None):
+    """
+
+    :param H:
+    :param W:
+    :param N_rays:
+    :param images: [B,3,H,W]
+    :param intrinsics: [B, 3, 3]
+    :param c2ws: [B, 4, 4]
+    :param depths: [B,H,W]
+    :param masks: [B,H,W]
+    :return:
+    """
+    assert len(images.shape) == 4
+
+    rays_o = []
+    rays_v = []
+    rays_color = []
+    rays_depth = []
+    rays_mask = []
+    for i in range(images.shape[0]):
+        sample = gen_random_rays_from_single_image(H, W, N_rays, images[i], intrinsics[i], c2ws[i],
+                                                   depth=depths[i] if depths is not None else None,
+                                                   mask=masks[i] if masks is not None else None)
+        rays_o.append(sample['rays_o'])
+        rays_v.append(sample['rays_v'])
+        rays_color.append(sample['rays_color'])
+        if depths is not None:
+            rays_depth.append(sample['rays_depth'])
+        if masks is not None:
+            rays_mask.append(sample['rays_mask'])
+
+    sample = {
+        'rays_o': torch.stack(rays_o, dim=0),  # [batch, N_rays, 3]
+        'rays_v': torch.stack(rays_v, dim=0),
+        'rays_color': torch.stack(rays_color, dim=0),
+        'rays_depth': torch.stack(rays_depth, dim=0) if depths is not None else None,
+        'rays_mask': torch.stack(rays_mask, dim=0) if masks is not None else None
+    }
+    return sample
+
+
+from scipy.spatial.transform import Rotation as Rot
+from scipy.spatial.transform import Slerp
+
+
+def gen_rays_between(c2w_0, c2w_1, intrinsic, ratio, H, W, resolution_level=1):
+    device = c2w_0.device
+
+    l = resolution_level
+    tx = torch.linspace(0, W - 1, W // l)
+    ty = torch.linspace(0, H - 1, H // l)
+    pixels_x, pixels_y = torch.meshgrid(tx, ty, indexing="ij")
+    p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).to(device)  # W, H, 3
+
+    intrinsic_inv = torch.inverse(intrinsic[:3, :3])
+    p = torch.matmul(intrinsic_inv[None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
+    rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
+    trans = c2w_0[:3, 3] * (1.0 - ratio) + c2w_1[:3, 3] * ratio
+
+    pose_0 = c2w_0.detach().cpu().numpy()
+    pose_1 = c2w_1.detach().cpu().numpy()
+    pose_0 = np.linalg.inv(pose_0)
+    pose_1 = np.linalg.inv(pose_1)
+    rot_0 = pose_0[:3, :3]
+    rot_1 = pose_1[:3, :3]
+    rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
+    key_times = [0, 1]
+    key_rots = [rot_0, rot_1]
+    slerp = Slerp(key_times, rots)
+    rot = slerp(ratio)
+    pose = np.diag([1.0, 1.0, 1.0, 1.0])
+    pose = pose.astype(np.float32)
+    pose[:3, :3] = rot.as_matrix()
+    pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
+    pose = np.linalg.inv(pose)
+
+    c2w = torch.from_numpy(pose).to(device)
+    rot = torch.from_numpy(pose[:3, :3]).cuda()
+    trans = torch.from_numpy(pose[:3, 3]).cuda()
+    rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
+    rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3
+    return c2w, rays_o.transpose(0, 1).contiguous().view(-1, 3), rays_v.transpose(0, 1).contiguous().view(-1, 3)

SparseNeuS_demo_v1/models/render_utils.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ops.back_project import cam2pixel
+
+
+def sample_pdf(bins, weights, n_samples, det=False):
+    '''
+    :param bins: tensor of shape [N_rays, M+1], M is the number of bins
+    :param weights: tensor of shape [N_rays, M]
+    :param N_samples: number of samples along each ray
+    :param det: if True, will perform deterministic sampling
+    :return: [N_rays, N_samples]
+    '''
+    device = weights.device
+
+    weights = weights + 1e-5  # prevent nans
+    pdf = weights / torch.sum(weights, -1, keepdim=True)
+    cdf = torch.cumsum(pdf, -1)
+    cdf = torch.cat([torch.zeros_like(cdf[..., :1]).to(device), cdf], -1)
+
+    # if bins.shape[1] != weights.shape[1]:  # - minor modification, add this constraint
+    #     cdf = torch.cat([torch.zeros_like(cdf[..., :1]).to(device), cdf], -1)
+    # Take uniform samples
+    if det:
+        u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples).to(device)
+        u = u.expand(list(cdf.shape[:-1]) + [n_samples])
+    else:
+        u = torch.rand(list(cdf.shape[:-1]) + [n_samples]).to(device)
+
+    # Invert CDF
+    u = u.contiguous()
+    # inds = searchsorted(cdf, u, side='right')
+    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])
+
+    # pdb.set_trace()
+    return samples
+
+
+def sample_ptsFeatures_from_featureVolume(pts, featureVolume, vol_dims=None, partial_vol_origin=None, vol_size=None):
+    """
+    sample feature of pts_wrd from featureVolume, all in world space
+    :param pts: [N_rays, n_samples, 3]
+    :param featureVolume: [C,wX,wY,wZ]
+    :param vol_dims: [3] "3" for dimX, dimY, dimZ
+    :param partial_vol_origin: [3]
+    :return: pts_feature: [N_rays, n_samples, C]
+    :return: valid_mask: [N_rays]
+    """
+
+    N_rays, n_samples, _ = pts.shape
+
+    if vol_dims is None:
+        pts_normalized = pts
+    else:
+        # normalized to (-1, 1)
+        pts_normalized = 2 * (pts - partial_vol_origin[None, None, :]) / (vol_size * (vol_dims[None, None, :] - 1)) - 1
+
+    valid_mask = (torch.abs(pts_normalized[:, :, 0]) < 1.0) & (
+            torch.abs(pts_normalized[:, :, 1]) < 1.0) & (
+                         torch.abs(pts_normalized[:, :, 2]) < 1.0)  # (N_rays, n_samples)
+
+    pts_normalized = torch.flip(pts_normalized, dims=[-1])  # ! reverse the xyz for grid_sample
+
+    # ! checked grid_sample, (x,y,z) is for (D,H,W), reverse for (W,H,D)
+    pts_feature = F.grid_sample(featureVolume[None, :, :, :, :], pts_normalized[None, None, :, :, :],
+                                padding_mode='zeros',
+                                align_corners=True).view(-1, N_rays, n_samples)  # [C, N_rays, n_samples]
+
+    pts_feature = pts_feature.permute(1, 2, 0)  # [N_rays, n_samples, C]
+    return pts_feature, valid_mask
+
+
+def sample_ptsFeatures_from_featureMaps(pts, featureMaps, w2cs, intrinsics, WH, proj_matrix=None, return_mask=False):
+    """
+    sample features of pts from 2d feature maps
+    :param pts: [N_rays, N_samples, 3]
+    :param featureMaps: [N_views, C, H, W]
+    :param w2cs: [N_views, 4, 4]
+    :param intrinsics: [N_views, 3, 3]
+    :param proj_matrix: [N_views, 4, 4]
+    :param HW:
+    :return:
+    """
+    # normalized to (-1, 1)
+    N_rays, n_samples, _ = pts.shape
+    N_views = featureMaps.shape[0]
+
+    if proj_matrix is None:
+        proj_matrix = torch.matmul(intrinsics, w2cs[:, :3, :])
+
+    pts = pts.permute(2, 0, 1).contiguous().view(1, 3, N_rays, n_samples).repeat(N_views, 1, 1, 1)
+    pixel_grids = cam2pixel(pts, proj_matrix[:, :3, :3], proj_matrix[:, :3, 3:],
+                            'zeros', sizeH=WH[1], sizeW=WH[0])  # (nviews, N_rays, n_samples, 2)
+
+    valid_mask = (torch.abs(pixel_grids[:, :, :, 0]) < 1.0) & (
+            torch.abs(pixel_grids[:, :, :, 1]) < 1.00)  # (nviews, N_rays, n_samples)
+
+    pts_feature = F.grid_sample(featureMaps, pixel_grids,
+                                padding_mode='zeros',
+                                align_corners=True)  # [N_views, C, N_rays, n_samples]
+
+    if return_mask:
+        return pts_feature, valid_mask
+    else:
+        return pts_feature

SparseNeuS_demo_v1/models/rendering_network.py

@@ -0,0 +1,129 @@
+# the codes are partly borrowed from IBRNet
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_set_profiling_mode(False)
+
+
+# default tensorflow initialization of linear layers
+def weights_init(m):
+    if isinstance(m, nn.Linear):
+        nn.init.kaiming_normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias.data)
+
+
+@torch.jit.script
+def fused_mean_variance(x, weight):
+    mean = torch.sum(x * weight, dim=2, keepdim=True)
+    var = torch.sum(weight * (x - mean) ** 2, dim=2, keepdim=True)
+    return mean, var
+
+
+class GeneralRenderingNetwork(nn.Module):
+    """
+    This model is not sensitive to finetuning
+    """
+
+    def __init__(self, in_geometry_feat_ch=8, in_rendering_feat_ch=56, anti_alias_pooling=True):
+        super(GeneralRenderingNetwork, self).__init__()
+
+        self.in_geometry_feat_ch = in_geometry_feat_ch
+        self.in_rendering_feat_ch = in_rendering_feat_ch
+        self.anti_alias_pooling = anti_alias_pooling
+
+        if self.anti_alias_pooling:
+            self.s = nn.Parameter(torch.tensor(0.2), requires_grad=True)
+        activation_func = nn.ELU(inplace=True)
+
+        self.ray_dir_fc = nn.Sequential(nn.Linear(4, 16),
+                                        activation_func,
+                                        nn.Linear(16, in_rendering_feat_ch + 3),
+                                        activation_func)
+
+        self.base_fc = nn.Sequential(nn.Linear((in_rendering_feat_ch + 3) * 3 + in_geometry_feat_ch, 64),
+                                     activation_func,
+                                     nn.Linear(64, 32),
+                                     activation_func)
+
+        self.vis_fc = nn.Sequential(nn.Linear(32, 32),
+                                    activation_func,
+                                    nn.Linear(32, 33),
+                                    activation_func,
+                                    )
+
+        self.vis_fc2 = nn.Sequential(nn.Linear(32, 32),
+                                     activation_func,
+                                     nn.Linear(32, 1),
+                                     nn.Sigmoid()
+                                     )
+
+        self.rgb_fc = nn.Sequential(nn.Linear(32 + 1 + 4, 16),
+                                    activation_func,
+                                    nn.Linear(16, 8),
+                                    activation_func,
+                                    nn.Linear(8, 1))
+
+        self.base_fc.apply(weights_init)
+        self.vis_fc2.apply(weights_init)
+        self.vis_fc.apply(weights_init)
+        self.rgb_fc.apply(weights_init)
+
+    def forward(self, geometry_feat, rgb_feat, ray_diff, mask):
+        '''
+        :param geometry_feat: geometry features indicates sdf  [n_rays, n_samples, n_feat]
+        :param rgb_feat: rgbs and image features [n_views, n_rays, n_samples, n_feat]
+        :param ray_diff: ray direction difference [n_views, n_rays, n_samples, 4], first 3 channels are directions,
+        last channel is inner product
+        :param mask: mask for whether each projection is valid or not. [n_views, n_rays, n_samples]
+        :return: rgb and density output, [n_rays, n_samples, 4]
+        '''
+
+        rgb_feat = rgb_feat.permute(1, 2, 0, 3).contiguous()
+        ray_diff = ray_diff.permute(1, 2, 0, 3).contiguous()
+        mask = mask[:, :, :, None].permute(1, 2, 0, 3).contiguous()
+        num_views = rgb_feat.shape[2]
+        geometry_feat = geometry_feat[:, :, None, :].repeat(1, 1, num_views, 1)
+
+        direction_feat = self.ray_dir_fc(ray_diff)
+        rgb_in = rgb_feat[..., :3]
+        rgb_feat = rgb_feat + direction_feat
+
+        if self.anti_alias_pooling:
+            _, dot_prod = torch.split(ray_diff, [3, 1], dim=-1)
+            exp_dot_prod = torch.exp(torch.abs(self.s) * (dot_prod - 1))
+            weight = (exp_dot_prod - torch.min(exp_dot_prod, dim=2, keepdim=True)[0]) * mask
+            weight = weight / (torch.sum(weight, dim=2, keepdim=True) + 1e-8)
+        else:
+            weight = mask / (torch.sum(mask, dim=2, keepdim=True) + 1e-8)
+
+        # compute mean and variance across different views for each point
+        mean, var = fused_mean_variance(rgb_feat, weight)  # [n_rays, n_samples, 1, n_feat]
+        globalfeat = torch.cat([mean, var], dim=-1)  # [n_rays, n_samples, 1, 2*n_feat]
+
+        x = torch.cat([geometry_feat, globalfeat.expand(-1, -1, num_views, -1), rgb_feat],
+                      dim=-1)  # [n_rays, n_samples, n_views, 3*n_feat+n_geo_feat]
+        x = self.base_fc(x)
+
+        x_vis = self.vis_fc(x * weight)
+        x_res, vis = torch.split(x_vis, [x_vis.shape[-1] - 1, 1], dim=-1)
+        vis = F.sigmoid(vis) * mask
+        x = x + x_res
+        vis = self.vis_fc2(x * vis) * mask
+
+        # rgb computation
+        x = torch.cat([x, vis, ray_diff], dim=-1)
+        x = self.rgb_fc(x)
+        x = x.masked_fill(mask == 0, -1e9)
+        blending_weights_valid = F.softmax(x, dim=2)  # color blending
+        rgb_out = torch.sum(rgb_in * blending_weights_valid, dim=2)
+
+        mask = mask.detach().to(rgb_out.dtype)  # [n_rays, n_samples, n_views, 1]
+        mask = torch.sum(mask, dim=2, keepdim=False)
+        mask = mask >= 2  # more than 2 views see the point
+        mask = torch.sum(mask.to(rgb_out.dtype), dim=1, keepdim=False)
+        valid_mask = mask > 8  # valid rays, more than 8 valid samples
+        return rgb_out, valid_mask  # (N_rays, n_samples, 3), (N_rays, 1)

SparseNeuS_demo_v1/models/sparse_neus_renderer.py

@@ -0,0 +1,985 @@
+"""
+The codes are heavily borrowed from NeuS
+"""
+
+import os
+import cv2 as cv
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import logging
+import mcubes
+from icecream import ic
+from models.render_utils import sample_pdf
+
+from models.projector import Projector
+from tsparse.torchsparse_utils import sparse_to_dense_channel
+
+from models.fast_renderer import FastRenderer
+
+from models.patch_projector import PatchProjector
+
+
+class SparseNeuSRenderer(nn.Module):
+    """
+    conditional neus render;
+    optimize on normalized world space;
+    warped by nn.Module to support DataParallel traning
+    """
+
+    def __init__(self,
+                 rendering_network_outside,
+                 sdf_network,
+                 variance_network,
+                 rendering_network,
+                 n_samples,
+                 n_importance,
+                 n_outside,
+                 perturb,
+                 alpha_type='div',
+                 conf=None
+                 ):
+        super(SparseNeuSRenderer, self).__init__()
+
+        self.conf = conf
+        self.base_exp_dir = conf['general.base_exp_dir']
+
+        # network setups
+        self.rendering_network_outside = rendering_network_outside
+        self.sdf_network = sdf_network
+        self.variance_network = variance_network
+        self.rendering_network = rendering_network
+
+        self.n_samples = n_samples
+        self.n_importance = n_importance
+        self.n_outside = n_outside
+        self.perturb = perturb
+        self.alpha_type = alpha_type
+
+        self.rendering_projector = Projector()  # used to obtain features for generalized rendering
+
+        self.h_patch_size = self.conf.get_int('model.h_patch_size', default=3)
+        self.patch_projector = PatchProjector(self.h_patch_size)
+
+        self.ray_tracer = FastRenderer()  # ray_tracer to extract depth maps from sdf_volume
+
+        # - fitted rendering or general rendering
+        try:
+            self.if_fitted_rendering = self.sdf_network.if_fitted_rendering
+        except:
+            self.if_fitted_rendering = False
+
+    def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_variance,
+                  conditional_valid_mask_volume=None):
+        device = rays_o.device
+        batch_size, n_samples = z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
+
+        if conditional_valid_mask_volume is not None:
+            pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), conditional_valid_mask_volume)
+            pts_mask = pts_mask.reshape(batch_size, n_samples)
+            pts_mask = pts_mask[:, :-1] * pts_mask[:, 1:]  # [batch_size, n_samples-1]
+        else:
+            pts_mask = torch.ones([batch_size, n_samples]).to(pts.device)
+
+        sdf = sdf.reshape(batch_size, n_samples)
+        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
+        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
+        mid_sdf = (prev_sdf + next_sdf) * 0.5
+        dot_val = None
+        if self.alpha_type == 'uniform':
+            dot_val = torch.ones([batch_size, n_samples - 1]) * -1.0
+        else:
+            dot_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
+            prev_dot_val = torch.cat([torch.zeros([batch_size, 1]).to(device), dot_val[:, :-1]], dim=-1)
+            dot_val = torch.stack([prev_dot_val, dot_val], dim=-1)
+            dot_val, _ = torch.min(dot_val, dim=-1, keepdim=False)
+            dot_val = dot_val.clip(-10.0, 0.0) * pts_mask
+        dist = (next_z_vals - prev_z_vals)
+        prev_esti_sdf = mid_sdf - dot_val * dist * 0.5
+        next_esti_sdf = mid_sdf + dot_val * dist * 0.5
+        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_variance)
+        next_cdf = torch.sigmoid(next_esti_sdf * inv_variance)
+        alpha_sdf = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
+
+        alpha = alpha_sdf
+
+        # - apply pts_mask
+        alpha = pts_mask * alpha
+
+        weights = alpha * torch.cumprod(
+            torch.cat([torch.ones([batch_size, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+
+        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
+        return z_samples
+
+    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, lod,
+                   sdf_network, gru_fusion,
+                   # * related to conditional feature
+                   conditional_volume=None,
+                   conditional_valid_mask_volume=None
+                   ):
+        device = rays_o.device
+        batch_size, n_samples = z_vals.shape
+        _, n_importance = new_z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
+
+        if conditional_valid_mask_volume is not None:
+            pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), conditional_valid_mask_volume)
+            pts_mask = pts_mask.reshape(batch_size, n_importance)
+            pts_mask_bool = (pts_mask > 0).view(-1)
+        else:
+            pts_mask = torch.ones([batch_size, n_importance]).to(pts.device)
+
+        new_sdf = torch.ones([batch_size * n_importance, 1]).to(pts.dtype).to(device) * 100
+
+        if torch.sum(pts_mask) > 1:
+            new_outputs = sdf_network.sdf(pts.reshape(-1, 3)[pts_mask_bool], conditional_volume, lod=lod)
+            new_sdf[pts_mask_bool] = new_outputs['sdf_pts_scale%d' % lod]  # .reshape(batch_size, n_importance)
+
+        new_sdf = new_sdf.view(batch_size, n_importance)
+
+        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
+        sdf = torch.cat([sdf, new_sdf], dim=-1)
+
+        z_vals, index = torch.sort(z_vals, dim=-1)
+        xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
+        index = index.reshape(-1)
+        sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
+
+        return z_vals, sdf
+
+    @torch.no_grad()
+    def get_pts_mask_for_conditional_volume(self, pts, mask_volume):
+        """
+
+        :param pts: [N, 3]
+        :param mask_volume: [1, 1, X, Y, Z]
+        :return:
+        """
+        num_pts = pts.shape[0]
+        pts = pts.view(1, 1, 1, num_pts, 3)  # - should be in range (-1, 1)
+
+        pts = torch.flip(pts, dims=[-1])
+
+        pts_mask = F.grid_sample(mask_volume, pts, mode='nearest')  # [1, c, 1, 1, num_pts]
+        pts_mask = pts_mask.view(-1, num_pts).permute(1, 0).contiguous()  # [num_pts, 1]
+
+        return pts_mask
+
+    def render_core(self,
+                    rays_o,
+                    rays_d,
+                    z_vals,
+                    sample_dist,
+                    lod,
+                    sdf_network,
+                    rendering_network,
+                    background_alpha=None,  # - no use here
+                    background_sampled_color=None,  # - no use here
+                    background_rgb=None,  # - no use here
+                    alpha_inter_ratio=0.0,
+                    # * related to conditional feature
+                    conditional_volume=None,
+                    conditional_valid_mask_volume=None,
+                    # * 2d feature maps
+                    feature_maps=None,
+                    color_maps=None,
+                    w2cs=None,
+                    intrinsics=None,
+                    img_wh=None,
+                    query_c2w=None,  # - used for testing
+                    if_general_rendering=True,
+                    if_render_with_grad=True,
+                    # * used for blending mlp rendering network
+                    img_index=None,
+                    rays_uv=None,
+                    # * used for clear bg and fg
+                    bg_num=0
+                    ):
+        device = rays_o.device
+        N_rays = rays_o.shape[0]
+        _, n_samples = z_vals.shape
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape).to(device)], -1)
+
+        mid_z_vals = z_vals + dists * 0.5
+        mid_dists = mid_z_vals[..., 1:] - mid_z_vals[..., :-1]
+
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
+        dirs = rays_d[:, None, :].expand(pts.shape)
+
+        pts = pts.reshape(-1, 3)
+        dirs = dirs.reshape(-1, 3)
+
+        # * if conditional_volume is restored from sparse volume, need mask for pts
+        if conditional_valid_mask_volume is not None:
+            pts_mask = self.get_pts_mask_for_conditional_volume(pts, conditional_valid_mask_volume)
+            pts_mask = pts_mask.reshape(N_rays, n_samples).float().detach()
+            pts_mask_bool = (pts_mask > 0).view(-1)
+
+            if torch.sum(pts_mask_bool.float()) < 1:  # ! when render out image, may meet this problem
+                pts_mask_bool[:100] = True
+
+        else:
+            pts_mask = torch.ones([N_rays, n_samples]).to(pts.device)
+        # import ipdb; ipdb.set_trace()
+        # pts_valid = pts[pts_mask_bool]
+        sdf_nn_output = sdf_network.sdf(pts[pts_mask_bool], conditional_volume, lod=lod)
+
+        sdf = torch.ones([N_rays * n_samples, 1]).to(pts.dtype).to(device) * 100
+        sdf[pts_mask_bool] = sdf_nn_output['sdf_pts_scale%d' % lod]  # [N_rays*n_samples, 1]
+        feature_vector_valid = sdf_nn_output['sdf_features_pts_scale%d' % lod]
+        feature_vector = torch.zeros([N_rays * n_samples, feature_vector_valid.shape[1]]).to(pts.dtype).to(device)
+        feature_vector[pts_mask_bool] = feature_vector_valid
+
+        # * estimate alpha from sdf
+        gradients = torch.zeros([N_rays * n_samples, 3]).to(pts.dtype).to(device)
+        # import ipdb; ipdb.set_trace()
+        gradients[pts_mask_bool] = sdf_network.gradient(
+            pts[pts_mask_bool], conditional_volume, lod=lod).squeeze()
+
+        sampled_color_mlp = None
+        rendering_valid_mask_mlp = None
+        sampled_color_patch = None
+        rendering_patch_mask = None
+
+        if self.if_fitted_rendering:  # used for fine-tuning
+            position_latent = sdf_nn_output['sampled_latent_scale%d' % lod]
+            sampled_color_mlp = torch.zeros([N_rays * n_samples, 3]).to(pts.dtype).to(device)
+            sampled_color_mlp_mask = torch.zeros([N_rays * n_samples, 1]).to(pts.dtype).to(device)
+
+            # - extract pixel
+            pts_pixel_color, pts_pixel_mask = self.patch_projector.pixel_warp(
+                pts[pts_mask_bool][:, None, :], color_maps, intrinsics,
+                w2cs, img_wh=None)  # [N_rays * n_samples,1, N_views,  3] , [N_rays*n_samples, 1, N_views]
+            pts_pixel_color = pts_pixel_color[:, 0, :, :]  # [N_rays * n_samples, N_views,  3]
+            pts_pixel_mask = pts_pixel_mask[:, 0, :]  # [N_rays*n_samples, N_views]
+
+            # - extract patch
+            if_patch_blending = False if rays_uv is None else True
+            pts_patch_color, pts_patch_mask = None, None
+            if if_patch_blending:
+                pts_patch_color, pts_patch_mask = self.patch_projector.patch_warp(
+                    pts.reshape([N_rays, n_samples, 3]),
+                    rays_uv, gradients.reshape([N_rays, n_samples, 3]),
+                    color_maps,
+                    intrinsics[0], intrinsics,
+                    query_c2w[0], torch.inverse(w2cs), img_wh=None
+                )  # (N_rays, n_samples, N_src, Npx, 3), (N_rays, n_samples, N_src, Npx)
+                N_src, Npx = pts_patch_mask.shape[2:]
+                pts_patch_color = pts_patch_color.view(N_rays * n_samples, N_src, Npx, 3)[pts_mask_bool]
+                pts_patch_mask = pts_patch_mask.view(N_rays * n_samples, N_src, Npx)[pts_mask_bool]
+
+                sampled_color_patch = torch.zeros([N_rays * n_samples, Npx, 3]).to(device)
+                sampled_color_patch_mask = torch.zeros([N_rays * n_samples, 1]).to(device)
+
+            sampled_color_mlp_, sampled_color_mlp_mask_, \
+            sampled_color_patch_, sampled_color_patch_mask_ = sdf_network.color_blend(
+                pts[pts_mask_bool],
+                position_latent,
+                gradients[pts_mask_bool],
+                dirs[pts_mask_bool],
+                feature_vector[pts_mask_bool],
+                img_index=img_index,
+                pts_pixel_color=pts_pixel_color,
+                pts_pixel_mask=pts_pixel_mask,
+                pts_patch_color=pts_patch_color,
+                pts_patch_mask=pts_patch_mask
+
+            )  # [n, 3], [n, 1]
+            sampled_color_mlp[pts_mask_bool] = sampled_color_mlp_
+            sampled_color_mlp_mask[pts_mask_bool] = sampled_color_mlp_mask_.float()
+            sampled_color_mlp = sampled_color_mlp.view(N_rays, n_samples, 3)
+            sampled_color_mlp_mask = sampled_color_mlp_mask.view(N_rays, n_samples)
+            rendering_valid_mask_mlp = torch.mean(pts_mask * sampled_color_mlp_mask, dim=-1, keepdim=True) > 0.5
+
+            # patch blending
+            if if_patch_blending:
+                sampled_color_patch[pts_mask_bool] = sampled_color_patch_
+                sampled_color_patch_mask[pts_mask_bool] = sampled_color_patch_mask_.float()
+                sampled_color_patch = sampled_color_patch.view(N_rays, n_samples, Npx, 3)
+                sampled_color_patch_mask = sampled_color_patch_mask.view(N_rays, n_samples)
+                rendering_patch_mask = torch.mean(pts_mask * sampled_color_patch_mask, dim=-1,
+                                                  keepdim=True) > 0.5  # [N_rays, 1]
+            else:
+                sampled_color_patch, rendering_patch_mask = None, None
+
+        if if_general_rendering:  # used for general training
+            # [512, 128, 16]; [4, 512, 128, 59]; [4, 512, 128, 4]
+            ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask, _, _ = self.rendering_projector.compute(
+                pts.view(N_rays, n_samples, 3),
+                # * 3d geometry feature volumes
+                geometryVolume=conditional_volume[0],
+                geometryVolumeMask=conditional_valid_mask_volume[0],
+                # * 2d rendering feature maps
+                rendering_feature_maps=feature_maps, # [n_views, 56, 256, 256]
+                color_maps=color_maps,
+                w2cs=w2cs,
+                intrinsics=intrinsics,
+                img_wh=img_wh,
+                query_img_idx=0,  # the index of the N_views dim for rendering
+                query_c2w=query_c2w,
+            )
+
+            # (N_rays, n_samples, 3)
+            if if_render_with_grad:
+                # import ipdb; ipdb.set_trace()
+                # [nrays, 3] [nrays, 1]
+                sampled_color, rendering_valid_mask = rendering_network(
+                    ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask)
+                # import ipdb; ipdb.set_trace()
+            else:
+                with torch.no_grad():
+                    sampled_color, rendering_valid_mask = rendering_network(
+                        ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask)
+        else:
+            sampled_color, rendering_valid_mask = None, None
+
+        inv_variance = self.variance_network(feature_vector)[:, :1].clip(1e-6, 1e6)
+
+        true_dot_val = (dirs * gradients).sum(-1, keepdim=True)  # * calculate
+
+        iter_cos = -(F.relu(-true_dot_val * 0.5 + 0.5) * (1.0 - alpha_inter_ratio) + F.relu(
+            -true_dot_val) * alpha_inter_ratio)  # always non-positive
+
+        iter_cos = iter_cos * pts_mask.view(-1, 1)
+
+        true_estimate_sdf_half_next = sdf + iter_cos.clip(-10.0, 10.0) * dists.reshape(-1, 1) * 0.5
+        true_estimate_sdf_half_prev = sdf - iter_cos.clip(-10.0, 10.0) * dists.reshape(-1, 1) * 0.5
+
+        prev_cdf = torch.sigmoid(true_estimate_sdf_half_prev * inv_variance)
+        next_cdf = torch.sigmoid(true_estimate_sdf_half_next * inv_variance)
+
+        p = prev_cdf - next_cdf
+        c = prev_cdf
+
+        if self.alpha_type == 'div':
+            alpha_sdf = ((p + 1e-5) / (c + 1e-5)).reshape(N_rays, n_samples).clip(0.0, 1.0)
+        elif self.alpha_type == 'uniform':
+            uniform_estimate_sdf_half_next = sdf - dists.reshape(-1, 1) * 0.5
+            uniform_estimate_sdf_half_prev = sdf + dists.reshape(-1, 1) * 0.5
+            uniform_prev_cdf = torch.sigmoid(uniform_estimate_sdf_half_prev * inv_variance)
+            uniform_next_cdf = torch.sigmoid(uniform_estimate_sdf_half_next * inv_variance)
+            uniform_alpha = F.relu(
+                (uniform_prev_cdf - uniform_next_cdf + 1e-5) / (uniform_prev_cdf + 1e-5)).reshape(
+                N_rays, n_samples).clip(0.0, 1.0)
+            alpha_sdf = uniform_alpha
+        else:
+            assert False
+
+        alpha = alpha_sdf
+
+        # - apply pts_mask
+        alpha = alpha * pts_mask
+
+        # pts_radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(N_rays, n_samples)
+        # inside_sphere = (pts_radius < 1.0).float().detach()
+        # relax_inside_sphere = (pts_radius < 1.2).float().detach()
+        inside_sphere = pts_mask
+        relax_inside_sphere = pts_mask
+
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([N_rays, 1]).to(device), 1. - alpha + 1e-7], -1), -1)[:,
+                          :-1]  # n_rays, n_samples
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+        alpha_sum = alpha.sum(dim=-1, keepdim=True)
+
+        if bg_num > 0:
+            weights_sum_fg = weights[:, :-bg_num].sum(dim=-1, keepdim=True)
+        else:
+            weights_sum_fg = weights_sum
+
+        if sampled_color is not None:
+            color = (sampled_color * weights[:, :, None]).sum(dim=1)
+        else:
+            color = None
+        # import ipdb; ipdb.set_trace()
+
+        if background_rgb is not None and color is not None:
+            color = color + background_rgb * (1.0 - weights_sum)
+            # print("color device:" + str(color.device))
+        # if color is not None:
+        #     # import ipdb; ipdb.set_trace()
+        #     color = color + (1.0 - weights_sum)
+
+
+        ###################*  mlp color rendering  #####################
+        color_mlp = None
+        # import ipdb; ipdb.set_trace()
+        if sampled_color_mlp is not None:
+            color_mlp = (sampled_color_mlp * weights[:, :, None]).sum(dim=1)
+
+        if background_rgb is not None and color_mlp is not None:
+            color_mlp = color_mlp + background_rgb * (1.0 - weights_sum)
+
+        ############################ *  patch blending  ################
+        blended_color_patch = None
+        if sampled_color_patch is not None:
+            blended_color_patch = (sampled_color_patch * weights[:, :, None, None]).sum(dim=1)  # [N_rays, Npx, 3]
+
+        ######################################################
+
+        gradient_error = (torch.linalg.norm(gradients.reshape(N_rays, n_samples, 3), ord=2,
+                                            dim=-1) - 1.0) ** 2
+        # ! the gradient normal should be masked out, the pts out of the bounding box should also be penalized
+        gradient_error = (pts_mask * gradient_error).sum() / (
+                (pts_mask).sum() + 1e-5)
+
+        depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
+        # print("[TEST]: weights_sum in render_core", weights_sum.mean())
+        # print("[TEST]: weights_sum in render_core NAN number", weights_sum.isnan().sum())
+        # if weights_sum.isnan().sum() > 0:
+        #     import ipdb; ipdb.set_trace()
+        return {
+            'color': color,
+            'color_mask': rendering_valid_mask,  # (N_rays, 1)
+            'color_mlp': color_mlp,
+            'color_mlp_mask': rendering_valid_mask_mlp,
+            'sdf': sdf,  # (N_rays, n_samples)
+            'depth': depth,  # (N_rays, 1)
+            'dists': dists,
+            'gradients': gradients.reshape(N_rays, n_samples, 3),
+            'variance': 1.0 / inv_variance,
+            'mid_z_vals': mid_z_vals,
+            'weights': weights,
+            'weights_sum': weights_sum,
+            'alpha_sum': alpha_sum,
+            'alpha_mean': alpha.mean(),
+            'cdf': c.reshape(N_rays, n_samples),
+            'gradient_error': gradient_error,
+            'inside_sphere': inside_sphere,
+            'blended_color_patch': blended_color_patch,
+            'blended_color_patch_mask': rendering_patch_mask,
+            'weights_sum_fg': weights_sum_fg
+        }
+
+    def render(self, rays_o, rays_d, near, far, sdf_network, rendering_network,
+               perturb_overwrite=-1,
+               background_rgb=None,
+               alpha_inter_ratio=0.0,
+               # * related to conditional feature
+               lod=None,
+               conditional_volume=None,
+               conditional_valid_mask_volume=None,
+               # * 2d feature maps
+               feature_maps=None,
+               color_maps=None,
+               w2cs=None,
+               intrinsics=None,
+               img_wh=None,
+               query_c2w=None,  # -used for testing
+               if_general_rendering=True,
+               if_render_with_grad=True,
+               # * used for blending mlp rendering network
+               img_index=None,
+               rays_uv=None,
+               # * importance sample for second lod network
+               pre_sample=False,  # no use here
+               # * for clear foreground
+               bg_ratio=0.0
+               ):
+        device = rays_o.device
+        N_rays = len(rays_o)
+        # sample_dist = 2.0 / self.n_samples
+        sample_dist = ((far - near) / self.n_samples).mean().item()
+        z_vals = torch.linspace(0.0, 1.0, self.n_samples).to(device)
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        bg_num = int(self.n_samples * bg_ratio)
+
+        if z_vals.shape[0] == 1:
+            z_vals = z_vals.repeat(N_rays, 1)
+
+        if bg_num > 0:
+            z_vals_bg = z_vals[:, self.n_samples - bg_num:]
+            z_vals = z_vals[:, :self.n_samples - bg_num]
+
+        n_samples = self.n_samples - bg_num
+        perturb = self.perturb
+
+        # - significantly speed up training, for the second lod network
+        if pre_sample:
+            z_vals = self.sample_z_vals_from_maskVolume(rays_o, rays_d, near, far,
+                                                        conditional_valid_mask_volume)
+
+        if perturb_overwrite >= 0:
+            perturb = perturb_overwrite
+        if perturb > 0:
+            # 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).to(device)
+            z_vals = lower + (upper - lower) * t_rand
+
+        background_alpha = None
+        background_sampled_color = None
+        z_val_before = z_vals.clone()
+        # Up sample
+        if self.n_importance > 0:
+            with torch.no_grad():
+                pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+
+                sdf_outputs = sdf_network.sdf(
+                    pts.reshape(-1, 3), conditional_volume, lod=lod)
+                # pdb.set_trace()
+                sdf = sdf_outputs['sdf_pts_scale%d' % lod].reshape(N_rays, self.n_samples - bg_num)
+
+                n_steps = 4
+                for i in range(n_steps):
+                    new_z_vals = self.up_sample(rays_o, rays_d, z_vals, sdf, self.n_importance // n_steps,
+                                                64 * 2 ** i,
+                                                conditional_valid_mask_volume=conditional_valid_mask_volume,
+                                                )
+
+                    # if new_z_vals.isnan().sum() > 0:
+                    #     import ipdb; ipdb.set_trace()
+
+                    z_vals, sdf = self.cat_z_vals(
+                        rays_o, rays_d, z_vals, new_z_vals, sdf, lod,
+                        sdf_network, gru_fusion=False,
+                        conditional_volume=conditional_volume,
+                        conditional_valid_mask_volume=conditional_valid_mask_volume,
+                    )
+
+                del sdf
+
+            n_samples = self.n_samples + self.n_importance
+
+        # Background
+        ret_outside = None
+
+        # Render
+        if bg_num > 0:
+            z_vals = torch.cat([z_vals, z_vals_bg], dim=1)
+        # if z_vals.isnan().sum() > 0:
+        #     import ipdb; ipdb.set_trace()
+        ret_fine = self.render_core(rays_o,
+                                    rays_d,
+                                    z_vals,
+                                    sample_dist,
+                                    lod,
+                                    sdf_network,
+                                    rendering_network,
+                                    background_rgb=background_rgb,
+                                    background_alpha=background_alpha,
+                                    background_sampled_color=background_sampled_color,
+                                    alpha_inter_ratio=alpha_inter_ratio,
+                                    # * related to conditional feature
+                                    conditional_volume=conditional_volume,
+                                    conditional_valid_mask_volume=conditional_valid_mask_volume,
+                                    # * 2d feature maps
+                                    feature_maps=feature_maps,
+                                    color_maps=color_maps,
+                                    w2cs=w2cs,
+                                    intrinsics=intrinsics,
+                                    img_wh=img_wh,
+                                    query_c2w=query_c2w,
+                                    if_general_rendering=if_general_rendering,
+                                    if_render_with_grad=if_render_with_grad,
+                                    # * used for blending mlp rendering network
+                                    img_index=img_index,
+                                    rays_uv=rays_uv
+                                    )
+
+        color_fine = ret_fine['color']
+
+        if self.n_outside > 0:
+            color_fine_mask = torch.logical_or(ret_fine['color_mask'], ret_outside['color_mask'])
+        else:
+            color_fine_mask = ret_fine['color_mask']
+
+        weights = ret_fine['weights']
+        weights_sum = ret_fine['weights_sum']
+
+        gradients = ret_fine['gradients']
+        mid_z_vals = ret_fine['mid_z_vals']
+
+        # depth = (mid_z_vals * weights[:, :n_samples]).sum(dim=1, keepdim=True)
+        depth = ret_fine['depth']
+        depth_varaince = ((mid_z_vals - depth) ** 2 * weights[:, :n_samples]).sum(dim=-1, keepdim=True)
+        variance = ret_fine['variance'].reshape(N_rays, n_samples).mean(dim=-1, keepdim=True)
+
+        # - randomly sample points from the volume, and maximize the sdf
+        pts_random = torch.rand([1024, 3]).float().to(device) * 2 - 1  # normalized to (-1, 1)
+        sdf_random = sdf_network.sdf(pts_random, conditional_volume, lod=lod)['sdf_pts_scale%d' % lod]
+
+        result = {
+            'depth': depth,
+            'color_fine': color_fine,
+            'color_fine_mask': color_fine_mask,
+            'color_outside': ret_outside['color'] if ret_outside is not None else None,
+            'color_outside_mask': ret_outside['color_mask'] if ret_outside is not None else None,
+            'color_mlp': ret_fine['color_mlp'],
+            'color_mlp_mask': ret_fine['color_mlp_mask'],
+            'variance': variance.mean(),
+            'cdf_fine': ret_fine['cdf'],
+            'depth_variance': depth_varaince,
+            'weights_sum': weights_sum,
+            'weights_max': torch.max(weights, dim=-1, keepdim=True)[0],
+            'alpha_sum': ret_fine['alpha_sum'].mean(),
+            'alpha_mean': ret_fine['alpha_mean'],
+            'gradients': gradients,
+            'weights': weights,
+            'gradient_error_fine': ret_fine['gradient_error'],
+            'inside_sphere': ret_fine['inside_sphere'],
+            'sdf': ret_fine['sdf'],
+            'sdf_random': sdf_random,
+            'blended_color_patch': ret_fine['blended_color_patch'],
+            'blended_color_patch_mask': ret_fine['blended_color_patch_mask'],
+            'weights_sum_fg': ret_fine['weights_sum_fg']
+        }
+
+        return result
+
+    @torch.no_grad()
+    def sample_z_vals_from_sdfVolume(self, rays_o, rays_d, near, far, sdf_volume, mask_volume):
+        # ? based on sdf to do importance sampling, seems that too biased on pre-estimation
+        device = rays_o.device
+        N_rays = len(rays_o)
+        n_samples = self.n_samples * 2
+
+        z_vals = torch.linspace(0.0, 1.0, n_samples).to(device)
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        if z_vals.shape[0] == 1:
+            z_vals = z_vals.repeat(N_rays, 1)
+
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+
+        sdf = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), sdf_volume).reshape([N_rays, n_samples])
+
+        new_z_vals = self.up_sample(rays_o, rays_d, z_vals, sdf, self.n_samples,
+                                    200,
+                                    conditional_valid_mask_volume=mask_volume,
+                                    )
+        return new_z_vals
+
+    @torch.no_grad()
+    def sample_z_vals_from_maskVolume(self, rays_o, rays_d, near, far, mask_volume):  # don't use
+        device = rays_o.device
+        N_rays = len(rays_o)
+        n_samples = self.n_samples * 2
+
+        z_vals = torch.linspace(0.0, 1.0, n_samples).to(device)
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        if z_vals.shape[0] == 1:
+            z_vals = z_vals.repeat(N_rays, 1)
+
+        mid_z_vals = (z_vals[:, 1:] + z_vals[:, :-1]) * 0.5
+
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]
+
+        pts_mask = self.get_pts_mask_for_conditional_volume(pts.view(-1, 3), mask_volume).reshape(
+            [N_rays, n_samples - 1])
+
+        # empty voxel set to 0.1, non-empty voxel set to 1
+        weights = torch.where(pts_mask > 0, torch.ones_like(pts_mask).to(device),
+                              0.1 * torch.ones_like(pts_mask).to(device))
+
+        # sample more pts in non-empty voxels
+        z_samples = sample_pdf(z_vals, weights, self.n_samples, det=True).detach()
+        return z_samples
+
+    @torch.no_grad()
+    def filter_pts_by_depthmaps(self, coords, pred_depth_maps, proj_matrices,
+                                partial_vol_origin, voxel_size,
+                                near, far, depth_interval, d_plane_nums):
+        """
+        Use the pred_depthmaps to remove redundant pts (pruned by sdf, sdf always have two sides, the back side is useless)
+        :param coords: [n, 3]  int coords
+        :param pred_depth_maps: [N_views, 1, h, w]
+        :param proj_matrices: [N_views, 4, 4]
+        :param partial_vol_origin: [3]
+        :param voxel_size: 1
+        :param near: 1
+        :param far: 1
+        :param depth_interval: 1
+        :param d_plane_nums: 1
+        :return:
+        """
+        device = pred_depth_maps.device
+        n_views, _, sizeH, sizeW = pred_depth_maps.shape
+
+        if len(partial_vol_origin.shape) == 1:
+            partial_vol_origin = partial_vol_origin[None, :]
+        pts = coords * voxel_size + partial_vol_origin
+
+        rs_grid = pts.unsqueeze(0).expand(n_views, -1, -1)
+        rs_grid = rs_grid.permute(0, 2, 1).contiguous()  # [n_views, 3, n_pts]
+        nV = rs_grid.shape[-1]
+        rs_grid = torch.cat([rs_grid, torch.ones([n_views, 1, nV]).to(device)], dim=1)  # [n_views, 4, n_pts]
+
+        # Project grid
+        im_p = proj_matrices @ rs_grid  # - transform world pts to image UV space   # [n_views, 4, n_pts]
+        im_x, im_y, im_z = im_p[:, 0], im_p[:, 1], im_p[:, 2]
+        im_x = im_x / im_z
+        im_y = im_y / im_z
+
+        im_grid = torch.stack([2 * im_x / (sizeW - 1) - 1, 2 * im_y / (sizeH - 1) - 1], dim=-1)
+
+        im_grid = im_grid.view(n_views, 1, -1, 2)
+        sampled_depths = torch.nn.functional.grid_sample(pred_depth_maps, im_grid, mode='bilinear',
+                                                         padding_mode='zeros',
+                                                         align_corners=True)[:, 0, 0, :]  # [n_views, n_pts]
+        sampled_depths_valid = (sampled_depths > 0.5 * near).float()
+        valid_d_min = (sampled_depths - d_plane_nums * depth_interval).clamp(near.item(),
+                                                                             far.item()) * sampled_depths_valid
+        valid_d_max = (sampled_depths + d_plane_nums * depth_interval).clamp(near.item(),
+                                                                             far.item()) * sampled_depths_valid
+
+        mask = im_grid.abs() <= 1
+        mask = mask[:, 0]  # [n_views, n_pts, 2]
+        mask = (mask.sum(dim=-1) == 2) & (im_z > valid_d_min) & (im_z < valid_d_max)
+
+        mask = mask.view(n_views, -1)
+        mask = mask.permute(1, 0).contiguous()  # [num_pts, nviews]
+
+        mask_final = torch.sum(mask.float(), dim=1, keepdim=False) > 0
+
+        return mask_final
+
+    @torch.no_grad()
+    def get_valid_sparse_coords_by_sdf_depthfilter(self, sdf_volume, coords_volume, mask_volume, feature_volume,
+                                                   pred_depth_maps, proj_matrices,
+                                                   partial_vol_origin, voxel_size,
+                                                   near, far, depth_interval, d_plane_nums,
+                                                   threshold=0.02, maximum_pts=110000):
+        """
+        assume batch size == 1, from the first lod to get sparse voxels
+        :param sdf_volume: [1, X, Y, Z]
+        :param coords_volume: [3, X, Y, Z]
+        :param mask_volume: [1, X, Y, Z]
+        :param feature_volume: [C, X, Y, Z]
+        :param threshold:
+        :return:
+        """
+        device = coords_volume.device
+        _, dX, dY, dZ = coords_volume.shape
+
+        def prune(sdf_pts, coords_pts, mask_volume, threshold):
+            occupancy_mask = (torch.abs(sdf_pts) < threshold).squeeze(1)  # [num_pts]
+            valid_coords = coords_pts[occupancy_mask]
+
+            # - filter backside surface by depth maps
+            mask_filtered = self.filter_pts_by_depthmaps(valid_coords, pred_depth_maps, proj_matrices,
+                                                         partial_vol_origin, voxel_size,
+                                                         near, far, depth_interval, d_plane_nums)
+            valid_coords = valid_coords[mask_filtered]
+
+            # - dilate
+            occupancy_mask = sparse_to_dense_channel(valid_coords, 1, [dX, dY, dZ], 1, 0, device)  # [dX, dY, dZ, 1]
+
+            # - dilate
+            occupancy_mask = occupancy_mask.float()
+            occupancy_mask = occupancy_mask.view(1, 1, dX, dY, dZ)
+            occupancy_mask = F.avg_pool3d(occupancy_mask, kernel_size=7, stride=1, padding=3)
+            occupancy_mask = occupancy_mask.view(-1, 1) > 0
+
+            final_mask = torch.logical_and(mask_volume, occupancy_mask)[:, 0]  # [num_pts]
+
+            return final_mask, torch.sum(final_mask.float())
+
+        C, dX, dY, dZ = feature_volume.shape
+        sdf_volume = sdf_volume.permute(1, 2, 3, 0).contiguous().view(-1, 1)
+        coords_volume = coords_volume.permute(1, 2, 3, 0).contiguous().view(-1, 3)
+        mask_volume = mask_volume.permute(1, 2, 3, 0).contiguous().view(-1, 1)
+        feature_volume = feature_volume.permute(1, 2, 3, 0).contiguous().view(-1, C)
+
+        # - for check
+        # sdf_volume = torch.rand_like(sdf_volume).float().to(sdf_volume.device) * 0.02
+
+        final_mask, valid_num = prune(sdf_volume, coords_volume, mask_volume, threshold)
+
+        while (valid_num > maximum_pts) and (threshold > 0.003):
+            threshold = threshold - 0.002
+            final_mask, valid_num = prune(sdf_volume, coords_volume, mask_volume, threshold)
+
+        valid_coords = coords_volume[final_mask]  # [N, 3]
+        valid_feature = feature_volume[final_mask]  # [N, C]
+
+        valid_coords = torch.cat([torch.ones([valid_coords.shape[0], 1]).to(valid_coords.device) * 0,
+                                  valid_coords], dim=1)  # [N, 4], append batch idx
+
+        # ! if the valid_num is still larger than maximum_pts, sample part of pts
+        if valid_num > maximum_pts:
+            valid_num = valid_num.long()
+            occupancy = torch.ones([valid_num]).to(device) > 0
+            choice = np.random.choice(valid_num.cpu().numpy(), valid_num.cpu().numpy() - maximum_pts,
+                                      replace=False)
+            ind = torch.nonzero(occupancy).to(device)
+            occupancy[ind[choice]] = False
+            valid_coords = valid_coords[occupancy]
+            valid_feature = valid_feature[occupancy]
+
+            print(threshold, "randomly sample to save memory")
+
+        return valid_coords, valid_feature
+
+    @torch.no_grad()
+    def get_valid_sparse_coords_by_sdf(self, sdf_volume, coords_volume, mask_volume, feature_volume, threshold=0.02,
+                                       maximum_pts=110000):
+        """
+        assume batch size == 1, from the first lod to get sparse voxels
+        :param sdf_volume: [num_pts, 1]
+        :param coords_volume: [3, X, Y, Z]
+        :param mask_volume: [1, X, Y, Z]
+        :param feature_volume: [C, X, Y, Z]
+        :param threshold:
+        :return:
+        """
+
+        def prune(sdf_volume, mask_volume, threshold):
+            occupancy_mask = torch.abs(sdf_volume) < threshold  # [num_pts, 1]
+
+            # - dilate
+            occupancy_mask = occupancy_mask.float()
+            occupancy_mask = occupancy_mask.view(1, 1, dX, dY, dZ)
+            occupancy_mask = F.avg_pool3d(occupancy_mask, kernel_size=7, stride=1, padding=3)
+            occupancy_mask = occupancy_mask.view(-1, 1) > 0
+
+            final_mask = torch.logical_and(mask_volume, occupancy_mask)[:, 0]  # [num_pts]
+
+            return final_mask, torch.sum(final_mask.float())
+
+        C, dX, dY, dZ = feature_volume.shape
+        coords_volume = coords_volume.permute(1, 2, 3, 0).contiguous().view(-1, 3)
+        mask_volume = mask_volume.permute(1, 2, 3, 0).contiguous().view(-1, 1)
+        feature_volume = feature_volume.permute(1, 2, 3, 0).contiguous().view(-1, C)
+
+        final_mask, valid_num = prune(sdf_volume, mask_volume, threshold)
+
+        while (valid_num > maximum_pts) and (threshold > 0.003):
+            threshold = threshold - 0.002
+            final_mask, valid_num = prune(sdf_volume, mask_volume, threshold)
+
+        valid_coords = coords_volume[final_mask]  # [N, 3]
+        valid_feature = feature_volume[final_mask]  # [N, C]
+
+        valid_coords = torch.cat([torch.ones([valid_coords.shape[0], 1]).to(valid_coords.device) * 0,
+                                  valid_coords], dim=1)  # [N, 4], append batch idx
+
+        # ! if the valid_num is still larger than maximum_pts, sample part of pts
+        if valid_num > maximum_pts:
+            device = sdf_volume.device
+            valid_num = valid_num.long()
+            occupancy = torch.ones([valid_num]).to(device) > 0
+            choice = np.random.choice(valid_num.cpu().numpy(), valid_num.cpu().numpy() - maximum_pts,
+                                      replace=False)
+            ind = torch.nonzero(occupancy).to(device)
+            occupancy[ind[choice]] = False
+            valid_coords = valid_coords[occupancy]
+            valid_feature = valid_feature[occupancy]
+
+            print(threshold, "randomly sample to save memory")
+
+        return valid_coords, valid_feature
+
+    @torch.no_grad()
+    def extract_fields(self, bound_min, bound_max, resolution, query_func, device,
+                       # * related to conditional feature
+                       **kwargs
+                       ):
+        N = 64
+        X = torch.linspace(bound_min[0], bound_max[0], resolution).to(device).split(N)
+        Y = torch.linspace(bound_min[1], bound_max[1], resolution).to(device).split(N)
+        Z = torch.linspace(bound_min[2], bound_max[2], resolution).to(device).split(N)
+
+        u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+        with torch.no_grad():
+            for xi, xs in enumerate(X):
+                for yi, ys in enumerate(Y):
+                    for zi, zs in enumerate(Z):
+                        xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing="ij")
+                        pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
+
+                        # ! attention, the query function is different for extract geometry and fields
+                        output = query_func(pts, **kwargs)
+                        sdf = output['sdf_pts_scale%d' % kwargs['lod']].reshape(len(xs), len(ys),
+                                                                                len(zs)).detach().cpu().numpy()
+
+                        u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = -1 * sdf
+        return u
+
+    @torch.no_grad()
+    def extract_geometry(self, sdf_network, bound_min, bound_max, resolution, threshold, device, occupancy_mask=None,
+                         # * 3d feature volume
+                         **kwargs
+                         ):
+        # logging.info('threshold: {}'.format(threshold))
+
+        u = self.extract_fields(bound_min, bound_max, resolution,
+                                lambda pts, **kwargs: sdf_network.sdf(pts, **kwargs),
+                                # - sdf need to be multiplied by -1
+                                device,
+                                # * 3d feature volume
+                                **kwargs
+                                )
+        if occupancy_mask is not None:
+            dX, dY, dZ = occupancy_mask.shape
+            empty_mask = 1 - occupancy_mask
+            empty_mask = empty_mask.view(1, 1, dX, dY, dZ)
+            # - dilation
+            # empty_mask = F.avg_pool3d(empty_mask, kernel_size=7, stride=1, padding=3)
+            empty_mask = F.interpolate(empty_mask, [resolution, resolution, resolution], mode='nearest')
+            empty_mask = empty_mask.view(resolution, resolution, resolution).cpu().numpy() > 0
+            u[empty_mask] = -100
+            del empty_mask
+
+        vertices, triangles = mcubes.marching_cubes(u, threshold)
+        b_max_np = bound_max.detach().cpu().numpy()
+        b_min_np = bound_min.detach().cpu().numpy()
+
+        vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+        return vertices, triangles, u
+
+    @torch.no_grad()
+    def extract_depth_maps(self, sdf_network, con_volume, intrinsics, c2ws, H, W, near, far):
+        """
+        extract depth maps from the density volume
+        :param con_volume: [1, 1+C, dX, dY, dZ]  can by con_volume or sdf_volume
+        :param c2ws: [B, 4, 4]
+        :param H:
+        :param W:
+        :param near:
+        :param far:
+        :return:
+        """
+        device = con_volume.device
+        batch_size = intrinsics.shape[0]
+
+        with torch.no_grad():
+            ys, xs = torch.meshgrid(torch.linspace(0, H - 1, H),
+                                    torch.linspace(0, W - 1, W), indexing="ij")  # pytorch's meshgrid has indexing='ij'
+            p = torch.stack([xs, ys, torch.ones_like(ys)], dim=-1)  # H, W, 3
+
+            intrinsics_inv = torch.inverse(intrinsics)
+
+            p = p.view(-1, 3).float().to(device)  # N_rays, 3
+            p = torch.matmul(intrinsics_inv[:, None, :3, :3], p[:, :, None]).squeeze()  # Batch, N_rays, 3
+            rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # Batch, N_rays, 3
+            rays_v = torch.matmul(c2ws[:, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # Batch, N_rays, 3
+            rays_o = c2ws[:, None, :3, 3].expand(rays_v.shape)  # Batch, N_rays, 3
+            rays_d = rays_v
+
+        rays_o = rays_o.contiguous().view(-1, 3)
+        rays_d = rays_d.contiguous().view(-1, 3)
+
+        ################## - sphere tracer to extract depth maps               ######################
+        depth_masks_sphere, depth_maps_sphere = self.ray_tracer.extract_depth_maps(
+            rays_o, rays_d,
+            near[None, :].repeat(rays_o.shape[0], 1),
+            far[None, :].repeat(rays_o.shape[0], 1),
+            sdf_network, con_volume
+        )
+
+        depth_maps = depth_maps_sphere.view(batch_size, 1, H, W)
+        depth_masks = depth_masks_sphere.view(batch_size, 1, H, W)
+
+        depth_maps = torch.where(depth_masks, depth_maps,
+                                 torch.zeros_like(depth_masks.float()).to(device))  # fill invalid pixels by 0
+
+        return depth_maps, depth_masks

SparseNeuS_demo_v1/models/sparse_sdf_network.py

@@ -0,0 +1,907 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsparse.tensor import PointTensor, SparseTensor
+import torchsparse.nn as spnn
+
+from tsparse.modules import SparseCostRegNet
+from tsparse.torchsparse_utils import sparse_to_dense_channel
+from ops.grid_sampler import grid_sample_3d, tricubic_sample_3d
+
+# from .gru_fusion import GRUFusion
+from ops.back_project import back_project_sparse_type
+from ops.generate_grids import generate_grid
+
+from inplace_abn import InPlaceABN
+
+from models.embedder import Embedding
+from models.featurenet import ConvBnReLU
+
+import pdb
+import random
+
+torch._C._jit_set_profiling_executor(False)
+torch._C._jit_set_profiling_mode(False)
+
+
+@torch.jit.script
+def fused_mean_variance(x, weight):
+    mean = torch.sum(x * weight, dim=1, keepdim=True)
+    var = torch.sum(weight * (x - mean) ** 2, dim=1, keepdim=True)
+    return mean, var
+
+
+class LatentSDFLayer(nn.Module):
+    def __init__(self,
+                 d_in=3,
+                 d_out=129,
+                 d_hidden=128,
+                 n_layers=4,
+                 skip_in=(4,),
+                 multires=0,
+                 bias=0.5,
+                 geometric_init=True,
+                 weight_norm=True,
+                 activation='softplus',
+                 d_conditional_feature=16):
+        super(LatentSDFLayer, self).__init__()
+
+        self.d_conditional_feature = d_conditional_feature
+
+        # concat latent code for ench layer input excepting the first layer and the last layer
+        dims_in = [d_in] + [d_hidden + d_conditional_feature for _ in range(n_layers - 2)] + [d_hidden]
+        dims_out = [d_hidden for _ in range(n_layers - 1)] + [d_out]
+
+        self.embed_fn_fine = None
+
+        if multires > 0:
+            embed_fn = Embedding(in_channels=d_in, N_freqs=multires)  # * include the input
+            self.embed_fn_fine = embed_fn
+            dims_in[0] = embed_fn.out_channels
+
+        self.num_layers = n_layers
+        self.skip_in = skip_in
+
+        for l in range(0, self.num_layers - 1):
+            if l in self.skip_in:
+                in_dim = dims_in[l] + dims_in[0]
+            else:
+                in_dim = dims_in[l]
+
+            out_dim = dims_out[l]
+            lin = nn.Linear(in_dim, out_dim)
+
+            if geometric_init:  # - from IDR code,
+                if l == self.num_layers - 2:
+                    torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(in_dim), std=0.0001)
+                    torch.nn.init.constant_(lin.bias, -bias)
+                    # the channels for latent codes are set to 0
+                    torch.nn.init.constant_(lin.weight[:, -d_conditional_feature:], 0.0)
+                    torch.nn.init.constant_(lin.bias[-d_conditional_feature:], 0.0)
+
+                elif multires > 0 and l == 0:  # the first layer
+                    torch.nn.init.constant_(lin.bias, 0.0)
+                    # * the channels for position embeddings are set to 0
+                    torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
+                    # * the channels for the xyz coordinate (3 channels) for initialized by normal distribution
+                    torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))
+                elif 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))
+                    # * the channels for position embeddings (and conditional_feature) are initialized to 0
+                    torch.nn.init.constant_(lin.weight[:, -(dims_in[0] - 3 + d_conditional_feature):], 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))
+                    # the channels for latent code are initialized to 0
+                    torch.nn.init.constant_(lin.weight[:, -d_conditional_feature:], 0.0)
+
+            if weight_norm:
+                lin = nn.utils.weight_norm(lin)
+
+            setattr(self, "lin" + str(l), lin)
+
+        if activation == 'softplus':
+            self.activation = nn.Softplus(beta=100)
+        else:
+            assert activation == 'relu'
+            self.activation = nn.ReLU()
+
+    def forward(self, inputs, latent):
+        inputs = inputs
+        if self.embed_fn_fine is not None:
+            inputs = self.embed_fn_fine(inputs)
+
+        # - only for lod1 network can use the pretrained params of lod0 network
+        if latent.shape[1] != self.d_conditional_feature:
+            latent = torch.cat([latent, latent], dim=1)
+
+        x = inputs
+        for l in range(0, self.num_layers - 1):
+            lin = getattr(self, "lin" + str(l))
+
+            # * due to the conditional bias, different from original neus version
+            if l in self.skip_in:
+                x = torch.cat([x, inputs], 1) / np.sqrt(2)
+
+            if 0 < l < self.num_layers - 1:
+                x = torch.cat([x, latent], 1)
+
+            x = lin(x)
+
+            if l < self.num_layers - 2:
+                x = self.activation(x)
+
+        return x
+
+
+class SparseSdfNetwork(nn.Module):
+    '''
+    Coarse-to-fine sparse cost regularization network
+    return sparse volume feature for extracting sdf
+    '''
+
+    def __init__(self, lod, ch_in, voxel_size, vol_dims,
+                 hidden_dim=128, activation='softplus',
+                 cost_type='variance_mean',
+                 d_pyramid_feature_compress=16,
+                 regnet_d_out=8, num_sdf_layers=4,
+                 multires=6,
+                 ):
+        super(SparseSdfNetwork, self).__init__()
+
+        self.lod = lod  # - gradually training, the current regularization lod
+        self.ch_in = ch_in
+        self.voxel_size = voxel_size  # - the voxel size of the current volume
+        self.vol_dims = torch.tensor(vol_dims)  # - the dims of the current volume
+
+        self.selected_views_num = 2  # the number of selected views for feature aggregation
+        self.hidden_dim = hidden_dim
+        self.activation = activation
+        self.cost_type = cost_type
+        self.d_pyramid_feature_compress = d_pyramid_feature_compress
+        self.gru_fusion = None
+
+        self.regnet_d_out = regnet_d_out
+        self.multires = multires
+
+        self.pos_embedder = Embedding(3, self.multires)
+
+        self.compress_layer = ConvBnReLU(
+            self.ch_in, self.d_pyramid_feature_compress, 3, 1, 1,
+            norm_act=InPlaceABN)
+        sparse_ch_in = self.d_pyramid_feature_compress * 2
+
+        sparse_ch_in = sparse_ch_in + 16 if self.lod > 0 else sparse_ch_in
+        self.sparse_costreg_net = SparseCostRegNet(
+            d_in=sparse_ch_in, d_out=self.regnet_d_out)
+        # self.regnet_d_out = self.sparse_costreg_net.d_out
+
+        if activation == 'softplus':
+            self.activation = nn.Softplus(beta=100)
+        else:
+            assert activation == 'relu'
+            self.activation = nn.ReLU()
+
+        self.sdf_layer = LatentSDFLayer(d_in=3,
+                                        d_out=self.hidden_dim + 1,
+                                        d_hidden=self.hidden_dim,
+                                        n_layers=num_sdf_layers,
+                                        multires=multires,
+                                        geometric_init=True,
+                                        weight_norm=True,
+                                        activation=activation,
+                                        d_conditional_feature=16  # self.regnet_d_out
+                                        )
+
+    def upsample(self, pre_feat, pre_coords, interval, num=8):
+        '''
+
+        :param pre_feat: (Tensor), features from last level, (N, C)
+        :param pre_coords: (Tensor), coordinates from last level, (N, 4) (4 : Batch ind, x, y, z)
+        :param interval: interval of voxels, interval = scale ** 2
+        :param num: 1 -> 8
+        :return: up_feat : (Tensor), upsampled features, (N*8, C)
+        :return: up_coords: (N*8, 4), upsampled coordinates, (4 : Batch ind, x, y, z)
+        '''
+        with torch.no_grad():
+            pos_list = [1, 2, 3, [1, 2], [1, 3], [2, 3], [1, 2, 3]]
+            n, c = pre_feat.shape
+            up_feat = pre_feat.unsqueeze(1).expand(-1, num, -1).contiguous()
+            up_coords = pre_coords.unsqueeze(1).repeat(1, num, 1).contiguous()
+            for i in range(num - 1):
+                up_coords[:, i + 1, pos_list[i]] += interval
+
+            up_feat = up_feat.view(-1, c)
+            up_coords = up_coords.view(-1, 4)
+
+        return up_feat, up_coords
+
+    def aggregate_multiview_features(self, multiview_features, multiview_masks):
+        """
+        aggregate mutli-view features by compute their cost variance
+        :param multiview_features: (num of voxels, num_of_views, c)
+        :param multiview_masks: (num of voxels, num_of_views)
+        :return:
+        """
+        num_pts, n_views, C = multiview_features.shape
+
+        counts = torch.sum(multiview_masks, dim=1, keepdim=False)  # [num_pts]
+
+        assert torch.all(counts > 0)  # the point is visible for at least 1 view
+
+        volume_sum = torch.sum(multiview_features, dim=1, keepdim=False)  # [num_pts, C]
+        volume_sq_sum = torch.sum(multiview_features ** 2, dim=1, keepdim=False)
+
+        if volume_sum.isnan().sum() > 0:
+            import ipdb; ipdb.set_trace()
+
+        del multiview_features
+
+        counts = 1. / (counts + 1e-5)
+        costvar = volume_sq_sum * counts[:, None] - (volume_sum * counts[:, None]) ** 2
+
+        costvar_mean = torch.cat([costvar, volume_sum * counts[:, None]], dim=1)
+        del volume_sum, volume_sq_sum, counts
+
+
+
+        return costvar_mean
+
+    def sparse_to_dense_volume(self, coords, feature, vol_dims, interval, device=None):
+        """
+        convert the sparse volume into dense volume to enable trilinear sampling
+        to save GPU memory;
+        :param coords: [num_pts, 3]
+        :param feature: [num_pts, C]
+        :param vol_dims: [3]  dX, dY, dZ
+        :param interval:
+        :return:
+        """
+
+        # * assume batch size is 1
+        if device is None:
+            device = feature.device
+
+        coords_int = (coords / interval).to(torch.int64)
+        vol_dims = (vol_dims / interval).to(torch.int64)
+
+        # - if stored in CPU, too slow
+        dense_volume = sparse_to_dense_channel(
+            coords_int.to(device), feature.to(device), vol_dims.to(device),
+            feature.shape[1], 0, device)  # [X, Y, Z, C]
+
+        valid_mask_volume = sparse_to_dense_channel(
+            coords_int.to(device),
+            torch.ones([feature.shape[0], 1]).to(feature.device),
+            vol_dims.to(device),
+            1, 0, device)  # [X, Y, Z, 1]
+
+        dense_volume = dense_volume.permute(3, 0, 1, 2).contiguous().unsqueeze(0)  # [1, C, X, Y, Z]
+        valid_mask_volume = valid_mask_volume.permute(3, 0, 1, 2).contiguous().unsqueeze(0)  # [1, 1, X, Y, Z]
+
+        return dense_volume, valid_mask_volume
+
+    def get_conditional_volume(self, feature_maps, partial_vol_origin, proj_mats, sizeH=None, sizeW=None, lod=0,
+                               pre_coords=None, pre_feats=None,
+                               ):
+        """
+
+        :param feature_maps: pyramid features (B,V,C0+C1+C2,H,W) fused pyramid features
+        :param partial_vol_origin: [B, 3]  the world coordinates of the volume origin (0,0,0)
+        :param proj_mats: projection matrix transform world pts into image space [B,V,4,4] suitable for original image size
+        :param sizeH: the H of original image size
+        :param sizeW: the W of original image size
+        :param pre_coords: the coordinates of sparse volume from the prior lod
+        :param pre_feats: the features of sparse volume from the prior lod
+        :return:
+        """
+        device = proj_mats.device
+        bs = feature_maps.shape[0]
+        N_views = feature_maps.shape[1]
+        minimum_visible_views = np.min([1, N_views - 1])
+        # import ipdb; ipdb.set_trace()
+        outputs = {}
+        pts_samples = []
+
+        # ----coarse to fine----
+
+        # * use fused pyramid feature maps are very important
+        if self.compress_layer is not None:
+            feats = self.compress_layer(feature_maps[0])
+        else:
+            feats = feature_maps[0]
+        feats = feats[:, None, :, :, :]  # [V, B, C, H, W]
+        KRcam = proj_mats.permute(1, 0, 2, 3).contiguous()  # [V, B, 4, 4]
+        interval = 1
+
+        if self.lod == 0:
+            # ----generate new coords----
+            coords = generate_grid(self.vol_dims, 1)[0]
+            coords = coords.view(3, -1).to(device)  # [3, num_pts]
+            up_coords = []
+            for b in range(bs):
+                up_coords.append(torch.cat([torch.ones(1, coords.shape[-1]).to(coords.device) * b, coords]))
+            up_coords = torch.cat(up_coords, dim=1).permute(1, 0).contiguous()
+            # * since we only estimate the geometry of input reference image at one time;
+            # * mask the outside of the camera frustum
+            # import ipdb; ipdb.set_trace()
+            frustum_mask = back_project_sparse_type(
+                up_coords, partial_vol_origin, self.voxel_size,
+                feats, KRcam, sizeH=sizeH, sizeW=sizeW, only_mask=True)  # [num_pts, n_views]
+            frustum_mask = torch.sum(frustum_mask, dim=-1) > minimum_visible_views  # ! here should be large
+            up_coords = up_coords[frustum_mask]  # [num_pts_valid, 4]
+
+        else:
+            # ----upsample coords----
+            assert pre_feats is not None
+            assert pre_coords is not None
+            up_feat, up_coords = self.upsample(pre_feats, pre_coords, 1)
+
+        # ----back project----
+        # give each valid 3d grid point all valid 2D features and masks
+        multiview_features, multiview_masks = back_project_sparse_type(
+            up_coords, partial_vol_origin, self.voxel_size, feats,
+            KRcam, sizeH=sizeH, sizeW=sizeW)  # (num of voxels, num_of_views, c), (num of voxels, num_of_views)
+                                              # num_of_views = all views
+        
+        # if multiview_features.isnan().sum() > 0:
+        #     import ipdb; ipdb.set_trace()
+
+        # import ipdb; ipdb.set_trace()
+        if self.lod > 0:
+            # ! need another invalid voxels filtering
+            frustum_mask = torch.sum(multiview_masks, dim=-1) > 1
+            up_feat = up_feat[frustum_mask]
+            up_coords = up_coords[frustum_mask]
+            multiview_features = multiview_features[frustum_mask]
+            multiview_masks = multiview_masks[frustum_mask]
+        # if multiview_features.isnan().sum() > 0:
+        #     import ipdb; ipdb.set_trace()
+        volume = self.aggregate_multiview_features(multiview_features, multiview_masks) # compute variance for all images features
+        # import ipdb; ipdb.set_trace()
+
+        # if volume.isnan().sum() > 0:
+        #     import ipdb; ipdb.set_trace()
+
+        del multiview_features, multiview_masks
+
+        # ----concat feature from last stage----
+        if self.lod != 0:
+            feat = torch.cat([volume, up_feat], dim=1)
+        else:
+            feat = volume
+
+        # batch index is in the last position
+        r_coords = up_coords[:, [1, 2, 3, 0]]
+
+        # if feat.isnan().sum() > 0:
+        #     print('feat has nan:', feat.isnan().sum())
+        #     import ipdb; ipdb.set_trace()
+
+        sparse_feat = SparseTensor(feat, r_coords.to(
+            torch.int32))  # - directly use sparse tensor to avoid point2voxel operations
+        # import ipdb; ipdb.set_trace()
+        feat = self.sparse_costreg_net(sparse_feat)
+
+        dense_volume, valid_mask_volume = self.sparse_to_dense_volume(up_coords[:, 1:], feat, self.vol_dims, interval,
+                                                                      device=None)  # [1, C/1, X, Y, Z]
+
+        # if dense_volume.isnan().sum() > 0:
+        #     import ipdb; ipdb.set_trace()
+
+
+        outputs['dense_volume_scale%d' % self.lod] = dense_volume # [1, 16, 96, 96, 96]
+        outputs['valid_mask_volume_scale%d' % self.lod] = valid_mask_volume # [1, 1, 96, 96, 96]
+        outputs['visible_mask_scale%d' % self.lod] = valid_mask_volume # [1, 1, 96, 96, 96]
+        outputs['coords_scale%d' % self.lod] = generate_grid(self.vol_dims, interval).to(device)
+        # import ipdb; ipdb.set_trace()
+        return outputs
+
+    def sdf(self, pts, conditional_volume, lod):
+        num_pts = pts.shape[0]
+        device = pts.device
+        pts_ = pts.clone()
+        pts = pts.view(1, 1, 1, num_pts, 3)  # - should be in range (-1, 1)
+
+        pts = torch.flip(pts, dims=[-1])
+        # import ipdb; ipdb.set_trace()
+        sampled_feature = grid_sample_3d(conditional_volume, pts)  # [1, c, 1, 1, num_pts]
+        sampled_feature = sampled_feature.view(-1, num_pts).permute(1, 0).contiguous().to(device)
+
+        sdf_pts = self.sdf_layer(pts_, sampled_feature)
+
+        outputs = {}
+        outputs['sdf_pts_scale%d' % lod] = sdf_pts[:, :1]
+        outputs['sdf_features_pts_scale%d' % lod] = sdf_pts[:, 1:]
+        outputs['sampled_latent_scale%d' % lod] = sampled_feature
+
+        return outputs
+
+    @torch.no_grad()
+    def sdf_from_sdfvolume(self, pts, sdf_volume, lod=0):
+        num_pts = pts.shape[0]
+        device = pts.device
+        pts_ = pts.clone()
+        pts = pts.view(1, 1, 1, num_pts, 3)  # - should be in range (-1, 1)
+
+        pts = torch.flip(pts, dims=[-1])
+
+        sdf = torch.nn.functional.grid_sample(sdf_volume, pts, mode='bilinear', align_corners=True,
+                                              padding_mode='border')
+        sdf = sdf.view(-1, num_pts).permute(1, 0).contiguous().to(device)
+
+        outputs = {}
+        outputs['sdf_pts_scale%d' % lod] = sdf
+
+        return outputs
+
+    @torch.no_grad()
+    def get_sdf_volume(self, conditional_volume, mask_volume, coords_volume, partial_origin):
+        """
+
+        :param conditional_volume: [1,C, dX,dY,dZ]
+        :param mask_volume: [1,1, dX,dY,dZ]
+        :param coords_volume: [1,3, dX,dY,dZ]
+        :return:
+        """
+        device = conditional_volume.device
+        chunk_size = 10240
+
+        _, C, dX, dY, dZ = conditional_volume.shape
+        conditional_volume = conditional_volume.view(C, dX * dY * dZ).permute(1, 0).contiguous()
+        mask_volume = mask_volume.view(-1)
+        coords_volume = coords_volume.view(3, dX * dY * dZ).permute(1, 0).contiguous()
+
+        pts = coords_volume * self.voxel_size + partial_origin  # [dX*dY*dZ, 3]
+
+        sdf_volume = torch.ones([dX * dY * dZ, 1]).float().to(device)
+
+        conditional_volume = conditional_volume[mask_volume > 0]
+        pts = pts[mask_volume > 0]
+        conditional_volume = conditional_volume.split(chunk_size)
+        pts = pts.split(chunk_size)
+
+        sdf_all = []
+        for pts_part, feature_part in zip(pts, conditional_volume):
+            sdf_part = self.sdf_layer(pts_part, feature_part)[:, :1]
+            sdf_all.append(sdf_part)
+
+        sdf_all = torch.cat(sdf_all, dim=0)
+        sdf_volume[mask_volume > 0] = sdf_all
+        sdf_volume = sdf_volume.view(1, 1, dX, dY, dZ)
+        return sdf_volume
+
+    def gradient(self, x, conditional_volume, lod):
+        """
+        return the gradient of specific lod
+        :param x:
+        :param lod:
+        :return:
+        """
+        x.requires_grad_(True)
+        # import ipdb; ipdb.set_trace()
+        with torch.enable_grad():
+            output = self.sdf(x, conditional_volume, lod)
+        y = output['sdf_pts_scale%d' % lod]
+
+        d_output = torch.ones_like(y, requires_grad=False, device=y.device)
+        # ! Distributed Data Parallel doesn’t work with torch.autograd.grad()
+        # ! (i.e. it will only work if gradients are to be accumulated in .grad attributes of parameters).
+        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)
+
+
+def sparse_to_dense_volume(coords, feature, vol_dims, interval, device=None):
+    """
+    convert the sparse volume into dense volume to enable trilinear sampling
+    to save GPU memory;
+    :param coords: [num_pts, 3]
+    :param feature: [num_pts, C]
+    :param vol_dims: [3]  dX, dY, dZ
+    :param interval:
+    :return:
+    """
+
+    # * assume batch size is 1
+    if device is None:
+        device = feature.device
+
+    coords_int = (coords / interval).to(torch.int64)
+    vol_dims = (vol_dims / interval).to(torch.int64)
+
+    # - if stored in CPU, too slow
+    dense_volume = sparse_to_dense_channel(
+        coords_int.to(device), feature.to(device), vol_dims.to(device),
+        feature.shape[1], 0, device)  # [X, Y, Z, C]
+
+    valid_mask_volume = sparse_to_dense_channel(
+        coords_int.to(device),
+        torch.ones([feature.shape[0], 1]).to(feature.device),
+        vol_dims.to(device),
+        1, 0, device)  # [X, Y, Z, 1]
+
+    dense_volume = dense_volume.permute(3, 0, 1, 2).contiguous().unsqueeze(0)  # [1, C, X, Y, Z]
+    valid_mask_volume = valid_mask_volume.permute(3, 0, 1, 2).contiguous().unsqueeze(0)  # [1, 1, X, Y, Z]
+
+    return dense_volume, valid_mask_volume
+
+
+class SdfVolume(nn.Module):
+    def __init__(self, volume, coords=None, type='dense'):
+        super(SdfVolume, self).__init__()
+        self.volume = torch.nn.Parameter(volume, requires_grad=True)
+        self.coords = coords
+        self.type = type
+
+    def forward(self):
+        return self.volume
+
+
+class FinetuneOctreeSdfNetwork(nn.Module):
+    '''
+    After obtain the conditional volume from generalized network;
+    directly optimize the conditional volume
+    The conditional volume is still sparse
+    '''
+
+    def __init__(self, voxel_size, vol_dims,
+                 origin=[-1., -1., -1.],
+                 hidden_dim=128, activation='softplus',
+                 regnet_d_out=8,
+                 multires=6,
+                 if_fitted_rendering=True,
+                 num_sdf_layers=4,
+                 ):
+        super(FinetuneOctreeSdfNetwork, self).__init__()
+
+        self.voxel_size = voxel_size  # - the voxel size of the current volume
+        self.vol_dims = torch.tensor(vol_dims)  # - the dims of the current volume
+
+        self.origin = torch.tensor(origin).to(torch.float32)
+
+        self.hidden_dim = hidden_dim
+        self.activation = activation
+
+        self.regnet_d_out = regnet_d_out
+
+        self.if_fitted_rendering = if_fitted_rendering
+        self.multires = multires
+        # d_in_embedding = self.regnet_d_out if self.pos_add_type == 'latent' else 3
+        # self.pos_embedder = Embedding(d_in_embedding, self.multires)
+
+        # - the optimized parameters
+        self.sparse_volume_lod0 = None
+        self.sparse_coords_lod0 = None
+
+        if activation == 'softplus':
+            self.activation = nn.Softplus(beta=100)
+        else:
+            assert activation == 'relu'
+            self.activation = nn.ReLU()
+
+        self.sdf_layer = LatentSDFLayer(d_in=3,
+                                        d_out=self.hidden_dim + 1,
+                                        d_hidden=self.hidden_dim,
+                                        n_layers=num_sdf_layers,
+                                        multires=multires,
+                                        geometric_init=True,
+                                        weight_norm=True,
+                                        activation=activation,
+                                        d_conditional_feature=16  # self.regnet_d_out
+                                        )
+
+        # - add mlp rendering when finetuning
+        self.renderer = None
+
+        d_in_renderer = 3 + self.regnet_d_out + 3 + 3
+        self.renderer = BlendingRenderingNetwork(
+            d_feature=self.hidden_dim - 1,
+            mode='idr',  # ! the view direction influence a lot
+            d_in=d_in_renderer,
+            d_out=50,  # maximum 50 images
+            d_hidden=self.hidden_dim,
+            n_layers=3,
+            weight_norm=True,
+            multires_view=4,
+            squeeze_out=True,
+        )
+
+    def initialize_conditional_volumes(self, dense_volume_lod0, dense_volume_mask_lod0,
+                                       sparse_volume_lod0=None, sparse_coords_lod0=None):
+        """
+
+        :param dense_volume_lod0: [1,C,dX,dY,dZ]
+        :param dense_volume_mask_lod0: [1,1,dX,dY,dZ]
+        :param dense_volume_lod1:
+        :param dense_volume_mask_lod1:
+        :return:
+        """
+
+        if sparse_volume_lod0 is None:
+            device = dense_volume_lod0.device
+            _, C, dX, dY, dZ = dense_volume_lod0.shape
+
+            dense_volume_lod0 = dense_volume_lod0.view(C, dX * dY * dZ).permute(1, 0).contiguous()
+            mask_lod0 = dense_volume_mask_lod0.view(dX * dY * dZ) > 0
+
+            self.sparse_volume_lod0 = SdfVolume(dense_volume_lod0[mask_lod0], type='sparse')
+
+            coords = generate_grid(self.vol_dims, 1)[0]  # [3, dX, dY, dZ]
+            coords = coords.view(3, dX * dY * dZ).permute(1, 0).to(device)
+            self.sparse_coords_lod0 = torch.nn.Parameter(coords[mask_lod0], requires_grad=False)
+        else:
+            self.sparse_volume_lod0 = SdfVolume(sparse_volume_lod0, type='sparse')
+            self.sparse_coords_lod0 = torch.nn.Parameter(sparse_coords_lod0, requires_grad=False)
+
+    def get_conditional_volume(self):
+        dense_volume, valid_mask_volume = sparse_to_dense_volume(
+            self.sparse_coords_lod0,
+            self.sparse_volume_lod0(), self.vol_dims, interval=1,
+            device=None)  # [1, C/1, X, Y, Z]
+
+        # valid_mask_volume = self.dense_volume_mask_lod0
+
+        outputs = {}
+        outputs['dense_volume_scale%d' % 0] = dense_volume
+        outputs['valid_mask_volume_scale%d' % 0] = valid_mask_volume
+
+        return outputs
+
+    def tv_regularizer(self):
+        dense_volume, valid_mask_volume = sparse_to_dense_volume(
+            self.sparse_coords_lod0,
+            self.sparse_volume_lod0(), self.vol_dims, interval=1,
+            device=None)  # [1, C/1, X, Y, Z]
+
+        dx = (dense_volume[:, :, 1:, :, :] - dense_volume[:, :, :-1, :, :]) ** 2  # [1, C/1, X-1, Y, Z]
+        dy = (dense_volume[:, :, :, 1:, :] - dense_volume[:, :, :, :-1, :]) ** 2  # [1, C/1, X, Y-1, Z]
+        dz = (dense_volume[:, :, :, :, 1:] - dense_volume[:, :, :, :, :-1]) ** 2  # [1, C/1, X, Y, Z-1]
+
+        tv = dx[:, :, :, :-1, :-1] + dy[:, :, :-1, :, :-1] + dz[:, :, :-1, :-1, :]  # [1, C/1, X-1, Y-1, Z-1]
+
+        mask = valid_mask_volume[:, :, :-1, :-1, :-1] * valid_mask_volume[:, :, 1:, :-1, :-1] * \
+               valid_mask_volume[:, :, :-1, 1:, :-1] * valid_mask_volume[:, :, :-1, :-1, 1:]
+
+        tv = torch.sqrt(tv + 1e-6).mean(dim=1, keepdim=True) * mask
+        # tv = tv.mean(dim=1, keepdim=True) * mask
+
+        assert torch.all(~torch.isnan(tv))
+
+        return torch.mean(tv)
+
+    def sdf(self, pts, conditional_volume, lod):
+
+        outputs = {}
+
+        num_pts = pts.shape[0]
+        device = pts.device
+        pts_ = pts.clone()
+        pts = pts.view(1, 1, 1, num_pts, 3)  # - should be in range (-1, 1)
+
+        pts = torch.flip(pts, dims=[-1])
+
+        sampled_feature = grid_sample_3d(conditional_volume, pts)  # [1, c, 1, 1, num_pts]
+        sampled_feature = sampled_feature.view(-1, num_pts).permute(1, 0).contiguous()
+        outputs['sampled_latent_scale%d' % lod] = sampled_feature
+
+        sdf_pts = self.sdf_layer(pts_, sampled_feature)
+
+        lod = 0
+        outputs['sdf_pts_scale%d' % lod] = sdf_pts[:, :1]
+        outputs['sdf_features_pts_scale%d' % lod] = sdf_pts[:, 1:]
+
+        return outputs
+
+    def color_blend(self, pts, position, normals, view_dirs, feature_vectors, img_index,
+                    pts_pixel_color, pts_pixel_mask, pts_patch_color=None, pts_patch_mask=None):
+
+        return self.renderer(torch.cat([pts, position], dim=-1), normals, view_dirs, feature_vectors,
+                             img_index, pts_pixel_color, pts_pixel_mask,
+                             pts_patch_color=pts_patch_color, pts_patch_mask=pts_patch_mask)
+
+    def gradient(self, x, conditional_volume, lod):
+        """
+        return the gradient of specific lod
+        :param x:
+        :param lod:
+        :return:
+        """
+        x.requires_grad_(True)
+        output = self.sdf(x, conditional_volume, lod)
+        y = output['sdf_pts_scale%d' % 0]
+
+        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)
+
+    @torch.no_grad()
+    def prune_dense_mask(self, threshold=0.02):
+        """
+        Just gradually prune the mask of dense volume to decrease the number of sdf network inference
+        :return:
+        """
+        chunk_size = 10240
+        coords = generate_grid(self.vol_dims_lod0, 1)[0]  # [3, dX, dY, dZ]
+
+        _, dX, dY, dZ = coords.shape
+
+        pts = coords.view(3, -1).permute(1,
+                                         0).contiguous() * self.voxel_size_lod0 + self.origin[None, :]  # [dX*dY*dZ, 3]
+
+        # dense_volume = self.dense_volume_lod0()  # [1,C,dX,dY,dZ]
+        dense_volume, _ = sparse_to_dense_volume(
+            self.sparse_coords_lod0,
+            self.sparse_volume_lod0(), self.vol_dims_lod0, interval=1,
+            device=None)  # [1, C/1, X, Y, Z]
+
+        sdf_volume = torch.ones([dX * dY * dZ, 1]).float().to(dense_volume.device) * 100
+
+        mask = self.dense_volume_mask_lod0.view(-1) > 0
+
+        pts_valid = pts[mask].to(dense_volume.device)
+        feature_valid = dense_volume.view(self.regnet_d_out, -1).permute(1, 0).contiguous()[mask]
+
+        pts_valid = pts_valid.split(chunk_size)
+        feature_valid = feature_valid.split(chunk_size)
+
+        sdf_list = []
+
+        for pts_part, feature_part in zip(pts_valid, feature_valid):
+            sdf_part = self.sdf_layer(pts_part, feature_part)[:, :1]
+            sdf_list.append(sdf_part)
+
+        sdf_list = torch.cat(sdf_list, dim=0)
+
+        sdf_volume[mask] = sdf_list
+
+        occupancy_mask = torch.abs(sdf_volume) < threshold  # [num_pts, 1]
+
+        # - dilate
+        occupancy_mask = occupancy_mask.float()
+        occupancy_mask = occupancy_mask.view(1, 1, dX, dY, dZ)
+        occupancy_mask = F.avg_pool3d(occupancy_mask, kernel_size=7, stride=1, padding=3)
+        occupancy_mask = occupancy_mask > 0
+
+        self.dense_volume_mask_lod0 = torch.logical_and(self.dense_volume_mask_lod0,
+                                                        occupancy_mask).float()  # (1, 1, dX, dY, dZ)
+
+
+class BlendingRenderingNetwork(nn.Module):
+    def __init__(
+            self,
+            d_feature,
+            mode,
+            d_in,
+            d_out,
+            d_hidden,
+            n_layers,
+            weight_norm=True,
+            multires_view=0,
+            squeeze_out=True,
+    ):
+        super(BlendingRenderingNetwork, self).__init__()
+
+        self.mode = mode
+        self.squeeze_out = squeeze_out
+        dims = [d_in + d_feature] + [d_hidden for _ in range(n_layers)] + [d_out]
+
+        self.embedder = None
+        if multires_view > 0:
+            self.embedder = Embedding(3, multires_view)
+            dims[0] += (self.embedder.out_channels - 3)
+
+        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 weight_norm:
+                lin = nn.utils.weight_norm(lin)
+
+            setattr(self, "lin" + str(l), lin)
+
+        self.relu = nn.ReLU()
+
+        self.color_volume = None
+
+        self.softmax = nn.Softmax(dim=1)
+
+        self.type = 'blending'
+
+    def sample_pts_from_colorVolume(self, pts):
+        device = pts.device
+        num_pts = pts.shape[0]
+        pts_ = pts.clone()
+        pts = pts.view(1, 1, 1, num_pts, 3)  # - should be in range (-1, 1)
+
+        pts = torch.flip(pts, dims=[-1])
+
+        sampled_color = grid_sample_3d(self.color_volume, pts)  # [1, c, 1, 1, num_pts]
+        sampled_color = sampled_color.view(-1, num_pts).permute(1, 0).contiguous().to(device)
+
+        return sampled_color
+
+    def forward(self, position, normals, view_dirs, feature_vectors, img_index,
+                pts_pixel_color, pts_pixel_mask, pts_patch_color=None, pts_patch_mask=None):
+        """
+
+        :param position: can be 3d coord or interpolated volume latent
+        :param normals:
+        :param view_dirs:
+        :param feature_vectors:
+        :param img_index: [N_views], used to extract corresponding weights
+        :param pts_pixel_color: [N_pts, N_views, 3]
+        :param pts_pixel_mask: [N_pts, N_views]
+        :param pts_patch_color: [N_pts, N_views, Npx, 3]
+        :return:
+        """
+        if self.embedder is not None:
+            view_dirs = self.embedder(view_dirs)
+
+        rendering_input = None
+
+        if self.mode == 'idr':
+            rendering_input = torch.cat([position, view_dirs, normals, feature_vectors], dim=-1)
+        elif self.mode == 'no_view_dir':
+            rendering_input = torch.cat([position, normals, feature_vectors], dim=-1)
+        elif self.mode == 'no_normal':
+            rendering_input = torch.cat([position, view_dirs, feature_vectors], dim=-1)
+        elif self.mode == 'no_points':
+            rendering_input = torch.cat([view_dirs, normals, feature_vectors], dim=-1)
+        elif self.mode == 'no_points_no_view_dir':
+            rendering_input = torch.cat([normals, feature_vectors], dim=-1)
+
+        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)  # [n_pts, d_out]
+
+        ## extract value based on img_index
+        x_extracted = torch.index_select(x, 1, img_index.long())
+
+        weights_pixel = self.softmax(x_extracted)  # [n_pts, N_views]
+        weights_pixel = weights_pixel * pts_pixel_mask
+        weights_pixel = weights_pixel / (
+                torch.sum(weights_pixel.float(), dim=1, keepdim=True) + 1e-8)  # [n_pts, N_views]
+        final_pixel_color = torch.sum(pts_pixel_color * weights_pixel[:, :, None], dim=1,
+                                      keepdim=False)  # [N_pts, 3]
+
+        final_pixel_mask = torch.sum(pts_pixel_mask.float(), dim=1, keepdim=True) > 0  # [N_pts, 1]
+
+        final_patch_color, final_patch_mask = None, None
+        # pts_patch_color  [N_pts, N_views, Npx, 3]; pts_patch_mask  [N_pts, N_views, Npx]
+        if pts_patch_color is not None:
+            N_pts, N_views, Npx, _ = pts_patch_color.shape
+            patch_mask = torch.sum(pts_patch_mask, dim=-1, keepdim=False) > Npx - 1  # [N_pts, N_views]
+
+            weights_patch = self.softmax(x_extracted)  # [N_pts, N_views]
+            weights_patch = weights_patch * patch_mask
+            weights_patch = weights_patch / (
+                    torch.sum(weights_patch.float(), dim=1, keepdim=True) + 1e-8)  # [n_pts, N_views]
+
+            final_patch_color = torch.sum(pts_patch_color * weights_patch[:, :, None, None], dim=1,
+                                          keepdim=False)  # [N_pts, Npx, 3]
+            final_patch_mask = torch.sum(patch_mask, dim=1, keepdim=True) > 0  # [N_pts, 1]  at least one image sees
+
+        return final_pixel_color, final_pixel_mask, final_patch_color, final_patch_mask

SparseNeuS_demo_v1/models/trainer_generic.py

@@ -0,0 +1,1207 @@
+"""
+decouple the trainer with the renderer
+"""
+import os
+import cv2 as cv
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+import trimesh
+from icecream import ic
+
+from utils.misc_utils import visualize_depth_numpy
+
+from loss.depth_metric import compute_depth_errors
+
+from loss.depth_loss import DepthLoss, DepthSmoothLoss
+
+from models.sparse_neus_renderer import SparseNeuSRenderer
+
+class GenericTrainer(nn.Module):
+    def __init__(self,
+                 rendering_network_outside,
+                 pyramid_feature_network_lod0,
+                 pyramid_feature_network_lod1,
+                 sdf_network_lod0,
+                 sdf_network_lod1,
+                 variance_network_lod0,
+                 variance_network_lod1,
+                 rendering_network_lod0,
+                 rendering_network_lod1,
+                 n_samples_lod0,
+                 n_importance_lod0,
+                 n_samples_lod1,
+                 n_importance_lod1,
+                 n_outside,
+                 perturb,
+                 alpha_type='div',
+                 conf=None,
+                 timestamp="",
+                 mode='train',
+                 base_exp_dir=None,
+                 ):
+        super(GenericTrainer, self).__init__()
+
+        self.conf = conf
+        self.timestamp = timestamp
+
+        
+        self.base_exp_dir = base_exp_dir 
+        
+
+        self.anneal_start = self.conf.get_float('train.anneal_start', default=0.0)
+        self.anneal_end = self.conf.get_float('train.anneal_end', default=0.0)
+        self.anneal_start_lod1 = self.conf.get_float('train.anneal_start_lod1', default=0.0)
+        self.anneal_end_lod1 = self.conf.get_float('train.anneal_end_lod1', default=0.0)
+
+        # network setups
+        self.rendering_network_outside = rendering_network_outside
+        self.pyramid_feature_network_geometry_lod0 = pyramid_feature_network_lod0  # 2D pyramid feature network for geometry
+        self.pyramid_feature_network_geometry_lod1 = pyramid_feature_network_lod1  # use differnet networks for the two lods
+
+        # when num_lods==2, may consume too much memeory
+        self.sdf_network_lod0 = sdf_network_lod0
+        self.sdf_network_lod1 = sdf_network_lod1
+
+        # - warpped by ModuleList to support DataParallel
+        self.variance_network_lod0 = variance_network_lod0
+        self.variance_network_lod1 = variance_network_lod1
+
+        self.rendering_network_lod0 = rendering_network_lod0
+        self.rendering_network_lod1 = rendering_network_lod1
+
+        self.n_samples_lod0 = n_samples_lod0
+        self.n_importance_lod0 = n_importance_lod0
+        self.n_samples_lod1 = n_samples_lod1
+        self.n_importance_lod1 = n_importance_lod1
+        self.n_outside = n_outside
+        self.num_lods = conf.get_int('model.num_lods')  # the number of octree lods
+        self.perturb = perturb
+        self.alpha_type = alpha_type
+
+        # - the two renderers
+        self.sdf_renderer_lod0 = SparseNeuSRenderer(
+            self.rendering_network_outside,
+            self.sdf_network_lod0,
+            self.variance_network_lod0,
+            self.rendering_network_lod0,
+            self.n_samples_lod0,
+            self.n_importance_lod0,
+            self.n_outside,
+            self.perturb,
+            alpha_type='div',
+            conf=self.conf)
+
+        self.sdf_renderer_lod1 = SparseNeuSRenderer(
+            self.rendering_network_outside,
+            self.sdf_network_lod1,
+            self.variance_network_lod1,
+            self.rendering_network_lod1,
+            self.n_samples_lod1,
+            self.n_importance_lod1,
+            self.n_outside,
+            self.perturb,
+            alpha_type='div',
+            conf=self.conf)
+
+        self.if_fix_lod0_networks = self.conf.get_bool('train.if_fix_lod0_networks')
+
+        # sdf network weights
+        self.sdf_igr_weight = self.conf.get_float('train.sdf_igr_weight')
+        self.sdf_sparse_weight = self.conf.get_float('train.sdf_sparse_weight', default=0)
+        self.sdf_decay_param = self.conf.get_float('train.sdf_decay_param', default=100)
+        self.fg_bg_weight = self.conf.get_float('train.fg_bg_weight', default=0.00)
+        self.bg_ratio = self.conf.get_float('train.bg_ratio', default=0.0)
+
+        self.depth_criterion = DepthLoss()
+
+        # - DataParallel mode, cannot modify attributes in forward()
+        # self.iter_step = 0
+        self.val_mesh_freq = self.conf.get_int('train.val_mesh_freq')
+
+        # - True for finetuning; False for general training
+        self.if_fitted_rendering = self.conf.get_bool('train.if_fitted_rendering', default=False)
+
+        self.prune_depth_filter = self.conf.get_bool('model.prune_depth_filter', default=False)
+
+    def get_trainable_params(self):
+        # set trainable params
+
+        self.params_to_train = []
+
+        if not self.if_fix_lod0_networks:
+            #  load pretrained featurenet
+            self.params_to_train += list(self.pyramid_feature_network_geometry_lod0.parameters())
+            self.params_to_train += list(self.sdf_network_lod0.parameters())
+            self.params_to_train += list(self.variance_network_lod0.parameters())
+
+            if self.rendering_network_lod0 is not None:
+                self.params_to_train += list(self.rendering_network_lod0.parameters())
+
+        if self.sdf_network_lod1 is not None:
+            #  load pretrained featurenet
+            self.params_to_train += list(self.pyramid_feature_network_geometry_lod1.parameters())
+
+            self.params_to_train += list(self.sdf_network_lod1.parameters())
+            self.params_to_train += list(self.variance_network_lod1.parameters())
+            if self.rendering_network_lod1 is not None:
+                self.params_to_train += list(self.rendering_network_lod1.parameters())
+
+        return self.params_to_train
+
+    def train_step(self, sample,
+                   perturb_overwrite=-1,
+                   background_rgb=None,
+                   alpha_inter_ratio_lod0=0.0,
+                   alpha_inter_ratio_lod1=0.0,
+                   iter_step=0,
+                   ):
+        # * only support batch_size==1
+        # ! attention: the list of string cannot be splited in DataParallel
+        batch_idx = sample['batch_idx'][0]
+        meta = sample['meta'][batch_idx]  # the scan lighting ref_view info
+
+        sizeW = sample['img_wh'][0][0]
+        sizeH = sample['img_wh'][0][1]
+        partial_vol_origin = sample['partial_vol_origin']  # [B, 3]
+        near, far = sample['near_fars'][0, 0, :1], sample['near_fars'][0, 0, 1:]
+
+        # the full-size ray variables
+        sample_rays = sample['rays']
+        rays_o = sample_rays['rays_o'][0]
+        rays_d = sample_rays['rays_v'][0]
+
+        imgs = sample['images'][0]
+        intrinsics = sample['intrinsics'][0]
+        intrinsics_l_4x = intrinsics.clone()
+        intrinsics_l_4x[:, :2] *= 0.25
+        w2cs = sample['w2cs'][0]
+        c2ws = sample['c2ws'][0]
+        proj_matrices = sample['affine_mats']
+        scale_mat = sample['scale_mat']
+        trans_mat = sample['trans_mat']
+
+        # ***********************     Lod==0     ***********************
+        if not self.if_fix_lod0_networks:
+            geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs)
+
+            conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
+                feature_maps=geometry_feature_maps[None, 1:, :, :, :],
+                partial_vol_origin=partial_vol_origin,
+                proj_mats=proj_matrices[:,1:],
+                # proj_mats=proj_matrices,
+                sizeH=sizeH,
+                sizeW=sizeW,
+                lod=0,
+            )
+
+        else:
+            with torch.no_grad():
+                geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
+                # geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
+                conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
+                    feature_maps=geometry_feature_maps[None, 1:, :, :, :],
+                    partial_vol_origin=partial_vol_origin,
+                    proj_mats=proj_matrices[:,1:],
+                    # proj_mats=proj_matrices,
+                    sizeH=sizeH,
+                    sizeW=sizeW,
+                    lod=0,
+                )
+        # print("Checker2:, construct cost volume")
+        con_volume_lod0 = conditional_features_lod0['dense_volume_scale0']
+
+        con_valid_mask_volume_lod0 = conditional_features_lod0['valid_mask_volume_scale0']
+        coords_lod0 = conditional_features_lod0['coords_scale0']  # [1,3,wX,wY,wZ]
+
+        # * extract depth maps for all the images
+        depth_maps_lod0, depth_masks_lod0 = None, None
+        if self.num_lods > 1:
+            sdf_volume_lod0 = self.sdf_network_lod0.get_sdf_volume(
+                con_volume_lod0, con_valid_mask_volume_lod0,
+                coords_lod0, partial_vol_origin)  # [1, 1, dX, dY, dZ]
+
+        if self.prune_depth_filter:
+            depth_maps_lod0_l4x, depth_masks_lod0_l4x = self.sdf_renderer_lod0.extract_depth_maps(
+                self.sdf_network_lod0, sdf_volume_lod0, intrinsics_l_4x, c2ws,
+                sizeH // 4, sizeW // 4, near * 1.5, far)
+            depth_maps_lod0 = F.interpolate(depth_maps_lod0_l4x, size=(sizeH, sizeW), mode='bilinear',
+                                            align_corners=True)
+            depth_masks_lod0 = F.interpolate(depth_masks_lod0_l4x.float(), size=(sizeH, sizeW), mode='nearest')
+
+        # *************** losses
+        loss_lod0, losses_lod0, depth_statis_lod0 = None, None, None
+
+        if not self.if_fix_lod0_networks:
+
+            render_out = self.sdf_renderer_lod0.render(
+                rays_o, rays_d, near, far,
+                self.sdf_network_lod0,
+                self.rendering_network_lod0,
+                background_rgb=background_rgb,
+                alpha_inter_ratio=alpha_inter_ratio_lod0,
+                # * related to conditional feature
+                lod=0,
+                conditional_volume=con_volume_lod0,
+                conditional_valid_mask_volume=con_valid_mask_volume_lod0,
+                # * 2d feature maps
+                feature_maps=geometry_feature_maps,
+                color_maps=imgs,
+                w2cs=w2cs,
+                intrinsics=intrinsics,
+                img_wh=[sizeW, sizeH],
+                if_general_rendering=True,
+                if_render_with_grad=True,
+            )
+
+            loss_lod0, losses_lod0, depth_statis_lod0 = self.cal_losses_sdf(render_out, sample_rays,
+                                                                                     iter_step, lod=0)
+
+        # ***********************     Lod==1     ***********************
+
+        loss_lod1, losses_lod1, depth_statis_lod1 = None, None, None
+
+        if self.num_lods > 1:
+            geometry_feature_maps_lod1 = self.obtain_pyramid_feature_maps(imgs, lod=1)
+            # geometry_feature_maps_lod1 = self.obtain_pyramid_feature_maps(imgs, lod=1)
+            if self.prune_depth_filter:
+                pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf_depthfilter(
+                    sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0],
+                    depth_maps_lod0, proj_matrices[0],
+                    partial_vol_origin, self.sdf_network_lod0.voxel_size,
+                    near, far, self.sdf_network_lod0.voxel_size, 12)
+            else:
+                pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf(
+                    sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0])
+
+            pre_coords[:, 1:] = pre_coords[:, 1:] * 2
+
+            # ? It seems that training gru_fusion, this part should be trainable too
+            conditional_features_lod1 = self.sdf_network_lod1.get_conditional_volume(
+                feature_maps=geometry_feature_maps_lod1[None, 1:, :, :, :],
+                partial_vol_origin=partial_vol_origin,
+                proj_mats=proj_matrices[:,1:],
+                # proj_mats=proj_matrices,
+                sizeH=sizeH,
+                sizeW=sizeW,
+                pre_coords=pre_coords,
+                pre_feats=pre_feats,
+            )
+
+            con_volume_lod1 = conditional_features_lod1['dense_volume_scale1']
+            con_valid_mask_volume_lod1 = conditional_features_lod1['valid_mask_volume_scale1']
+
+            # if not self.if_gru_fusion_lod1:
+            render_out_lod1 = self.sdf_renderer_lod1.render(
+                rays_o, rays_d, near, far,
+                self.sdf_network_lod1,
+                self.rendering_network_lod1,
+                background_rgb=background_rgb,
+                alpha_inter_ratio=alpha_inter_ratio_lod1,
+                # * related to conditional feature
+                lod=1,
+                conditional_volume=con_volume_lod1,
+                conditional_valid_mask_volume=con_valid_mask_volume_lod1,
+                # * 2d feature maps
+                feature_maps=geometry_feature_maps_lod1,
+                color_maps=imgs,
+                w2cs=w2cs,
+                intrinsics=intrinsics,
+                img_wh=[sizeW, sizeH],
+                bg_ratio=self.bg_ratio,
+            )
+            loss_lod1, losses_lod1, depth_statis_lod1 = self.cal_losses_sdf(render_out_lod1, sample_rays,
+                                                                                     iter_step, lod=1)
+
+        # print("Checker3:, compute losses")
+        # # - extract mesh
+        if iter_step % self.val_mesh_freq == 0:
+            torch.cuda.empty_cache()
+            self.validate_mesh(self.sdf_network_lod0,
+                               self.sdf_renderer_lod0.extract_geometry,
+                               conditional_volume=con_volume_lod0, lod=0,
+                               threshold=0,
+                               # occupancy_mask=con_valid_mask_volume_lod0[0, 0],
+                               mode='train_bg', meta=meta,
+                               iter_step=iter_step, scale_mat=scale_mat,
+                               trans_mat=trans_mat)
+            torch.cuda.empty_cache()
+
+            if self.num_lods > 1:
+                self.validate_mesh(self.sdf_network_lod1,
+                                   self.sdf_renderer_lod1.extract_geometry,
+                                   conditional_volume=con_volume_lod1, lod=1,
+                                   # occupancy_mask=con_valid_mask_volume_lod1[0, 0].detach(),
+                                   mode='train_bg', meta=meta,
+                                   iter_step=iter_step, scale_mat=scale_mat,
+                                   trans_mat=trans_mat)
+        losses = {
+            # - lod 0
+            'loss_lod0': loss_lod0,
+            'losses_lod0': losses_lod0,
+            'depth_statis_lod0': depth_statis_lod0,
+
+            # - lod 1
+            'loss_lod1': loss_lod1,
+            'losses_lod1': losses_lod1,
+            'depth_statis_lod1': depth_statis_lod1,
+
+        }
+
+        return losses
+
+    def val_step(self, sample,
+                 perturb_overwrite=-1,
+                 background_rgb=None,
+                 alpha_inter_ratio_lod0=0.0,
+                 alpha_inter_ratio_lod1=0.0,
+                 iter_step=0,
+                 chunk_size=512,
+                 save_vis=False,
+                 ):
+        # * only support batch_size==1
+        # ! attention: the list of string cannot be splited in DataParallel
+        batch_idx = sample['batch_idx'][0]
+        meta = sample['meta'][batch_idx]  # the scan lighting ref_view info
+
+        sizeW = sample['img_wh'][0][0]
+        sizeH = sample['img_wh'][0][1]
+        H, W = sizeH, sizeW
+
+        partial_vol_origin = sample['partial_vol_origin']  # [B, 3]
+        near, far = sample['query_near_far'][0, :1], sample['query_near_far'][0, 1:]
+
+        # the ray variables
+        sample_rays = sample['rays']
+        rays_o = sample_rays['rays_o'][0]
+        rays_d = sample_rays['rays_v'][0]
+        rays_ndc_uv = sample_rays['rays_ndc_uv'][0]
+
+        imgs = sample['images'][0]
+        intrinsics = sample['intrinsics'][0]
+        intrinsics_l_4x = intrinsics.clone()
+        intrinsics_l_4x[:, :2] *= 0.25
+        w2cs = sample['w2cs'][0]
+        c2ws = sample['c2ws'][0]
+        proj_matrices = sample['affine_mats']
+
+        # render_img_idx = sample['render_img_idx'][0]
+        # true_img = sample['images'][0][render_img_idx]
+
+        # - the image to render
+        scale_mat = sample['scale_mat']  # [1,4,4]  used to convert mesh into true scale
+        trans_mat = sample['trans_mat']
+        query_c2w = sample['query_c2w']  # [1,4,4]
+        query_w2c = sample['query_w2c']  # [1,4,4]
+        true_img = sample['query_image'][0]
+        true_img = np.uint8(true_img.permute(1, 2, 0).cpu().numpy() * 255)
+
+        depth_min, depth_max = near.cpu().numpy(), far.cpu().numpy()
+
+        scale_factor = sample['scale_factor'][0].cpu().numpy()
+        true_depth = sample['query_depth'] if 'query_depth' in sample.keys() else None
+        if true_depth is not None:
+            true_depth = true_depth[0].cpu().numpy()
+            true_depth_colored = visualize_depth_numpy(true_depth, [depth_min, depth_max])[0]
+        else:
+            true_depth_colored = None
+
+        rays_o = rays_o.reshape(-1, 3).split(chunk_size)
+        rays_d = rays_d.reshape(-1, 3).split(chunk_size)
+
+        # - obtain conditional features
+        with torch.no_grad():
+            # - obtain conditional features
+            geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
+            # - lod 0
+            conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
+                feature_maps=geometry_feature_maps[None, :, :, :, :],
+                partial_vol_origin=partial_vol_origin,
+                proj_mats=proj_matrices,
+                sizeH=sizeH,
+                sizeW=sizeW,
+                lod=0,
+            )
+
+        con_volume_lod0 = conditional_features_lod0['dense_volume_scale0']
+        con_valid_mask_volume_lod0 = conditional_features_lod0['valid_mask_volume_scale0']
+        coords_lod0 = conditional_features_lod0['coords_scale0']  # [1,3,wX,wY,wZ]
+
+        if self.num_lods > 1:
+            sdf_volume_lod0 = self.sdf_network_lod0.get_sdf_volume(
+                con_volume_lod0, con_valid_mask_volume_lod0,
+                coords_lod0, partial_vol_origin)  # [1, 1, dX, dY, dZ]
+
+        depth_maps_lod0, depth_masks_lod0 = None, None
+        if self.prune_depth_filter:
+            depth_maps_lod0_l4x, depth_masks_lod0_l4x = self.sdf_renderer_lod0.extract_depth_maps(
+                self.sdf_network_lod0, sdf_volume_lod0,
+                intrinsics_l_4x, c2ws,
+                sizeH // 4, sizeW // 4, near * 1.5, far)  # - near*1.5 is a experienced number
+            depth_maps_lod0 = F.interpolate(depth_maps_lod0_l4x, size=(sizeH, sizeW), mode='bilinear',
+                                            align_corners=True)
+            depth_masks_lod0 = F.interpolate(depth_masks_lod0_l4x.float(), size=(sizeH, sizeW), mode='nearest')
+
+            #### visualize the depth_maps_lod0 for checking
+            colored_depth_maps_lod0 = []
+            for i in range(depth_maps_lod0.shape[0]):
+                colored_depth_maps_lod0.append(
+                    visualize_depth_numpy(depth_maps_lod0[i, 0].cpu().numpy(), [depth_min, depth_max])[0])
+
+            colored_depth_maps_lod0 = np.concatenate(colored_depth_maps_lod0, axis=0).astype(np.uint8)
+            os.makedirs(os.path.join(self.base_exp_dir, 'depth_maps_lod0'), exist_ok=True)
+            cv.imwrite(os.path.join(self.base_exp_dir, 'depth_maps_lod0',
+                                    '{:0>8d}_{}.png'.format(iter_step, meta)),
+                       colored_depth_maps_lod0[:, :, ::-1])
+
+        if self.num_lods > 1:
+            geometry_feature_maps_lod1 = self.obtain_pyramid_feature_maps(imgs, lod=1)
+
+            if self.prune_depth_filter:
+                pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf_depthfilter(
+                    sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0],
+                    depth_maps_lod0, proj_matrices[0],
+                    partial_vol_origin, self.sdf_network_lod0.voxel_size,
+                    near, far, self.sdf_network_lod0.voxel_size, 12)
+            else:
+                pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf(
+                    sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0])
+
+            pre_coords[:, 1:] = pre_coords[:, 1:] * 2
+
+            with torch.no_grad():
+                conditional_features_lod1 = self.sdf_network_lod1.get_conditional_volume(
+                    feature_maps=geometry_feature_maps_lod1[None, :, :, :, :],
+                    partial_vol_origin=partial_vol_origin,
+                    proj_mats=proj_matrices,
+                    sizeH=sizeH,
+                    sizeW=sizeW,
+                    pre_coords=pre_coords,
+                    pre_feats=pre_feats,
+                )
+
+            con_volume_lod1 = conditional_features_lod1['dense_volume_scale1']
+            con_valid_mask_volume_lod1 = conditional_features_lod1['valid_mask_volume_scale1']
+
+        out_rgb_fine = []
+        out_normal_fine = []
+        out_depth_fine = []
+
+        out_rgb_fine_lod1 = []
+        out_normal_fine_lod1 = []
+        out_depth_fine_lod1 = []
+
+        # out_depth_fine_explicit = []
+        if save_vis:
+            for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
+
+                # ****** lod 0 ****
+                render_out = self.sdf_renderer_lod0.render(
+                    rays_o_batch, rays_d_batch, near, far,
+                    self.sdf_network_lod0,
+                    self.rendering_network_lod0,
+                    background_rgb=background_rgb,
+                    alpha_inter_ratio=alpha_inter_ratio_lod0,
+                    # * related to conditional feature
+                    lod=0,
+                    conditional_volume=con_volume_lod0,
+                    conditional_valid_mask_volume=con_valid_mask_volume_lod0,
+                    # * 2d feature maps
+                    feature_maps=geometry_feature_maps,
+                    color_maps=imgs,
+                    w2cs=w2cs,
+                    intrinsics=intrinsics,
+                    img_wh=[sizeW, sizeH],
+                    query_c2w=query_c2w,
+                    if_render_with_grad=False,
+                )
+
+                feasible = lambda key: ((key in render_out) and (render_out[key] is not None))
+
+                if feasible('depth'):
+                    out_depth_fine.append(render_out['depth'].detach().cpu().numpy())
+
+                # if render_out['color_coarse'] is not None:
+                if feasible('color_fine'):
+                    out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())
+                if feasible('gradients') and feasible('weights'):
+                    if render_out['inside_sphere'] is not None:
+                        out_normal_fine.append((render_out['gradients'] * render_out['weights'][:,
+                                                                          :self.n_samples_lod0 + self.n_importance_lod0,
+                                                                          None] * render_out['inside_sphere'][
+                                                    ..., None]).sum(dim=1).detach().cpu().numpy())
+                    else:
+                        out_normal_fine.append((render_out['gradients'] * render_out['weights'][:,
+                                                                          :self.n_samples_lod0 + self.n_importance_lod0,
+                                                                          None]).sum(dim=1).detach().cpu().numpy())
+                del render_out
+
+            # ****************** lod 1 **************************
+            if self.num_lods > 1:
+                for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
+                    render_out_lod1 = self.sdf_renderer_lod1.render(
+                        rays_o_batch, rays_d_batch, near, far,
+                        self.sdf_network_lod1,
+                        self.rendering_network_lod1,
+                        background_rgb=background_rgb,
+                        alpha_inter_ratio=alpha_inter_ratio_lod1,
+                        # * related to conditional feature
+                        lod=1,
+                        conditional_volume=con_volume_lod1,
+                        conditional_valid_mask_volume=con_valid_mask_volume_lod1,
+                        # * 2d feature maps
+                        feature_maps=geometry_feature_maps_lod1,
+                        color_maps=imgs,
+                        w2cs=w2cs,
+                        intrinsics=intrinsics,
+                        img_wh=[sizeW, sizeH],
+                        query_c2w=query_c2w,
+                        if_render_with_grad=False,
+                    )
+
+                    feasible = lambda key: ((key in render_out_lod1) and (render_out_lod1[key] is not None))
+
+                    if feasible('depth'):
+                        out_depth_fine_lod1.append(render_out_lod1['depth'].detach().cpu().numpy())
+
+                    # if render_out['color_coarse'] is not None:
+                    if feasible('color_fine'):
+                        out_rgb_fine_lod1.append(render_out_lod1['color_fine'].detach().cpu().numpy())
+                    if feasible('gradients') and feasible('weights'):
+                        if render_out_lod1['inside_sphere'] is not None:
+                            out_normal_fine_lod1.append((render_out_lod1['gradients'] * render_out_lod1['weights'][:,
+                                                                                        :self.n_samples_lod1 + self.n_importance_lod1,
+                                                                                        None] *
+                                                         render_out_lod1['inside_sphere'][
+                                                             ..., None]).sum(dim=1).detach().cpu().numpy())
+                        else:
+                            out_normal_fine_lod1.append((render_out_lod1['gradients'] * render_out_lod1['weights'][:,
+                                                                                        :self.n_samples_lod1 + self.n_importance_lod1,
+                                                                                        None]).sum(
+                                dim=1).detach().cpu().numpy())
+                    del render_out_lod1
+
+            # - save visualization of lod 0
+
+            self.save_visualization(true_img, true_depth_colored, out_depth_fine, out_normal_fine,
+                                    query_w2c[0], out_rgb_fine, H, W,
+                                    depth_min, depth_max, iter_step, meta, "val_lod0", true_depth=true_depth, scale_factor=scale_factor)
+
+            if self.num_lods > 1:
+                self.save_visualization(true_img, true_depth_colored, out_depth_fine_lod1, out_normal_fine_lod1,
+                                        query_w2c[0], out_rgb_fine_lod1, H, W,
+                                        depth_min, depth_max, iter_step, meta, "val_lod1", true_depth=true_depth, scale_factor=scale_factor)
+
+        # - extract mesh
+        if (iter_step % self.val_mesh_freq == 0):
+            torch.cuda.empty_cache()
+            self.validate_mesh(self.sdf_network_lod0,
+                               self.sdf_renderer_lod0.extract_geometry,
+                               conditional_volume=con_volume_lod0, lod=0,
+                               threshold=0,
+                               # occupancy_mask=con_valid_mask_volume_lod0[0, 0],
+                               mode='val_bg', meta=meta,
+                               iter_step=iter_step, scale_mat=scale_mat, trans_mat=trans_mat)
+            torch.cuda.empty_cache()
+
+            if self.num_lods > 1:
+                self.validate_mesh(self.sdf_network_lod1,
+                                   self.sdf_renderer_lod1.extract_geometry,
+                                   conditional_volume=con_volume_lod1, lod=1,
+                                   # occupancy_mask=con_valid_mask_volume_lod1[0, 0].detach(),
+                                   mode='val_bg', meta=meta,
+                                   iter_step=iter_step, scale_mat=scale_mat, trans_mat=trans_mat)
+
+            torch.cuda.empty_cache()
+
+
+
+    def export_mesh_step(self, sample,
+                        perturb_overwrite=-1,
+                        background_rgb=None,
+                        alpha_inter_ratio_lod0=0.0,
+                        alpha_inter_ratio_lod1=0.0,
+                        iter_step=0,
+                        chunk_size=512,
+                        save_vis=False,
+                        ):
+        # * only support batch_size==1
+        # ! attention: the list of string cannot be splited in DataParallel
+        batch_idx = sample['batch_idx'][0]
+        meta = sample['meta'][batch_idx]  # the scan lighting ref_view info
+
+        sizeW = sample['img_wh'][0][0]
+        sizeH = sample['img_wh'][0][1]
+        H, W = sizeH, sizeW
+
+        partial_vol_origin = sample['partial_vol_origin']  # [B, 3]
+        near, far = sample['query_near_far'][0, :1], sample['query_near_far'][0, 1:]
+
+        # the ray variables
+        sample_rays = sample['rays']
+        rays_o = sample_rays['rays_o'][0]
+        rays_d = sample_rays['rays_v'][0]
+        rays_ndc_uv = sample_rays['rays_ndc_uv'][0]
+
+        imgs = sample['images'][0]
+        intrinsics = sample['intrinsics'][0]
+        intrinsics_l_4x = intrinsics.clone()
+        intrinsics_l_4x[:, :2] *= 0.25
+        w2cs = sample['w2cs'][0]
+        c2ws = sample['c2ws'][0]
+        # target_candidate_w2cs = sample['target_candidate_w2cs'][0]
+        proj_matrices = sample['affine_mats']
+
+
+        # - the image to render
+        scale_mat = sample['scale_mat']  # [1,4,4]  used to convert mesh into true scale
+        trans_mat = sample['trans_mat']
+        query_c2w = sample['query_c2w']  # [1,4,4]
+        query_w2c = sample['query_w2c']  # [1,4,4]
+        true_img = sample['query_image'][0]
+        true_img = np.uint8(true_img.permute(1, 2, 0).cpu().numpy() * 255)
+
+        depth_min, depth_max = near.cpu().numpy(), far.cpu().numpy()
+
+        scale_factor = sample['scale_factor'][0].cpu().numpy()
+        true_depth = sample['query_depth'] if 'query_depth' in sample.keys() else None
+        if true_depth is not None:
+            true_depth = true_depth[0].cpu().numpy()
+            true_depth_colored = visualize_depth_numpy(true_depth, [depth_min, depth_max])[0]
+        else:
+            true_depth_colored = None
+
+        rays_o = rays_o.reshape(-1, 3).split(chunk_size)
+        rays_d = rays_d.reshape(-1, 3).split(chunk_size)
+        # import time
+        # jha_begin1 = time.time()
+        # - obtain conditional features
+        with torch.no_grad():
+            # - obtain conditional features
+            geometry_feature_maps = self.obtain_pyramid_feature_maps(imgs, lod=0)
+            # - lod 0
+            conditional_features_lod0 = self.sdf_network_lod0.get_conditional_volume(
+                feature_maps=geometry_feature_maps[None, :, :, :, :],
+                partial_vol_origin=partial_vol_origin,
+                proj_mats=proj_matrices,
+                sizeH=sizeH,
+                sizeW=sizeW,
+                lod=0,
+            )
+        # jha_end1 = time.time()
+        # print("get_conditional_volume: ", jha_end1 - jha_begin1)
+        # jha_begin2 = time.time()
+        con_volume_lod0 = conditional_features_lod0['dense_volume_scale0']
+        con_valid_mask_volume_lod0 = conditional_features_lod0['valid_mask_volume_scale0']
+        coords_lod0 = conditional_features_lod0['coords_scale0']  # [1,3,wX,wY,wZ]
+
+        if self.num_lods > 1:
+            sdf_volume_lod0 = self.sdf_network_lod0.get_sdf_volume(
+                con_volume_lod0, con_valid_mask_volume_lod0,
+                coords_lod0, partial_vol_origin)  # [1, 1, dX, dY, dZ]
+
+        depth_maps_lod0, depth_masks_lod0 = None, None
+
+
+        if self.num_lods > 1:
+            geometry_feature_maps_lod1 = self.obtain_pyramid_feature_maps(imgs, lod=1)
+
+            if self.prune_depth_filter:
+                pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf_depthfilter(
+                    sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0],
+                    depth_maps_lod0, proj_matrices[0],
+                    partial_vol_origin, self.sdf_network_lod0.voxel_size,
+                    near, far, self.sdf_network_lod0.voxel_size, 12)
+            else:
+                pre_coords, pre_feats = self.sdf_renderer_lod0.get_valid_sparse_coords_by_sdf(
+                    sdf_volume_lod0[0], coords_lod0[0], con_valid_mask_volume_lod0[0], con_volume_lod0[0])
+
+            pre_coords[:, 1:] = pre_coords[:, 1:] * 2
+
+            with torch.no_grad():
+                conditional_features_lod1 = self.sdf_network_lod1.get_conditional_volume(
+                    feature_maps=geometry_feature_maps_lod1[None, :, :, :, :],
+                    partial_vol_origin=partial_vol_origin,
+                    proj_mats=proj_matrices,
+                    sizeH=sizeH,
+                    sizeW=sizeW,
+                    pre_coords=pre_coords,
+                    pre_feats=pre_feats,
+                )
+
+            con_volume_lod1 = conditional_features_lod1['dense_volume_scale1']
+            con_valid_mask_volume_lod1 = conditional_features_lod1['valid_mask_volume_scale1']
+
+        out_rgb_fine = []
+        out_normal_fine = []
+        out_depth_fine = []
+
+        out_rgb_fine_lod1 = []
+        out_normal_fine_lod1 = []
+        out_depth_fine_lod1 = []
+
+        # jha_end2 = time.time()
+        # print("interval before starting mesh export: ", jha_end2 - jha_begin2)
+
+        # - extract mesh
+        if (iter_step % self.val_mesh_freq == 0):
+            torch.cuda.empty_cache()
+            # jha_begin3 = time.time()
+            self.validate_colored_mesh(
+                                        density_or_sdf_network=self.sdf_network_lod0,
+                                        func_extract_geometry=self.sdf_renderer_lod0.extract_geometry,
+                                        conditional_volume=con_volume_lod0,
+                                        conditional_valid_mask_volume = con_valid_mask_volume_lod0,
+                                        feature_maps=geometry_feature_maps,
+                                        color_maps=imgs,
+                                        w2cs=w2cs,
+                                        target_candidate_w2cs=None,
+                                        intrinsics=intrinsics,
+                                        rendering_network=self.rendering_network_lod0,
+                                        rendering_projector=self.sdf_renderer_lod0.rendering_projector,
+                                        lod=0,
+                                        threshold=0,
+                                        query_c2w=query_c2w,
+                                        mode='val_bg', meta=meta,
+                                        iter_step=iter_step, scale_mat=scale_mat, trans_mat=trans_mat
+                                    )
+            torch.cuda.empty_cache()
+            # jha_end3 = time.time()
+            # print("validate_colored_mesh_test_time: ", jha_end3 - jha_begin3)
+
+            if self.num_lods > 1:
+                self.validate_colored_mesh(
+                            density_or_sdf_network=self.sdf_network_lod1,
+                            func_extract_geometry=self.sdf_renderer_lod1.extract_geometry,
+                            conditional_volume=con_volume_lod1,
+                            conditional_valid_mask_volume = con_valid_mask_volume_lod1,
+                            feature_maps=geometry_feature_maps,
+                            color_maps=imgs,
+                            w2cs=w2cs,
+                            target_candidate_w2cs=None,
+                            intrinsics=intrinsics,
+                            rendering_network=self.rendering_network_lod1,
+                            rendering_projector=self.sdf_renderer_lod1.rendering_projector,
+                            lod=1,
+                            threshold=0,
+                            query_c2w=query_c2w,
+                            mode='val_bg', meta=meta,
+                            iter_step=iter_step, scale_mat=scale_mat, trans_mat=trans_mat
+                        )
+            torch.cuda.empty_cache()
+
+
+
+    def save_visualization(self, true_img, true_colored_depth, out_depth, out_normal, w2cs, out_color, H, W,
+                           depth_min, depth_max, iter_step, meta, comment, out_color_mlp=[], true_depth=None, scale_factor=1.0):
+        if len(out_color) > 0:
+            img_fine = (np.concatenate(out_color, axis=0).reshape([H, W, 3]) * 256).clip(0, 255)
+
+        if len(out_color_mlp) > 0:
+            img_mlp = (np.concatenate(out_color_mlp, axis=0).reshape([H, W, 3]) * 256).clip(0, 255)
+
+        if len(out_normal) > 0:
+            normal_img = np.concatenate(out_normal, axis=0)
+            rot = w2cs[:3, :3].detach().cpu().numpy()
+            # - convert normal from world space to camera space
+            normal_img = (np.matmul(rot[None, :, :],
+                                    normal_img[:, :, None]).reshape([H, W, 3]) * 128 + 128).clip(0, 255)
+        if len(out_depth) > 0:
+            pred_depth = np.concatenate(out_depth, axis=0).reshape([H, W])
+            pred_depth_colored = visualize_depth_numpy(pred_depth, [depth_min, depth_max])[0]
+
+        if len(out_depth) > 0:
+            os.makedirs(os.path.join(self.base_exp_dir, 'depths_' + comment), exist_ok=True)
+            if true_colored_depth is not None:
+
+                if true_depth is not None:
+                    depth_error_map = np.abs(true_depth - pred_depth) * 2.0 / scale_factor
+                    # [256, 256, 1] -> [256, 256, 3]
+                    depth_error_map = np.tile(depth_error_map[:, :, None], [1, 1, 3])
+                    print("meta: ", meta)
+                    print("scale_factor: ", scale_factor)
+                    print("depth_error_mean: ", depth_error_map.mean())
+                    depth_visualized = np.concatenate(
+                            [(depth_error_map * 255).astype(np.uint8), true_colored_depth, pred_depth_colored, true_img], axis=1)[:, :, ::-1]
+                    # print("depth_visualized.shape: ", depth_visualized.shape)
+                    # write  depth error result text on img, the input is a numpy array of [256, 1024, 3]
+                    # cv.putText(depth_visualized.copy(), "depth_error_mean: {:.4f}".format(depth_error_map.mean()), (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+                else:
+                    depth_visualized = np.concatenate(
+                            [true_colored_depth, pred_depth_colored, true_img])[:, :, ::-1]
+                cv.imwrite(
+                    os.path.join(self.base_exp_dir, 'depths_' + comment,
+                                 '{:0>8d}_{}.png'.format(iter_step, meta)), depth_visualized
+                    )
+            else:
+                cv.imwrite(
+                    os.path.join(self.base_exp_dir, 'depths_' + comment,
+                                 '{:0>8d}_{}.png'.format(iter_step, meta)),
+                    np.concatenate(
+                        [pred_depth_colored, true_img])[:, :, ::-1])
+        if len(out_color) > 0:
+            os.makedirs(os.path.join(self.base_exp_dir, 'synthesized_color_' + comment), exist_ok=True)
+            cv.imwrite(os.path.join(self.base_exp_dir, 'synthesized_color_' + comment,
+                                    '{:0>8d}_{}.png'.format(iter_step, meta)),
+                       np.concatenate(
+                           [img_fine, true_img])[:, :, ::-1])  # bgr2rgb
+            # compute psnr (image pixel lie in [0, 255])
+            mse_loss = np.mean((img_fine - true_img) ** 2)
+            psnr = 10 * np.log10(255 ** 2 / mse_loss)
+            
+            print("PSNR: ", psnr)
+
+        if len(out_color_mlp) > 0:
+            os.makedirs(os.path.join(self.base_exp_dir, 'synthesized_color_mlp_' + comment), exist_ok=True)
+            cv.imwrite(os.path.join(self.base_exp_dir, 'synthesized_color_mlp_' + comment,
+                                    '{:0>8d}_{}.png'.format(iter_step, meta)),
+                       np.concatenate(
+                           [img_mlp, true_img])[:, :, ::-1])  # bgr2rgb
+
+        if len(out_normal) > 0:
+            os.makedirs(os.path.join(self.base_exp_dir, 'normals_' + comment), exist_ok=True)
+            cv.imwrite(os.path.join(self.base_exp_dir, 'normals_' + comment,
+                                    '{:0>8d}_{}.png'.format(iter_step, meta)),
+                       normal_img[:, :, ::-1])
+
+    def forward(self, sample,
+                perturb_overwrite=-1,
+                background_rgb=None,
+                alpha_inter_ratio_lod0=0.0,
+                alpha_inter_ratio_lod1=0.0,
+                iter_step=0,
+                mode='train',
+                save_vis=False,
+                ):
+
+        if mode == 'train':
+            return self.train_step(sample,
+                                   perturb_overwrite=perturb_overwrite,
+                                   background_rgb=background_rgb,
+                                   alpha_inter_ratio_lod0=alpha_inter_ratio_lod0,
+                                   alpha_inter_ratio_lod1=alpha_inter_ratio_lod1,
+                                   iter_step=iter_step
+                                   )
+        elif mode == 'val':
+            import time
+            begin = time.time()
+            result =  self.val_step(sample,
+                                 perturb_overwrite=perturb_overwrite,
+                                 background_rgb=background_rgb,
+                                 alpha_inter_ratio_lod0=alpha_inter_ratio_lod0,
+                                 alpha_inter_ratio_lod1=alpha_inter_ratio_lod1,
+                                 iter_step=iter_step,
+                                 save_vis=save_vis,
+                                 )
+            end = time.time()
+            print("val_step time: ", end - begin)
+            return result
+        elif mode == 'export_mesh':
+            import time
+            begin = time.time()
+            result =  self.export_mesh_step(sample,
+                                 perturb_overwrite=perturb_overwrite,
+                                 background_rgb=background_rgb,
+                                 alpha_inter_ratio_lod0=alpha_inter_ratio_lod0,
+                                 alpha_inter_ratio_lod1=alpha_inter_ratio_lod1,
+                                 iter_step=iter_step,
+                                 save_vis=save_vis,
+                                 )
+            end = time.time()
+            print("export mesh time: ", end - begin)
+            return result
+    def obtain_pyramid_feature_maps(self, imgs, lod=0):
+        """
+        get feature maps of all conditional images
+        :param imgs:
+        :return:
+        """
+
+        if lod == 0:
+            extractor = self.pyramid_feature_network_geometry_lod0
+        elif lod >= 1:
+            extractor = self.pyramid_feature_network_geometry_lod1
+
+        pyramid_feature_maps = extractor(imgs)
+
+        # * the pyramid features are very important, if only use the coarst features, hard to optimize
+        fused_feature_maps = torch.cat([
+            F.interpolate(pyramid_feature_maps[0], scale_factor=4, mode='bilinear', align_corners=True),
+            F.interpolate(pyramid_feature_maps[1], scale_factor=2, mode='bilinear', align_corners=True),
+            pyramid_feature_maps[2]
+        ], dim=1)
+
+        return fused_feature_maps
+
+    def cal_losses_sdf(self, render_out, sample_rays, iter_step=-1, lod=0):
+
+        # loss weight schedule; the regularization terms should be added in later training stage
+        def get_weight(iter_step, weight):
+            if lod == 1:
+                anneal_start = self.anneal_end if lod == 0 else self.anneal_end_lod1
+                anneal_end = self.anneal_end if lod == 0 else self.anneal_end_lod1
+                anneal_end = anneal_end * 2
+            else:
+                anneal_start = self.anneal_start if lod == 0 else self.anneal_start_lod1
+                anneal_end = self.anneal_end if lod == 0 else self.anneal_end_lod1
+                anneal_end = anneal_end * 2
+
+            if iter_step < 0:
+                return weight
+
+            if anneal_end == 0.0:
+                return weight
+            elif iter_step < anneal_start:
+                return 0.0
+            else:
+                return np.min(
+                    [1.0,
+                     (iter_step - anneal_start) / (anneal_end - anneal_start)]) * weight
+
+        rays_o = sample_rays['rays_o'][0]
+        rays_d = sample_rays['rays_v'][0]
+        true_rgb = sample_rays['rays_color'][0]
+
+        if 'rays_depth' in sample_rays.keys():
+            true_depth = sample_rays['rays_depth'][0]
+        else:
+            true_depth = None
+        mask = sample_rays['rays_mask'][0]
+
+        color_fine = render_out['color_fine']
+        color_fine_mask = render_out['color_fine_mask']
+        depth_pred = render_out['depth']
+
+        variance = render_out['variance']
+        cdf_fine = render_out['cdf_fine']
+        weight_sum = render_out['weights_sum']
+
+        gradient_error_fine = render_out['gradient_error_fine']
+
+        sdf = render_out['sdf']
+
+        # * color generated by mlp
+        color_mlp = render_out['color_mlp']
+        color_mlp_mask = render_out['color_mlp_mask']
+
+        if color_fine is not None:
+            # Color loss
+            color_mask = color_fine_mask if color_fine_mask is not None else mask
+            color_mask = color_mask[..., 0]
+            color_error = (color_fine[color_mask] - true_rgb[color_mask])
+            # print("Nan number", torch.isnan(color_error).sum())
+            # print("Color error shape", color_error.shape)
+            color_fine_loss = F.l1_loss(color_error, torch.zeros_like(color_error).to(color_error.device),
+                                        reduction='mean')
+            # print(color_fine_loss)
+            psnr = 20.0 * torch.log10(
+                1.0 / (((color_fine[color_mask] - true_rgb[color_mask]) ** 2).mean() / (3.0)).sqrt())
+        else:
+            color_fine_loss = 0.
+            psnr = 0.
+
+        if color_mlp is not None:
+            # Color loss
+            color_mlp_mask = color_mlp_mask[..., 0]
+            color_error_mlp = (color_mlp[color_mlp_mask] - true_rgb[color_mlp_mask])
+            color_mlp_loss = F.l1_loss(color_error_mlp,
+                                       torch.zeros_like(color_error_mlp).to(color_error_mlp.device),
+                                       reduction='mean')
+
+            psnr_mlp = 20.0 * torch.log10(
+                1.0 / (((color_mlp[color_mlp_mask] - true_rgb[color_mlp_mask]) ** 2).mean() / (3.0)).sqrt())
+        else:
+            color_mlp_loss = 0.
+            psnr_mlp = 0.
+
+        # depth loss is only used for inference, not included in total loss
+        if true_depth is not None:
+            # depth_loss = self.depth_criterion(depth_pred, true_depth, mask)
+            depth_loss = self.depth_criterion(depth_pred, true_depth)
+
+            # # depth evaluation
+            # depth_statis = compute_depth_errors(depth_pred.detach().cpu().numpy(), true_depth.cpu().numpy())
+            # depth_statis = numpy2tensor(depth_statis, device=rays_o.device)
+            depth_statis = None
+        else:
+            depth_loss = 0.
+            depth_statis = None
+
+        sparse_loss_1 = torch.exp(
+            -1 * torch.abs(render_out['sdf_random']) * self.sdf_decay_param).mean()  # - should equal
+        sparse_loss_2 = torch.exp(-1 * torch.abs(sdf) * self.sdf_decay_param).mean()
+        sparse_loss = (sparse_loss_1 + sparse_loss_2) / 2
+
+        sdf_mean = torch.abs(sdf).mean()
+        sparseness_1 = (torch.abs(sdf) < 0.01).to(torch.float32).mean()
+        sparseness_2 = (torch.abs(sdf) < 0.02).to(torch.float32).mean()
+
+        # Eikonal loss
+        gradient_error_loss = gradient_error_fine
+
+        # ! the first 50k, don't use bg constraint
+        fg_bg_weight = 0.0 if iter_step < 50000 else get_weight(iter_step, self.fg_bg_weight)
+
+        # Mask loss, optional
+        # The images of DTU dataset contain large black regions (0 rgb values),
+        # can use this data prior to make fg more clean
+        background_loss = 0.0
+        fg_bg_loss = 0.0
+        if self.fg_bg_weight > 0 and torch.mean((mask < 0.5).to(torch.float32)) > 0.02:
+            weights_sum_fg = render_out['weights_sum_fg']
+            fg_bg_error = (weights_sum_fg - mask)[mask < 0.5]
+            fg_bg_loss = F.l1_loss(fg_bg_error,
+                                   torch.zeros_like(fg_bg_error).to(fg_bg_error.device),
+                                   reduction='mean')
+
+
+
+        loss = 1.0 * depth_loss + color_fine_loss + color_mlp_loss + \
+               sparse_loss * get_weight(iter_step, self.sdf_sparse_weight) + \
+               fg_bg_loss * fg_bg_weight + \
+               gradient_error_loss * self.sdf_igr_weight  # ! gradient_error_loss need a mask
+
+        losses = {
+            "loss": loss,
+            "depth_loss": depth_loss,
+            "color_fine_loss": color_fine_loss,
+            "color_mlp_loss": color_mlp_loss,
+            "gradient_error_loss": gradient_error_loss,
+            "background_loss": background_loss,
+            "sparse_loss": sparse_loss,
+            "sparseness_1": sparseness_1,
+            "sparseness_2": sparseness_2,
+            "sdf_mean": sdf_mean,
+            "psnr": psnr,
+            "psnr_mlp": psnr_mlp,
+            "weights_sum": render_out['weights_sum'],
+            "weights_sum_fg": render_out['weights_sum_fg'],
+            "alpha_sum": render_out['alpha_sum'],
+            "variance": render_out['variance'],
+            "sparse_weight": get_weight(iter_step, self.sdf_sparse_weight),
+            "fg_bg_weight": fg_bg_weight,
+            "fg_bg_loss": fg_bg_loss, # added by jha, bug of sparseNeuS
+        }
+        losses = torch.tensor(losses, device=rays_o.device)
+        return loss, losses, depth_statis
+
+    @torch.no_grad()
+    def validate_mesh(self, density_or_sdf_network, func_extract_geometry, world_space=True, resolution=360,
+                      threshold=0.0, mode='val',
+                      # * 3d feature volume
+                      conditional_volume=None, lod=None, occupancy_mask=None,
+                      bound_min=[-1, -1, -1], bound_max=[1, 1, 1], meta='', iter_step=0, scale_mat=None,
+                      trans_mat=None
+                      ):
+
+        bound_min = torch.tensor(bound_min, dtype=torch.float32)
+        bound_max = torch.tensor(bound_max, dtype=torch.float32)
+
+        vertices, triangles, fields = func_extract_geometry(
+            density_or_sdf_network,
+            bound_min, bound_max, resolution=resolution,
+            threshold=threshold, device=conditional_volume.device,
+            # * 3d feature volume
+            conditional_volume=conditional_volume, lod=lod,
+            occupancy_mask=occupancy_mask
+        )
+
+
+        if scale_mat is not None:
+            scale_mat_np = scale_mat.cpu().numpy()
+            vertices = vertices * scale_mat_np[0][0, 0] + scale_mat_np[0][:3, 3][None]
+
+        if trans_mat is not None: # w2c_ref_inv
+            trans_mat_np = trans_mat.cpu().numpy()
+            vertices_homo = np.concatenate([vertices, np.ones_like(vertices[:, :1])], axis=1)
+            vertices = np.matmul(trans_mat_np, vertices_homo[:, :, None])[:, :3, 0]
+
+        mesh = trimesh.Trimesh(vertices, triangles)
+        os.makedirs(os.path.join(self.base_exp_dir, 'meshes_' + mode), exist_ok=True)
+        mesh.export(os.path.join(self.base_exp_dir, 'meshes_' + mode,
+                                 'mesh_{:0>8d}_{}_lod{:0>1d}.ply'.format(iter_step, meta, lod)))
+
+
+
+    def validate_colored_mesh(self, density_or_sdf_network, func_extract_geometry, world_space=True, resolution=360,
+                                threshold=0.0, mode='val',
+                                # * 3d feature volume
+                                conditional_volume=None,
+                                conditional_valid_mask_volume=None,
+                                feature_maps=None,
+                                color_maps = None,
+                                w2cs=None,
+                                target_candidate_w2cs=None,
+                                intrinsics=None,
+                                rendering_network=None,
+                                rendering_projector=None,
+                                query_c2w=None,
+                                lod=None, occupancy_mask=None,
+                                bound_min=[-1, -1, -1], bound_max=[1, 1, 1], meta='', iter_step=0, scale_mat=None,
+                                trans_mat=None
+                                ):
+
+        bound_min = torch.tensor(bound_min, dtype=torch.float32)
+        bound_max = torch.tensor(bound_max, dtype=torch.float32)
+        # import time
+        # jha_begin4 = time.time()
+        vertices, triangles, fields = func_extract_geometry(
+            density_or_sdf_network,
+            bound_min, bound_max, resolution=resolution,
+            threshold=threshold, device=conditional_volume.device,
+            # * 3d feature volume
+            conditional_volume=conditional_volume, lod=lod,
+            occupancy_mask=occupancy_mask
+        )
+        # jha_end4 = time.time()
+        # print("[TEST]: func_extract_geometry time", jha_end4 - jha_begin4)
+
+        # import time
+        # jha_begin5 = time.time()
+
+
+        with torch.no_grad():
+            ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask, _, _ = rendering_projector.compute_view_independent(
+                torch.tensor(vertices).to(conditional_volume),
+                lod=lod, # JHA EDITED
+                # * 3d geometry feature volumes
+                geometryVolume=conditional_volume[0],
+                geometryVolumeMask=conditional_valid_mask_volume[0],
+                sdf_network=density_or_sdf_network,
+                # * 2d rendering feature maps
+                rendering_feature_maps=feature_maps, # [n_view, 56, 256, 256]
+                color_maps=color_maps,
+                w2cs=w2cs,
+                target_candidate_w2cs=target_candidate_w2cs,
+                intrinsics=intrinsics,
+                img_wh=[256,256],
+                query_img_idx=0,  # the index of the N_views dim for rendering
+                query_c2w=query_c2w,
+            )
+
+
+            vertices_color, rendering_valid_mask = rendering_network(
+                ren_geo_feats, ren_rgb_feats, ren_ray_diff, ren_mask)
+        
+        # jha_end5 = time.time()
+        # print("[TEST]: rendering_network time", jha_end5 - jha_begin5)
+
+        if scale_mat is not None:
+            scale_mat_np = scale_mat.cpu().numpy()
+            vertices = vertices * scale_mat_np[0][0, 0] + scale_mat_np[0][:3, 3][None]
+
+        if trans_mat is not None: # w2c_ref_inv
+            trans_mat_np = trans_mat.cpu().numpy()
+            vertices_homo = np.concatenate([vertices, np.ones_like(vertices[:, :1])], axis=1)
+            vertices = np.matmul(trans_mat_np, vertices_homo[:, :, None])[:, :3, 0]
+
+        vertices_color = np.array(vertices_color.squeeze(0).cpu() * 255, dtype=np.uint8)
+        mesh = trimesh.Trimesh(vertices, triangles, vertex_colors=vertices_color)
+        os.makedirs(os.path.join(self.base_exp_dir, 'meshes_' + mode, 'lod{:0>1d}'.format(lod)), exist_ok=True)
+        # mesh.export(os.path.join(self.base_exp_dir, 'meshes_' + mode, 'lod{:0>1d}'.format(lod),
+        #                          'mesh_{:0>8d}_{}_lod{:0>1d}.ply'.format(iter_step, meta, lod)))
+        # MODIFIED
+        mesh.export(os.path.join(self.base_exp_dir, 'meshes_' + mode, 'lod{:0>1d}'.format(lod),
+                                 'mesh_{:0>8d}_gradio_lod{:0>1d}.ply'.format(iter_step, lod)))

SparseNeuS_demo_v1/ops/back_project.py

@@ -0,0 +1,175 @@
+import torch
+from torch.nn.functional import grid_sample
+
+
+def back_project_sparse_type(coords, origin, voxel_size, feats, KRcam, sizeH=None, sizeW=None, only_mask=False,
+                             with_proj_z=False):
+    # - modified version from NeuRecon
+    '''
+    Unproject the image fetures to form a 3D (sparse) feature volume
+
+    :param coords: coordinates of voxels,
+    dim: (num of voxels, 4) (4 : batch ind, x, y, z)
+    :param origin: origin of the partial voxel volume (xyz position of voxel (0, 0, 0))
+    dim: (batch size, 3) (3: x, y, z)
+    :param voxel_size: floats specifying the size of a voxel
+    :param feats: image features
+    dim: (num of views, batch size, C, H, W)
+    :param KRcam: projection matrix
+    dim: (num of views, batch size, 4, 4)
+    :return: feature_volume_all: 3D feature volumes
+    dim: (num of voxels, num_of_views, c)
+    :return: mask_volume_all: indicate the voxel of sampled feature volume is valid or not
+    dim: (num of voxels, num_of_views)
+    '''
+    n_views, bs, c, h, w = feats.shape
+    device = feats.device
+
+    if sizeH is None:
+        sizeH, sizeW = h, w  # - if the KRcam is not suitable for the current feats
+
+    feature_volume_all = torch.zeros(coords.shape[0], n_views, c).to(device)
+    mask_volume_all = torch.zeros([coords.shape[0], n_views], dtype=torch.int32).to(device)
+    # import ipdb; ipdb.set_trace()
+    for batch in range(bs):
+        # import ipdb; ipdb.set_trace()
+        batch_ind = torch.nonzero(coords[:, 0] == batch).squeeze(1)
+        coords_batch = coords[batch_ind][:, 1:]
+
+        coords_batch = coords_batch.view(-1, 3)
+        origin_batch = origin[batch].unsqueeze(0)
+        feats_batch = feats[:, batch]
+        proj_batch = KRcam[:, batch]
+
+        grid_batch = coords_batch * voxel_size + origin_batch.float()
+        rs_grid = grid_batch.unsqueeze(0).expand(n_views, -1, -1)
+        rs_grid = rs_grid.permute(0, 2, 1).contiguous()
+        nV = rs_grid.shape[-1]
+        rs_grid = torch.cat([rs_grid, torch.ones([n_views, 1, nV]).to(device)], dim=1)
+
+        # Project grid
+        im_p = proj_batch @ rs_grid  # - transform world pts to image UV space
+        im_x, im_y, im_z = im_p[:, 0], im_p[:, 1], im_p[:, 2]
+
+        im_z[im_z >= 0] = im_z[im_z >= 0].clamp(min=1e-6)
+
+        im_x = im_x / im_z
+        im_y = im_y / im_z
+
+        im_grid = torch.stack([2 * im_x / (sizeW - 1) - 1, 2 * im_y / (sizeH - 1) - 1], dim=-1)
+        mask = im_grid.abs() <= 1
+        mask = (mask.sum(dim=-1) == 2) & (im_z > 0)
+
+        mask = mask.view(n_views, -1)
+        mask = mask.permute(1, 0).contiguous()  # [num_pts, nviews]
+
+        mask_volume_all[batch_ind] = mask.to(torch.int32)
+
+        if only_mask:
+            return mask_volume_all
+
+        feats_batch = feats_batch.view(n_views, c, h, w)
+        im_grid = im_grid.view(n_views, 1, -1, 2)
+        features = grid_sample(feats_batch, im_grid, padding_mode='zeros', align_corners=True)
+        # if features.isnan().sum() > 0:
+        #     import ipdb; ipdb.set_trace()
+        features = features.view(n_views, c, -1)
+        features = features.permute(2, 0, 1).contiguous()  # [num_pts, nviews, c]
+
+        feature_volume_all[batch_ind] = features
+
+        if with_proj_z:
+            im_z = im_z.view(n_views, 1, -1).permute(2, 0, 1).contiguous()  # [num_pts, nviews, 1]
+            return feature_volume_all, mask_volume_all, im_z
+    # if feature_volume_all.isnan().sum() > 0:
+    #     import ipdb; ipdb.set_trace()
+    return feature_volume_all, mask_volume_all
+
+
+def cam2pixel(cam_coords, proj_c2p_rot, proj_c2p_tr, padding_mode, sizeH=None, sizeW=None, with_depth=False):
+    """Transform coordinates in the camera frame to the pixel frame.
+    Args:
+        cam_coords: pixel coordinates defined in the first camera coordinates system -- [B, 3, H, W]
+        proj_c2p_rot: rotation matrix of cameras -- [B, 3, 3]
+        proj_c2p_tr: translation vectors of cameras -- [B, 3, 1]
+    Returns:
+        array of [-1,1] coordinates -- [B, H, W, 2]
+    """
+    b, _, h, w = cam_coords.size()
+    if sizeH is None:
+        sizeH = h
+        sizeW = w
+
+    cam_coords_flat = cam_coords.view(b, 3, -1)  # [B, 3, H*W]
+    if proj_c2p_rot is not None:
+        pcoords = proj_c2p_rot.bmm(cam_coords_flat)
+    else:
+        pcoords = cam_coords_flat
+
+    if proj_c2p_tr is not None:
+        pcoords = pcoords + proj_c2p_tr  # [B, 3, H*W]
+    X = pcoords[:, 0]
+    Y = pcoords[:, 1]
+    Z = pcoords[:, 2].clamp(min=1e-3)
+
+    X_norm = 2 * (X / Z) / (sizeW - 1) - 1  # Normalized, -1 if on extreme left,
+    # 1 if on extreme right (x = w-1) [B, H*W]
+    Y_norm = 2 * (Y / Z) / (sizeH - 1) - 1  # Idem [B, H*W]
+    if padding_mode == 'zeros':
+        X_mask = ((X_norm > 1) + (X_norm < -1)).detach()
+        X_norm[X_mask] = 2  # make sure that no point in warped image is a combinaison of im and gray
+        Y_mask = ((Y_norm > 1) + (Y_norm < -1)).detach()
+        Y_norm[Y_mask] = 2
+
+    if with_depth:
+        pixel_coords = torch.stack([X_norm, Y_norm, Z], dim=2)  # [B, H*W, 3]
+        return pixel_coords.view(b, h, w, 3)
+    else:
+        pixel_coords = torch.stack([X_norm, Y_norm], dim=2)  # [B, H*W, 2]
+        return pixel_coords.view(b, h, w, 2)
+
+
+# * have already checked, should check whether proj_matrix is for right coordinate system and resolution
+def back_project_dense_type(coords, origin, voxel_size, feats, proj_matrix, sizeH=None, sizeW=None):
+    '''
+    Unproject the image fetures to form a 3D (dense) feature volume
+
+    :param coords: coordinates of voxels,
+    dim: (batch, nviews, 3, X,Y,Z)
+    :param origin: origin of the partial voxel volume (xyz position of voxel (0, 0, 0))
+    dim: (batch size, 3) (3: x, y, z)
+    :param voxel_size: floats specifying the size of a voxel
+    :param feats: image features
+    dim: (batch size, num of views,  C, H, W)
+    :param proj_matrix: projection matrix
+    dim: (batch size, num of views, 4, 4)
+    :return: feature_volume_all: 3D feature volumes
+    dim: (batch, nviews, C, X,Y,Z)
+    :return: count: number of times each voxel can be seen
+    dim: (batch, nviews, 1, X,Y,Z)
+    '''
+
+    batch, nviews, _, wX, wY, wZ = coords.shape
+
+    if sizeH is None:
+        sizeH, sizeW = feats.shape[-2:]
+    proj_matrix = proj_matrix.view(batch * nviews, *proj_matrix.shape[2:])
+
+    coords_wrd = coords * voxel_size + origin.view(batch, 1, 3, 1, 1, 1)
+    coords_wrd = coords_wrd.view(batch * nviews, 3, wX * wY * wZ, 1)  # (b*nviews,3,wX*wY*wZ, 1)
+
+    pixel_grids = cam2pixel(coords_wrd, proj_matrix[:, :3, :3], proj_matrix[:, :3, 3:],
+                            'zeros', sizeH=sizeH, sizeW=sizeW)  # (b*nviews,wX*wY*wZ, 2)
+    pixel_grids = pixel_grids.view(batch * nviews, 1, wX * wY * wZ, 2)
+
+    feats = feats.view(batch * nviews, *feats.shape[2:])  # (b*nviews,c,h,w)
+
+    ones = torch.ones((batch * nviews, 1, *feats.shape[2:])).to(feats.dtype).to(feats.device)
+
+    features_volume = torch.nn.functional.grid_sample(feats, pixel_grids, padding_mode='zeros', align_corners=True)
+    counts_volume = torch.nn.functional.grid_sample(ones, pixel_grids, padding_mode='zeros', align_corners=True)
+
+    features_volume = features_volume.view(batch, nviews, -1, wX, wY, wZ)  # (batch, nviews, C, X,Y,Z)
+    counts_volume = counts_volume.view(batch, nviews, -1, wX, wY, wZ)
+    return features_volume, counts_volume
+

SparseNeuS_demo_v1/ops/generate_grids.py

@@ -0,0 +1,33 @@
+import torch
+
+
+def generate_grid(n_vox, interval):
+    """
+    generate grid
+    if 3D volume, grid[:,:,x,y,z]  = (x,y,z)
+    :param n_vox:
+    :param interval:
+    :return:
+    """
+    with torch.no_grad():
+        # Create voxel grid
+        grid_range = [torch.arange(0, n_vox[axis], interval) for axis in range(3)]
+        grid = torch.stack(torch.meshgrid(grid_range[0], grid_range[1], grid_range[2], indexing="ij"))  # 3 dx dy dz
+        # ! don't create tensor on gpu; imbalanced gpu memory in ddp mode
+        grid = grid.unsqueeze(0).type(torch.float32)  # 1 3 dx dy dz
+
+    return grid
+
+
+if __name__ == "__main__":
+    import torch.nn.functional as F
+    grid = generate_grid([5, 6, 8], 1)
+
+    pts = 2 * torch.tensor([1, 2, 3]) / (torch.tensor([5, 6, 8]) - 1) - 1
+    pts = pts.view(1, 1, 1, 1, 3)
+
+    pts = torch.flip(pts, dims=[-1])
+
+    sampled = F.grid_sample(grid, pts, mode='nearest')
+
+    print(sampled)

SparseNeuS_demo_v1/ops/grid_sampler.py

@@ -0,0 +1,467 @@
+"""
+pytorch grid_sample doesn't support second-order derivative
+implement custom version
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+def grid_sample_2d(image, optical):
+    N, C, IH, IW = image.shape
+    _, H, W, _ = optical.shape
+
+    ix = optical[..., 0]
+    iy = optical[..., 1]
+
+    ix = ((ix + 1) / 2) * (IW - 1);
+    iy = ((iy + 1) / 2) * (IH - 1);
+    with torch.no_grad():
+        ix_nw = torch.floor(ix);
+        iy_nw = torch.floor(iy);
+        ix_ne = ix_nw + 1;
+        iy_ne = iy_nw;
+        ix_sw = ix_nw;
+        iy_sw = iy_nw + 1;
+        ix_se = ix_nw + 1;
+        iy_se = iy_nw + 1;
+
+    nw = (ix_se - ix) * (iy_se - iy)
+    ne = (ix - ix_sw) * (iy_sw - iy)
+    sw = (ix_ne - ix) * (iy - iy_ne)
+    se = (ix - ix_nw) * (iy - iy_nw)
+
+    with torch.no_grad():
+        torch.clamp(ix_nw, 0, IW - 1, out=ix_nw)
+        torch.clamp(iy_nw, 0, IH - 1, out=iy_nw)
+
+        torch.clamp(ix_ne, 0, IW - 1, out=ix_ne)
+        torch.clamp(iy_ne, 0, IH - 1, out=iy_ne)
+
+        torch.clamp(ix_sw, 0, IW - 1, out=ix_sw)
+        torch.clamp(iy_sw, 0, IH - 1, out=iy_sw)
+
+        torch.clamp(ix_se, 0, IW - 1, out=ix_se)
+        torch.clamp(iy_se, 0, IH - 1, out=iy_se)
+
+    image = image.view(N, C, IH * IW)
+
+    nw_val = torch.gather(image, 2, (iy_nw * IW + ix_nw).long().view(N, 1, H * W).repeat(1, C, 1))
+    ne_val = torch.gather(image, 2, (iy_ne * IW + ix_ne).long().view(N, 1, H * W).repeat(1, C, 1))
+    sw_val = torch.gather(image, 2, (iy_sw * IW + ix_sw).long().view(N, 1, H * W).repeat(1, C, 1))
+    se_val = torch.gather(image, 2, (iy_se * IW + ix_se).long().view(N, 1, H * W).repeat(1, C, 1))
+
+    out_val = (nw_val.view(N, C, H, W) * nw.view(N, 1, H, W) +
+               ne_val.view(N, C, H, W) * ne.view(N, 1, H, W) +
+               sw_val.view(N, C, H, W) * sw.view(N, 1, H, W) +
+               se_val.view(N, C, H, W) * se.view(N, 1, H, W))
+
+    return out_val
+
+
+# - checked for correctness
+def grid_sample_3d(volume, optical):
+    """
+    bilinear sampling cannot guarantee continuous first-order gradient
+    mimic pytorch grid_sample function
+    The 8 corner points of a volume noted as: 4 points (front view); 4 points (back view)
+    fnw (front north west) point
+    bse (back south east) point
+    :param volume: [B, C, X, Y, Z]
+    :param optical: [B, x, y, z, 3]
+    :return:
+    """
+    N, C, ID, IH, IW = volume.shape
+    _, D, H, W, _ = optical.shape
+
+    ix = optical[..., 0]
+    iy = optical[..., 1]
+    iz = optical[..., 2]
+
+    ix = ((ix + 1) / 2) * (IW - 1)
+    iy = ((iy + 1) / 2) * (IH - 1)
+    iz = ((iz + 1) / 2) * (ID - 1)
+
+    mask_x = (ix > 0) & (ix < IW)
+    mask_y = (iy > 0) & (iy < IH)
+    mask_z = (iz > 0) & (iz < ID)
+
+    mask = mask_x & mask_y & mask_z  # [B, x, y, z]
+    mask = mask[:, None, :, :, :].repeat(1, C, 1, 1, 1)  # [B, C, x, y, z]
+
+    with torch.no_grad():
+        # back north west
+        ix_bnw = torch.floor(ix)
+        iy_bnw = torch.floor(iy)
+        iz_bnw = torch.floor(iz)
+
+        ix_bne = ix_bnw + 1
+        iy_bne = iy_bnw
+        iz_bne = iz_bnw
+
+        ix_bsw = ix_bnw
+        iy_bsw = iy_bnw + 1
+        iz_bsw = iz_bnw
+
+        ix_bse = ix_bnw + 1
+        iy_bse = iy_bnw + 1
+        iz_bse = iz_bnw
+
+        # front view
+        ix_fnw = ix_bnw
+        iy_fnw = iy_bnw
+        iz_fnw = iz_bnw + 1
+
+        ix_fne = ix_bnw + 1
+        iy_fne = iy_bnw
+        iz_fne = iz_bnw + 1
+
+        ix_fsw = ix_bnw
+        iy_fsw = iy_bnw + 1
+        iz_fsw = iz_bnw + 1
+
+        ix_fse = ix_bnw + 1
+        iy_fse = iy_bnw + 1
+        iz_fse = iz_bnw + 1
+
+    # back view
+    bnw = (ix_fse - ix) * (iy_fse - iy) * (iz_fse - iz)  # smaller volume, larger weight
+    bne = (ix - ix_fsw) * (iy_fsw - iy) * (iz_fsw - iz)
+    bsw = (ix_fne - ix) * (iy - iy_fne) * (iz_fne - iz)
+    bse = (ix - ix_fnw) * (iy - iy_fnw) * (iz_fnw - iz)
+
+    # front view
+    fnw = (ix_bse - ix) * (iy_bse - iy) * (iz - iz_bse)  # smaller volume, larger weight
+    fne = (ix - ix_bsw) * (iy_bsw - iy) * (iz - iz_bsw)
+    fsw = (ix_bne - ix) * (iy - iy_bne) * (iz - iz_bne)
+    fse = (ix - ix_bnw) * (iy - iy_bnw) * (iz - iz_bnw)
+
+    with torch.no_grad():
+        # back view
+        torch.clamp(ix_bnw, 0, IW - 1, out=ix_bnw)
+        torch.clamp(iy_bnw, 0, IH - 1, out=iy_bnw)
+        torch.clamp(iz_bnw, 0, ID - 1, out=iz_bnw)
+
+        torch.clamp(ix_bne, 0, IW - 1, out=ix_bne)
+        torch.clamp(iy_bne, 0, IH - 1, out=iy_bne)
+        torch.clamp(iz_bne, 0, ID - 1, out=iz_bne)
+
+        torch.clamp(ix_bsw, 0, IW - 1, out=ix_bsw)
+        torch.clamp(iy_bsw, 0, IH - 1, out=iy_bsw)
+        torch.clamp(iz_bsw, 0, ID - 1, out=iz_bsw)
+
+        torch.clamp(ix_bse, 0, IW - 1, out=ix_bse)
+        torch.clamp(iy_bse, 0, IH - 1, out=iy_bse)
+        torch.clamp(iz_bse, 0, ID - 1, out=iz_bse)
+
+        # front view
+        torch.clamp(ix_fnw, 0, IW - 1, out=ix_fnw)
+        torch.clamp(iy_fnw, 0, IH - 1, out=iy_fnw)
+        torch.clamp(iz_fnw, 0, ID - 1, out=iz_fnw)
+
+        torch.clamp(ix_fne, 0, IW - 1, out=ix_fne)
+        torch.clamp(iy_fne, 0, IH - 1, out=iy_fne)
+        torch.clamp(iz_fne, 0, ID - 1, out=iz_fne)
+
+        torch.clamp(ix_fsw, 0, IW - 1, out=ix_fsw)
+        torch.clamp(iy_fsw, 0, IH - 1, out=iy_fsw)
+        torch.clamp(iz_fsw, 0, ID - 1, out=iz_fsw)
+
+        torch.clamp(ix_fse, 0, IW - 1, out=ix_fse)
+        torch.clamp(iy_fse, 0, IH - 1, out=iy_fse)
+        torch.clamp(iz_fse, 0, ID - 1, out=iz_fse)
+
+    # xxx = volume[:, :, iz_bnw.long(), iy_bnw.long(), ix_bnw.long()]
+    volume = volume.view(N, C, ID * IH * IW)
+    # yyy = volume[:, :, (iz_bnw * ID + iy_bnw * IW + ix_bnw).long()]
+
+    # back view
+    bnw_val = torch.gather(volume, 2,
+                           (iz_bnw * ID ** 2 + iy_bnw * IW + ix_bnw).long().view(N, 1, D * H * W).repeat(1, C, 1))
+    bne_val = torch.gather(volume, 2,
+                           (iz_bne * ID ** 2 + iy_bne * IW + ix_bne).long().view(N, 1, D * H * W).repeat(1, C, 1))
+    bsw_val = torch.gather(volume, 2,
+                           (iz_bsw * ID ** 2 + iy_bsw * IW + ix_bsw).long().view(N, 1, D * H * W).repeat(1, C, 1))
+    bse_val = torch.gather(volume, 2,
+                           (iz_bse * ID ** 2 + iy_bse * IW + ix_bse).long().view(N, 1, D * H * W).repeat(1, C, 1))
+
+    # front view
+    fnw_val = torch.gather(volume, 2,
+                           (iz_fnw * ID ** 2 + iy_fnw * IW + ix_fnw).long().view(N, 1, D * H * W).repeat(1, C, 1))
+    fne_val = torch.gather(volume, 2,
+                           (iz_fne * ID ** 2 + iy_fne * IW + ix_fne).long().view(N, 1, D * H * W).repeat(1, C, 1))
+    fsw_val = torch.gather(volume, 2,
+                           (iz_fsw * ID ** 2 + iy_fsw * IW + ix_fsw).long().view(N, 1, D * H * W).repeat(1, C, 1))
+    fse_val = torch.gather(volume, 2,
+                           (iz_fse * ID ** 2 + iy_fse * IW + ix_fse).long().view(N, 1, D * H * W).repeat(1, C, 1))
+
+    out_val = (
+        # back
+            bnw_val.view(N, C, D, H, W) * bnw.view(N, 1, D, H, W) +
+            bne_val.view(N, C, D, H, W) * bne.view(N, 1, D, H, W) +
+            bsw_val.view(N, C, D, H, W) * bsw.view(N, 1, D, H, W) +
+            bse_val.view(N, C, D, H, W) * bse.view(N, 1, D, H, W) +
+            # front
+            fnw_val.view(N, C, D, H, W) * fnw.view(N, 1, D, H, W) +
+            fne_val.view(N, C, D, H, W) * fne.view(N, 1, D, H, W) +
+            fsw_val.view(N, C, D, H, W) * fsw.view(N, 1, D, H, W) +
+            fse_val.view(N, C, D, H, W) * fse.view(N, 1, D, H, W)
+
+    )
+
+    # * zero padding
+    out_val = torch.where(mask, out_val, torch.zeros_like(out_val).float().to(out_val.device))
+
+    return out_val
+
+
+# Interpolation kernel
+def get_weight(s, a=-0.5):
+    mask_0 = (torch.abs(s) >= 0) & (torch.abs(s) <= 1)
+    mask_1 = (torch.abs(s) > 1) & (torch.abs(s) <= 2)
+    mask_2 = torch.abs(s) > 2
+
+    weight = torch.zeros_like(s).to(s.device)
+    weight = torch.where(mask_0, (a + 2) * (torch.abs(s) ** 3) - (a + 3) * (torch.abs(s) ** 2) + 1, weight)
+    weight = torch.where(mask_1,
+                         a * (torch.abs(s) ** 3) - (5 * a) * (torch.abs(s) ** 2) + (8 * a) * torch.abs(s) - 4 * a,
+                         weight)
+
+    # if (torch.abs(s) >= 0) & (torch.abs(s) <= 1):
+    #     return (a + 2) * (torch.abs(s) ** 3) - (a + 3) * (torch.abs(s) ** 2) + 1
+    #
+    # elif (torch.abs(s) > 1) & (torch.abs(s) <= 2):
+    #     return a * (torch.abs(s) ** 3) - (5 * a) * (torch.abs(s) ** 2) + (8 * a) * torch.abs(s) - 4 * a
+    # return 0
+
+    return weight
+
+
+def cubic_interpolate(p, x):
+    """
+    one dimensional cubic interpolation
+    :param p: [N, 4]  (4) should be in order
+    :param x: [N]
+    :return:
+    """
+    return p[:, 1] + 0.5 * x * (p[:, 2] - p[:, 0] + x * (
+            2.0 * p[:, 0] - 5.0 * p[:, 1] + 4.0 * p[:, 2] - p[:, 3] + x * (
+            3.0 * (p[:, 1] - p[:, 2]) + p[:, 3] - p[:, 0])))
+
+
+def bicubic_interpolate(p, x, y, if_batch=True):
+    """
+    two dimensional cubic interpolation
+    :param p: [N, 4, 4]
+    :param x: [N]
+    :param y: [N]
+    :return:
+    """
+    num = p.shape[0]
+
+    if not if_batch:
+        arr0 = cubic_interpolate(p[:, 0, :], x)  # [N]
+        arr1 = cubic_interpolate(p[:, 1, :], x)
+        arr2 = cubic_interpolate(p[:, 2, :], x)
+        arr3 = cubic_interpolate(p[:, 3, :], x)
+        return cubic_interpolate(torch.stack([arr0, arr1, arr2, arr3], dim=-1), y)  # [N]
+    else:
+        x = x[:, None].repeat(1, 4).view(-1)
+        p = p.contiguous().view(num * 4, 4)
+        arr = cubic_interpolate(p, x)
+        arr = arr.view(num, 4)
+
+        return cubic_interpolate(arr, y)
+
+
+def tricubic_interpolate(p, x, y, z):
+    """
+    three dimensional cubic interpolation
+    :param p: [N,4,4,4]
+    :param x: [N]
+    :param y: [N]
+    :param z: [N]
+    :return:
+    """
+    num = p.shape[0]
+
+    arr0 = bicubic_interpolate(p[:, 0, :, :], x, y)  # [N]
+    arr1 = bicubic_interpolate(p[:, 1, :, :], x, y)
+    arr2 = bicubic_interpolate(p[:, 2, :, :], x, y)
+    arr3 = bicubic_interpolate(p[:, 3, :, :], x, y)
+
+    return cubic_interpolate(torch.stack([arr0, arr1, arr2, arr3], dim=-1), z)  # [N]
+
+
+def cubic_interpolate_batch(p, x):
+    """
+    one dimensional cubic interpolation
+    :param p: [B, N, 4]  (4) should be in order
+    :param x: [B, N]
+    :return:
+    """
+    return p[:, :, 1] + 0.5 * x * (p[:, :, 2] - p[:, :, 0] + x * (
+            2.0 * p[:, :, 0] - 5.0 * p[:, :, 1] + 4.0 * p[:, :, 2] - p[:, :, 3] + x * (
+            3.0 * (p[:, :, 1] - p[:, :, 2]) + p[:, :, 3] - p[:, :, 0])))
+
+
+def bicubic_interpolate_batch(p, x, y):
+    """
+    two dimensional cubic interpolation
+    :param p: [B, N, 4, 4]
+    :param x: [B, N]
+    :param y: [B, N]
+    :return:
+    """
+    B, N, _, _ = p.shape
+
+    x = x[:, :, None].repeat(1, 1, 4).view(B, N * 4)  # [B, N*4]
+    arr = cubic_interpolate_batch(p.contiguous().view(B, N * 4, 4), x)
+    arr = arr.view(B, N, 4)
+    return cubic_interpolate_batch(arr, y)  # [B, N]
+
+
+# * batch version cannot speed up training
+def tricubic_interpolate_batch(p, x, y, z):
+    """
+    three dimensional cubic interpolation
+    :param p: [N,4,4,4]
+    :param x: [N]
+    :param y: [N]
+    :param z: [N]
+    :return:
+    """
+    N = p.shape[0]
+
+    x = x[None, :].repeat(4, 1)
+    y = y[None, :].repeat(4, 1)
+
+    p = p.permute(1, 0, 2, 3).contiguous()
+
+    arr = bicubic_interpolate_batch(p[:, :, :, :], x, y)  # [4, N]
+
+    arr = arr.permute(1, 0).contiguous()  # [N, 4]
+
+    return cubic_interpolate(arr, z)  # [N]
+
+
+def tricubic_sample_3d(volume, optical):
+    """
+    tricubic sampling; can guarantee continuous gradient  (interpolation border)
+    :param volume: [B, C, ID, IH, IW]
+    :param optical: [B, D, H, W, 3]
+    :param sample_num:
+    :return:
+    """
+
+    @torch.no_grad()
+    def get_shifts(x):
+        x1 = -1 * (1 + x - torch.floor(x))
+        x2 = -1 * (x - torch.floor(x))
+        x3 = torch.floor(x) + 1 - x
+        x4 = torch.floor(x) + 2 - x
+
+        return torch.stack([x1, x2, x3, x4], dim=-1)  # (B,d,h,w,4)
+
+    N, C, ID, IH, IW = volume.shape
+    _, D, H, W, _ = optical.shape
+
+    device = volume.device
+
+    ix = optical[..., 0]
+    iy = optical[..., 1]
+    iz = optical[..., 2]
+
+    ix = ((ix + 1) / 2) * (IW - 1)  # (B,d,h,w)
+    iy = ((iy + 1) / 2) * (IH - 1)
+    iz = ((iz + 1) / 2) * (ID - 1)
+
+    ix = ix.view(-1)
+    iy = iy.view(-1)
+    iz = iz.view(-1)
+
+    with torch.no_grad():
+        shifts_x = get_shifts(ix).view(-1, 4)  # (B*d*h*w,4)
+        shifts_y = get_shifts(iy).view(-1, 4)
+        shifts_z = get_shifts(iz).view(-1, 4)
+
+        perm_weights = torch.ones([N * D * H * W, 4 * 4 * 4]).long().to(device)
+        perm = torch.cumsum(perm_weights, dim=-1) - 1  # (B*d*h*w,64)
+
+        perm_z = perm // 16  # [N*D*H*W, num]
+        perm_y = (perm - perm_z * 16) // 4
+        perm_x = (perm - perm_z * 16 - perm_y * 4)
+
+        shifts_x = torch.gather(shifts_x, 1, perm_x)  # [N*D*H*W, num]
+        shifts_y = torch.gather(shifts_y, 1, perm_y)
+        shifts_z = torch.gather(shifts_z, 1, perm_z)
+
+        ix_target = (ix[:, None] + shifts_x).long()  # [N*D*H*W, num]
+        iy_target = (iy[:, None] + shifts_y).long()
+        iz_target = (iz[:, None] + shifts_z).long()
+
+        torch.clamp(ix_target, 0, IW - 1, out=ix_target)
+        torch.clamp(iy_target, 0, IH - 1, out=iy_target)
+        torch.clamp(iz_target, 0, ID - 1, out=iz_target)
+
+    local_dist_x = ix - ix_target[:, 1]  # ! attention here is [:, 1]
+    local_dist_y = iy - iy_target[:, 1 + 4]
+    local_dist_z = iz - iz_target[:, 1 + 16]
+
+    local_dist_x = local_dist_x.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
+    local_dist_y = local_dist_y.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
+    local_dist_z = local_dist_z.view(N, 1, D * H * W).repeat(1, C, 1).view(-1)
+
+    # ! attention: IW is correct
+    idx_target = iz_target * ID ** 2 + iy_target * IW + ix_target  # [N*D*H*W, num]
+
+    volume = volume.view(N, C, ID * IH * IW)
+
+    out = torch.gather(volume, 2,
+                       idx_target.view(N, 1, D * H * W * 64).repeat(1, C, 1))
+    out = out.view(N * C * D * H * W, 4, 4, 4)
+
+    # - tricubic_interpolate() is a bit faster than tricubic_interpolate_batch()
+    final = tricubic_interpolate(out, local_dist_x, local_dist_y, local_dist_z).view(N, C, D, H, W)  # [N,C,D,H,W]
+
+    return final
+
+
+
+if __name__ == "__main__":
+    # image = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).view(1, 3, 1, 3)
+    #
+    # optical = torch.Tensor([0.9, 0.5, 0.6, -0.7]).view(1, 1, 2, 2)
+    #
+    # print(grid_sample_2d(image, optical))
+    #
+    # print(F.grid_sample(image, optical, padding_mode='border', align_corners=True))
+
+    from ops.generate_grids import generate_grid
+
+    p = torch.tensor([x for x in range(4)]).view(1, 4).float()
+
+    v = cubic_interpolate(p, torch.tensor([0.5]).view(1))
+    # v = bicubic_interpolate(p, torch.tensor([2/3]).view(1) , torch.tensor([2/3]).view(1))
+
+    vsize = 9
+    volume = generate_grid([vsize, vsize, vsize], 1)  # [1,3,10,10,10]
+    # volume = torch.tensor([x for x in range(1000)]).view(1, 1, 10, 10, 10).float()
+    X, Y, Z = 0, 0, 6
+    x = 2 * X / (vsize - 1) - 1
+    y = 2 * Y / (vsize - 1) - 1
+    z = 2 * Z / (vsize - 1) - 1
+
+    # print(volume[:, :, Z, Y, X])
+
+    # volume = volume.view(1, 3, -1)
+    # xx = volume[:, :, Z * 9*9 + Y * 9 + X]
+
+    optical = torch.Tensor([-0.6, -0.7, 0.5, 0.3, 0.5, 0.5]).view(1, 1, 1, 2, 3)
+
+    print(F.grid_sample(volume, optical, padding_mode='border', align_corners=True))
+    print(grid_sample_3d(volume, optical))
+    print(tricubic_sample_3d(volume, optical))
+    # target, relative_coords = implicit_sample_3d(volume, optical, 1)
+    # print(target)

SparseNeuS_demo_v1/tsparse/modules.py

@@ -0,0 +1,326 @@
+import torch
+import torch.nn as nn
+import torchsparse
+import torchsparse.nn as spnn
+from torchsparse.tensor import PointTensor
+
+from tsparse.torchsparse_utils import *
+
+
+# __all__ = ['SPVCNN', 'SConv3d', 'SparseConvGRU']
+
+
+class ConvBnReLU(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=3, stride=1, pad=1):
+        super(ConvBnReLU, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels,
+                              kernel_size, stride=stride, padding=pad, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.activation = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.activation(self.bn(self.conv(x)))
+
+
+class ConvBnReLU3D(nn.Module):
+    def __init__(self, in_channels, out_channels,
+                 kernel_size=3, stride=1, pad=1):
+        super(ConvBnReLU3D, self).__init__()
+        self.conv = nn.Conv3d(in_channels, out_channels,
+                              kernel_size, stride=stride, padding=pad, bias=False)
+        self.bn = nn.BatchNorm3d(out_channels)
+        self.activation = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.activation(self.bn(self.conv(x)))
+
+
+###################################  feature net  ######################################
+class FeatureNet(nn.Module):
+    """
+    output 3 levels of features using a FPN structure
+    """
+
+    def __init__(self):
+        super(FeatureNet, self).__init__()
+
+        self.conv0 = nn.Sequential(
+            ConvBnReLU(3, 8, 3, 1, 1),
+            ConvBnReLU(8, 8, 3, 1, 1))
+
+        self.conv1 = nn.Sequential(
+            ConvBnReLU(8, 16, 5, 2, 2),
+            ConvBnReLU(16, 16, 3, 1, 1),
+            ConvBnReLU(16, 16, 3, 1, 1))
+
+        self.conv2 = nn.Sequential(
+            ConvBnReLU(16, 32, 5, 2, 2),
+            ConvBnReLU(32, 32, 3, 1, 1),
+            ConvBnReLU(32, 32, 3, 1, 1))
+
+        self.toplayer = nn.Conv2d(32, 32, 1)
+        self.lat1 = nn.Conv2d(16, 32, 1)
+        self.lat0 = nn.Conv2d(8, 32, 1)
+
+        # to reduce channel size of the outputs from FPN
+        self.smooth1 = nn.Conv2d(32, 16, 3, padding=1)
+        self.smooth0 = nn.Conv2d(32, 8, 3, padding=1)
+
+    def _upsample_add(self, x, y):
+        return torch.nn.functional.interpolate(x, scale_factor=2,
+                                               mode="bilinear", align_corners=True) + y
+
+    def forward(self, x):
+        # x: (B, 3, H, W)
+        conv0 = self.conv0(x)  # (B, 8, H, W)
+        conv1 = self.conv1(conv0)  # (B, 16, H//2, W//2)
+        conv2 = self.conv2(conv1)  # (B, 32, H//4, W//4)
+        feat2 = self.toplayer(conv2)  # (B, 32, H//4, W//4)
+        feat1 = self._upsample_add(feat2, self.lat1(conv1))  # (B, 32, H//2, W//2)
+        feat0 = self._upsample_add(feat1, self.lat0(conv0))  # (B, 32, H, W)
+
+        # reduce output channels
+        feat1 = self.smooth1(feat1)  # (B, 16, H//2, W//2)
+        feat0 = self.smooth0(feat0)  # (B, 8, H, W)
+
+        # feats = {"level_0": feat0,
+        #          "level_1": feat1,
+        #          "level_2": feat2}
+
+        return [feat2, feat1, feat0]  # coarser to finer features
+
+
+class BasicSparseConvolutionBlock(nn.Module):
+    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
+        super().__init__()
+        self.net = nn.Sequential(
+            spnn.Conv3d(inc,
+                        outc,
+                        kernel_size=ks,
+                        dilation=dilation,
+                        stride=stride),
+            spnn.BatchNorm(outc),
+            spnn.ReLU(True))
+
+    def forward(self, x):
+        out = self.net(x)
+        return out
+
+
+class BasicSparseDeconvolutionBlock(nn.Module):
+    def __init__(self, inc, outc, ks=3, stride=1):
+        super().__init__()
+        self.net = nn.Sequential(
+            spnn.Conv3d(inc,
+                        outc,
+                        kernel_size=ks,
+                        stride=stride,
+                        transposed=True),
+            spnn.BatchNorm(outc),
+            spnn.ReLU(True))
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class SparseResidualBlock(nn.Module):
+    def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
+        super().__init__()
+        self.net = nn.Sequential(
+            spnn.Conv3d(inc,
+                        outc,
+                        kernel_size=ks,
+                        dilation=dilation,
+                        stride=stride), spnn.BatchNorm(outc),
+            spnn.ReLU(True),
+            spnn.Conv3d(outc,
+                        outc,
+                        kernel_size=ks,
+                        dilation=dilation,
+                        stride=1), spnn.BatchNorm(outc))
+
+        self.downsample = nn.Sequential() if (inc == outc and stride == 1) else \
+            nn.Sequential(
+                spnn.Conv3d(inc, outc, kernel_size=1, dilation=1, stride=stride),
+                spnn.BatchNorm(outc)
+            )
+
+        self.relu = spnn.ReLU(True)
+
+    def forward(self, x):
+        out = self.relu(self.net(x) + self.downsample(x))
+        return out
+
+
+class SPVCNN(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+
+        self.dropout = kwargs['dropout']
+
+        cr = kwargs.get('cr', 1.0)
+        cs = [32, 64, 128, 96, 96]
+        cs = [int(cr * x) for x in cs]
+
+        if 'pres' in kwargs and 'vres' in kwargs:
+            self.pres = kwargs['pres']
+            self.vres = kwargs['vres']
+
+        self.stem = nn.Sequential(
+            spnn.Conv3d(kwargs['in_channels'], cs[0], kernel_size=3, stride=1),
+            spnn.BatchNorm(cs[0]), spnn.ReLU(True)
+        )
+
+        self.stage1 = nn.Sequential(
+            BasicSparseConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
+            SparseResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
+            SparseResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
+        )
+
+        self.stage2 = nn.Sequential(
+            BasicSparseConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
+            SparseResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
+            SparseResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
+        )
+
+        self.up1 = nn.ModuleList([
+            BasicSparseDeconvolutionBlock(cs[2], cs[3], ks=2, stride=2),
+            nn.Sequential(
+                SparseResidualBlock(cs[3] + cs[1], cs[3], ks=3, stride=1,
+                                    dilation=1),
+                SparseResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
+            )
+        ])
+
+        self.up2 = nn.ModuleList([
+            BasicSparseDeconvolutionBlock(cs[3], cs[4], ks=2, stride=2),
+            nn.Sequential(
+                SparseResidualBlock(cs[4] + cs[0], cs[4], ks=3, stride=1,
+                                    dilation=1),
+                SparseResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
+            )
+        ])
+
+        self.point_transforms = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(cs[0], cs[2]),
+                nn.BatchNorm1d(cs[2]),
+                nn.ReLU(True),
+            ),
+            nn.Sequential(
+                nn.Linear(cs[2], cs[4]),
+                nn.BatchNorm1d(cs[4]),
+                nn.ReLU(True),
+            )
+        ])
+
+        self.weight_initialization()
+
+        if self.dropout:
+            self.dropout = nn.Dropout(0.3, True)
+
+    def weight_initialization(self):
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, z):
+        # x: SparseTensor z: PointTensor
+        x0 = initial_voxelize(z, self.pres, self.vres)
+
+        x0 = self.stem(x0)
+        z0 = voxel_to_point(x0, z, nearest=False)
+        z0.F = z0.F
+
+        x1 = point_to_voxel(x0, z0)
+        x1 = self.stage1(x1)
+        x2 = self.stage2(x1)
+        z1 = voxel_to_point(x2, z0)
+        z1.F = z1.F + self.point_transforms[0](z0.F)
+
+        y3 = point_to_voxel(x2, z1)
+        if self.dropout:
+            y3.F = self.dropout(y3.F)
+        y3 = self.up1[0](y3)
+        y3 = torchsparse.cat([y3, x1])
+        y3 = self.up1[1](y3)
+
+        y4 = self.up2[0](y3)
+        y4 = torchsparse.cat([y4, x0])
+        y4 = self.up2[1](y4)
+        z3 = voxel_to_point(y4, z1)
+        z3.F = z3.F + self.point_transforms[1](z1.F)
+
+        return z3.F
+
+
+class SparseCostRegNet(nn.Module):
+    """
+    Sparse cost regularization network;
+    require sparse tensors as input
+    """
+
+    def __init__(self, d_in, d_out=8):
+        super(SparseCostRegNet, self).__init__()
+        self.d_in = d_in
+        self.d_out = d_out
+
+        self.conv0 = BasicSparseConvolutionBlock(d_in, d_out)
+
+        self.conv1 = BasicSparseConvolutionBlock(d_out, 16, stride=2)
+        self.conv2 = BasicSparseConvolutionBlock(16, 16)
+
+        self.conv3 = BasicSparseConvolutionBlock(16, 32, stride=2)
+        self.conv4 = BasicSparseConvolutionBlock(32, 32)
+
+        self.conv5 = BasicSparseConvolutionBlock(32, 64, stride=2)
+        self.conv6 = BasicSparseConvolutionBlock(64, 64)
+
+        self.conv7 = BasicSparseDeconvolutionBlock(64, 32, ks=3, stride=2)
+
+        self.conv9 = BasicSparseDeconvolutionBlock(32, 16, ks=3, stride=2)
+
+        self.conv11 = BasicSparseDeconvolutionBlock(16, d_out, ks=3, stride=2)
+
+    def forward(self, x):
+        """
+
+        :param x: sparse tensor
+        :return: sparse tensor
+        """
+        conv0 = self.conv0(x)
+        conv2 = self.conv2(self.conv1(conv0))
+        conv4 = self.conv4(self.conv3(conv2))
+
+        x = self.conv6(self.conv5(conv4))
+        x = conv4 + self.conv7(x)
+        del conv4
+        x = conv2 + self.conv9(x)
+        del conv2
+        x = conv0 + self.conv11(x)
+        del conv0
+        return x.F
+
+
+class SConv3d(nn.Module):
+    def __init__(self, inc, outc, pres, vres, ks=3, stride=1, dilation=1):
+        super().__init__()
+        self.net = spnn.Conv3d(inc,
+                               outc,
+                               kernel_size=ks,
+                               dilation=dilation,
+                               stride=stride)
+        self.point_transforms = nn.Sequential(
+            nn.Linear(inc, outc),
+        )
+        self.pres = pres
+        self.vres = vres
+
+    def forward(self, z):
+        x = initial_voxelize(z, self.pres, self.vres)
+        x = self.net(x)
+        out = voxel_to_point(x, z, nearest=False)
+        out.F = out.F + self.point_transforms(z.F)
+        return out

SparseNeuS_demo_v1/tsparse/torchsparse_utils.py

@@ -0,0 +1,137 @@
+"""
+Copied from:
+https://github.com/mit-han-lab/spvnas/blob/b24f50379ed888d3a0e784508a809d4e92e820c0/core/models/utils.py
+"""
+import torch
+import torchsparse.nn.functional as F
+from torchsparse import PointTensor, SparseTensor
+from torchsparse.nn.utils import get_kernel_offsets
+
+import numpy as np
+
+# __all__ = ['initial_voxelize', 'point_to_voxel', 'voxel_to_point']
+
+
+# z: PointTensor
+# return: SparseTensor
+def initial_voxelize(z, init_res, after_res):
+    new_float_coord = torch.cat(
+        [(z.C[:, :3] * init_res) / after_res, z.C[:, -1].view(-1, 1)], 1)
+
+    pc_hash = F.sphash(torch.floor(new_float_coord).int())
+    sparse_hash = torch.unique(pc_hash)
+    idx_query = F.sphashquery(pc_hash, sparse_hash)
+    counts = F.spcount(idx_query.int(), len(sparse_hash))
+
+    inserted_coords = F.spvoxelize(torch.floor(new_float_coord), idx_query,
+                                   counts)
+    inserted_coords = torch.round(inserted_coords).int()
+    inserted_feat = F.spvoxelize(z.F, idx_query, counts)
+
+    new_tensor = SparseTensor(inserted_feat, inserted_coords, 1)
+    new_tensor.cmaps.setdefault(new_tensor.stride, new_tensor.coords)
+    z.additional_features['idx_query'][1] = idx_query
+    z.additional_features['counts'][1] = counts
+    z.C = new_float_coord
+
+    return new_tensor
+
+
+# x: SparseTensor, z: PointTensor
+# return: SparseTensor
+def point_to_voxel(x, z):
+    if z.additional_features is None or z.additional_features.get('idx_query') is None \
+            or z.additional_features['idx_query'].get(x.s) is None:
+        # pc_hash = hash_gpu(torch.floor(z.C).int())
+        pc_hash = F.sphash(
+            torch.cat([
+                torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
+                z.C[:, -1].int().view(-1, 1)
+            ], 1))
+        sparse_hash = F.sphash(x.C)
+        idx_query = F.sphashquery(pc_hash, sparse_hash)
+        counts = F.spcount(idx_query.int(), x.C.shape[0])
+        z.additional_features['idx_query'][x.s] = idx_query
+        z.additional_features['counts'][x.s] = counts
+    else:
+        idx_query = z.additional_features['idx_query'][x.s]
+        counts = z.additional_features['counts'][x.s]
+
+    inserted_feat = F.spvoxelize(z.F, idx_query, counts)
+    new_tensor = SparseTensor(inserted_feat, x.C, x.s)
+    new_tensor.cmaps = x.cmaps
+    new_tensor.kmaps = x.kmaps
+
+    return new_tensor
+
+
+# x: SparseTensor, z: PointTensor
+# return: PointTensor
+def voxel_to_point(x, z, nearest=False):
+    if z.idx_query is None or z.weights is None or z.idx_query.get(
+            x.s) is None or z.weights.get(x.s) is None:
+        off = get_kernel_offsets(2, x.s, 1, device=z.F.device)
+        # old_hash = kernel_hash_gpu(torch.floor(z.C).int(), off)
+        old_hash = F.sphash(
+            torch.cat([
+                torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
+                z.C[:, -1].int().view(-1, 1)
+            ], 1), off)
+        mm = x.C.to(z.F.device)
+        pc_hash = F.sphash(x.C.to(z.F.device))
+        idx_query = F.sphashquery(old_hash, pc_hash)
+        weights = F.calc_ti_weights(z.C, idx_query,
+                                    scale=x.s[0]).transpose(0, 1).contiguous()
+        idx_query = idx_query.transpose(0, 1).contiguous()
+        if nearest:
+            weights[:, 1:] = 0.
+            idx_query[:, 1:] = -1
+        new_feat = F.spdevoxelize(x.F, idx_query, weights)
+        new_tensor = PointTensor(new_feat,
+                                 z.C,
+                                 idx_query=z.idx_query,
+                                 weights=z.weights)
+        new_tensor.additional_features = z.additional_features
+        new_tensor.idx_query[x.s] = idx_query
+        new_tensor.weights[x.s] = weights
+        z.idx_query[x.s] = idx_query
+        z.weights[x.s] = weights
+
+    else:
+        new_feat = F.spdevoxelize(x.F, z.idx_query.get(x.s),
+                                  z.weights.get(x.s))   # - sparse trilinear interpoltation operation
+        new_tensor = PointTensor(new_feat,
+                                 z.C,
+                                 idx_query=z.idx_query,
+                                 weights=z.weights)
+        new_tensor.additional_features = z.additional_features
+
+    return new_tensor
+
+
+def sparse_to_dense_torch_batch(locs, values, dim, default_val):
+    dense = torch.full([dim[0], dim[1], dim[2], dim[3]], float(default_val), device=locs.device)
+    dense[locs[:, 0], locs[:, 1], locs[:, 2], locs[:, 3]] = values
+    return dense
+
+
+def sparse_to_dense_torch(locs, values, dim, default_val, device):
+    dense = torch.full([dim[0], dim[1], dim[2]], float(default_val), device=device)
+    if locs.shape[0] > 0:
+        dense[locs[:, 0], locs[:, 1], locs[:, 2]] = values
+    return dense
+
+
+def sparse_to_dense_channel(locs, values, dim, c, default_val, device):
+    locs = locs.to(torch.int64)
+    dense = torch.full([dim[0], dim[1], dim[2], c], float(default_val), device=device)
+    if locs.shape[0] > 0:
+        dense[locs[:, 0], locs[:, 1], locs[:, 2]] = values
+    return dense
+
+
+def sparse_to_dense_np(locs, values, dim, default_val):
+    dense = np.zeros([dim[0], dim[1], dim[2]], dtype=values.dtype)
+    dense.fill(default_val)
+    dense[locs[:, 0], locs[:, 1], locs[:, 2]] = values
+    return dense

SparseNeuS_demo_v1/utils/misc_utils.py

@@ -0,0 +1,219 @@
+import os, torch, cv2, re
+import numpy as np
+
+from PIL import Image
+import torch.nn.functional as F
+import torchvision.transforms as T
+
+# Misc
+img2mse = lambda x, y: torch.mean((x - y) ** 2)
+mse2psnr = lambda x: -10. * torch.log(x) / torch.log(torch.Tensor([10.]))
+to8b = lambda x: (255 * np.clip(x, 0, 1)).astype(np.uint8)
+mse2psnr2 = lambda x: -10. * np.log(x) / np.log(10.)
+
+
+def get_psnr(imgs_pred, imgs_gt):
+    psnrs = []
+    for (img, tar) in zip(imgs_pred, imgs_gt):
+        psnrs.append(mse2psnr2(np.mean((img - tar.cpu().numpy()) ** 2)))
+    return np.array(psnrs)
+
+
+def init_log(log, keys):
+    for key in keys:
+        log[key] = torch.tensor([0.0], dtype=float)
+    return log
+
+
+def visualize_depth_numpy(depth, minmax=None, cmap=cv2.COLORMAP_JET):
+    """
+    depth: (H, W)
+    """
+
+    x = np.nan_to_num(depth)  # change nan to 0
+    if minmax is None:
+        mi = np.min(x[x > 0])  # get minimum positive depth (ignore background)
+        ma = np.max(x)
+    else:
+        mi, ma = minmax
+
+    x = (x - mi) / (ma - mi + 1e-8)  # normalize to 0~1
+    x = (255 * x).astype(np.uint8)
+    x_ = cv2.applyColorMap(x, cmap)
+    return x_, [mi, ma]
+
+
+def visualize_depth(depth, minmax=None, cmap=cv2.COLORMAP_JET):
+    """
+    depth: (H, W)
+    """
+    if type(depth) is not np.ndarray:
+        depth = depth.cpu().numpy()
+
+    x = np.nan_to_num(depth)  # change nan to 0
+    if minmax is None:
+        mi = np.min(x[x > 0])  # get minimum positive depth (ignore background)
+        ma = np.max(x)
+    else:
+        mi, ma = minmax
+
+    x = (x - mi) / (ma - mi + 1e-8)  # normalize to 0~1
+    x = (255 * x).astype(np.uint8)
+    x_ = Image.fromarray(cv2.applyColorMap(x, cmap))
+    x_ = T.ToTensor()(x_)  # (3, H, W)
+    return x_, [mi, ma]
+
+
+def abs_error_numpy(depth_pred, depth_gt, mask):
+    depth_pred, depth_gt = depth_pred[mask], depth_gt[mask]
+    return np.abs(depth_pred - depth_gt)
+
+
+def abs_error(depth_pred, depth_gt, mask):
+    depth_pred, depth_gt = depth_pred[mask], depth_gt[mask]
+    err = depth_pred - depth_gt
+    return np.abs(err) if type(depth_pred) is np.ndarray else err.abs()
+
+
+def acc_threshold(depth_pred, depth_gt, mask, threshold):
+    """
+    computes the percentage of pixels whose depth error is less than @threshold
+    """
+    errors = abs_error(depth_pred, depth_gt, mask)
+    acc_mask = errors < threshold
+    return acc_mask.astype('float') if type(depth_pred) is np.ndarray else acc_mask.float()
+
+
+def to_tensor_cuda(data, device, filter):
+    for item in data.keys():
+
+        if item in filter:
+            continue
+
+        if type(data[item]) is np.ndarray:
+            data[item] = torch.tensor(data[item], dtype=torch.float32, device=device)
+        else:
+            data[item] = data[item].float().to(device)
+    return data
+
+
+def to_cuda(data, device, filter):
+    for item in data.keys():
+        if item in filter:
+            continue
+
+        data[item] = data[item].float().to(device)
+    return data
+
+
+def tensor_unsqueeze(data, filter):
+    for item in data.keys():
+        if item in filter:
+            continue
+
+        data[item] = data[item][None]
+    return data
+
+
+def filter_keys(dict):
+    dict.pop('N_samples')
+    if 'ndc' in dict.keys():
+        dict.pop('ndc')
+    if 'lindisp' in dict.keys():
+        dict.pop('lindisp')
+    return dict
+
+
+def sub_selete_data(data_batch, device, idx, filtKey=[],
+                    filtIndex=['view_ids_all', 'c2ws_all', 'scan', 'bbox', 'w2ref', 'ref2w', 'light_id', 'ckpt',
+                               'idx']):
+    data_sub_selete = {}
+    for item in data_batch.keys():
+        data_sub_selete[item] = data_batch[item][:, idx].float() if (
+                item not in filtIndex and torch.is_tensor(item) and item.dim() > 2) else data_batch[item].float()
+        if not data_sub_selete[item].is_cuda:
+            data_sub_selete[item] = data_sub_selete[item].to(device)
+    return data_sub_selete
+
+
+def detach_data(dictionary):
+    dictionary_new = {}
+    for key in dictionary.keys():
+        dictionary_new[key] = dictionary[key].detach().clone()
+    return dictionary_new
+
+
+def read_pfm(filename):
+    file = open(filename, 'rb')
+    color = None
+    width = None
+    height = None
+    scale = None
+    endian = None
+
+    header = file.readline().decode('utf-8').rstrip()
+    if header == 'PF':
+        color = True
+    elif header == 'Pf':
+        color = False
+    else:
+        raise Exception('Not a PFM file.')
+
+    dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode('utf-8'))
+    if dim_match:
+        width, height = map(int, dim_match.groups())
+    else:
+        raise Exception('Malformed PFM header.')
+
+    scale = float(file.readline().rstrip())
+    if scale < 0:  # little-endian
+        endian = '<'
+        scale = -scale
+    else:
+        endian = '>'  # big-endian
+
+    data = np.fromfile(file, endian + 'f')
+    shape = (height, width, 3) if color else (height, width)
+
+    data = np.reshape(data, shape)
+    data = np.flipud(data)
+    file.close()
+    return data, scale
+
+
+from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
+
+
+# from warmup_scheduler import GradualWarmupScheduler
+def get_scheduler(hparams, optimizer):
+    eps = 1e-8
+    if hparams.lr_scheduler == 'steplr':
+        scheduler = MultiStepLR(optimizer, milestones=hparams.decay_step,
+                                gamma=hparams.decay_gamma)
+    elif hparams.lr_scheduler == 'cosine':
+        scheduler = CosineAnnealingLR(optimizer, T_max=hparams.num_epochs, eta_min=eps)
+
+    else:
+        raise ValueError('scheduler not recognized!')
+
+    # if hparams.warmup_epochs > 0 and hparams.optimizer not in ['radam', 'ranger']:
+    #     scheduler = GradualWarmupScheduler(optimizer, multiplier=hparams.warmup_multiplier,
+    #                                        total_epoch=hparams.warmup_epochs, after_scheduler=scheduler)
+    return scheduler
+
+
+####  pairing  ####
+def get_nearest_pose_ids(tar_pose, ref_poses, num_select):
+    '''
+    Args:
+        tar_pose: target pose [N, 4, 4]
+        ref_poses: reference poses [M, 4, 4]
+        num_select: the number of nearest views to select
+    Returns: the selected indices
+    '''
+
+    dists = np.linalg.norm(tar_pose[:, None, :3, 3] - ref_poses[None, :, :3, 3], axis=-1)
+
+    sorted_ids = np.argsort(dists, axis=-1)
+    selected_ids = sorted_ids[:, :num_select]
+    return selected_ids

configs/sd-objaverse-finetune-c_concat-256.yaml

@@ -0,0 +1,117 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "image_target"
+    cond_stage_key: "image_cond"
+    image_size: 32
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: hybrid
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 100 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 8
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPImageEmbedder
+
+
+data:
+  target: ldm.data.simple.ObjaverseDataModuleFromConfig
+  params:
+    root_dir: 'views_whole_sphere'
+    batch_size: 192
+    num_workers: 16
+    total_view: 4
+    train:
+      validation: False
+      image_transforms:
+        size: 256
+
+    validation:
+      validation: True
+      image_transforms:
+        size: 256
+
+
+lightning:
+  find_unused_parameters: false
+  metrics_over_trainsteps_checkpoint: True
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 5000
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        batch_frequency: 500
+        max_images: 32
+        increase_log_steps: False
+        log_first_step: True
+        log_images_kwargs:
+          use_ema_scope: False
+          inpaint: False
+          plot_progressive_rows: False
+          plot_diffusion_rows: False
+          N: 32
+          unconditional_guidance_scale: 3.0
+          unconditional_guidance_label: [""]
+
+  trainer:
+    benchmark: True
+    val_check_interval: 5000000 # really sorry
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 1

ldm/data/base.py

@@ -0,0 +1,40 @@
+import os
+import numpy as np
+from abc import abstractmethod
+from torch.utils.data import Dataset, ConcatDataset, ChainDataset, IterableDataset
+
+
+class Txt2ImgIterableBaseDataset(IterableDataset):
+    '''
+    Define an interface to make the IterableDatasets for text2img data chainable
+    '''
+    def __init__(self, num_records=0, valid_ids=None, size=256):
+        super().__init__()
+        self.num_records = num_records
+        self.valid_ids = valid_ids
+        self.sample_ids = valid_ids
+        self.size = size
+
+        print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')
+
+    def __len__(self):
+        return self.num_records
+
+    @abstractmethod
+    def __iter__(self):
+        pass
+
+
+class PRNGMixin(object):
+    """
+    Adds a prng property which is a numpy RandomState which gets
+    reinitialized whenever the pid changes to avoid synchronized sampling
+    behavior when used in conjunction with multiprocessing.
+    """
+    @property
+    def prng(self):
+        currentpid = os.getpid()
+        if getattr(self, "_initpid", None) != currentpid:
+            self._initpid = currentpid
+            self._prng = np.random.RandomState()
+        return self._prng

ldm/data/coco.py

@@ -0,0 +1,253 @@
+import os
+import json
+import albumentations
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from torch.utils.data import Dataset
+from abc import abstractmethod
+
+
+class CocoBase(Dataset):
+    """needed for (image, caption, segmentation) pairs"""
+    def __init__(self, size=None, dataroot="", datajson="", onehot_segmentation=False, use_stuffthing=False,
+                 crop_size=None, force_no_crop=False, given_files=None, use_segmentation=True,crop_type=None):
+        self.split = self.get_split()
+        self.size = size
+        if crop_size is None:
+            self.crop_size = size
+        else:
+            self.crop_size = crop_size
+
+        assert crop_type in [None, 'random', 'center']
+        self.crop_type = crop_type
+        self.use_segmenation = use_segmentation
+        self.onehot = onehot_segmentation       # return segmentation as rgb or one hot
+        self.stuffthing = use_stuffthing        # include thing in segmentation
+        if self.onehot and not self.stuffthing:
+            raise NotImplemented("One hot mode is only supported for the "
+                                 "stuffthings version because labels are stored "
+                                 "a bit different.")
+
+        data_json = datajson
+        with open(data_json) as json_file:
+            self.json_data = json.load(json_file)
+            self.img_id_to_captions = dict()
+            self.img_id_to_filepath = dict()
+            self.img_id_to_segmentation_filepath = dict()
+
+        assert data_json.split("/")[-1] in [f"captions_train{self.year()}.json",
+                                            f"captions_val{self.year()}.json"]
+        # TODO currently hardcoded paths, would be better to follow logic in
+        # cocstuff pixelmaps
+        if self.use_segmenation:
+            if self.stuffthing:
+                self.segmentation_prefix = (
+                    f"data/cocostuffthings/val{self.year()}" if
+                    data_json.endswith(f"captions_val{self.year()}.json") else
+                    f"data/cocostuffthings/train{self.year()}")
+            else:
+                self.segmentation_prefix = (
+                    f"data/coco/annotations/stuff_val{self.year()}_pixelmaps" if
+                    data_json.endswith(f"captions_val{self.year()}.json") else
+                    f"data/coco/annotations/stuff_train{self.year()}_pixelmaps")
+
+        imagedirs = self.json_data["images"]
+        self.labels = {"image_ids": list()}
+        for imgdir in tqdm(imagedirs, desc="ImgToPath"):
+            self.img_id_to_filepath[imgdir["id"]] = os.path.join(dataroot, imgdir["file_name"])
+            self.img_id_to_captions[imgdir["id"]] = list()
+            pngfilename = imgdir["file_name"].replace("jpg", "png")
+            if self.use_segmenation:
+                self.img_id_to_segmentation_filepath[imgdir["id"]] = os.path.join(
+                    self.segmentation_prefix, pngfilename)
+            if given_files is not None:
+                if pngfilename in given_files:
+                    self.labels["image_ids"].append(imgdir["id"])
+            else:
+                self.labels["image_ids"].append(imgdir["id"])
+
+        capdirs = self.json_data["annotations"]
+        for capdir in tqdm(capdirs, desc="ImgToCaptions"):
+            # there are in average 5 captions per image
+            #self.img_id_to_captions[capdir["image_id"]].append(np.array([capdir["caption"]]))
+            self.img_id_to_captions[capdir["image_id"]].append(capdir["caption"])
+
+        self.rescaler = albumentations.SmallestMaxSize(max_size=self.size)
+        if self.split=="validation":
+            self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
+        else:
+            # default option for train is random crop
+            if self.crop_type in [None, 'random']:
+                self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size)
+            else:
+                self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size)
+        self.preprocessor = albumentations.Compose(
+            [self.rescaler, self.cropper],
+            additional_targets={"segmentation": "image"})
+        if force_no_crop:
+            self.rescaler = albumentations.Resize(height=self.size, width=self.size)
+            self.preprocessor = albumentations.Compose(
+                [self.rescaler],
+                additional_targets={"segmentation": "image"})
+
+    @abstractmethod
+    def year(self):
+        raise NotImplementedError()
+
+    def __len__(self):
+        return len(self.labels["image_ids"])
+
+    def preprocess_image(self, image_path, segmentation_path=None):
+        image = Image.open(image_path)
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+        image = np.array(image).astype(np.uint8)
+        if segmentation_path:
+            segmentation = Image.open(segmentation_path)
+            if not self.onehot and not segmentation.mode == "RGB":
+                segmentation = segmentation.convert("RGB")
+            segmentation = np.array(segmentation).astype(np.uint8)
+            if self.onehot:
+                assert self.stuffthing
+                # stored in caffe format: unlabeled==255. stuff and thing from
+                # 0-181. to be compatible with the labels in
+                # https://github.com/nightrome/cocostuff/blob/master/labels.txt
+                # we shift stuffthing one to the right and put unlabeled in zero
+                # as long as segmentation is uint8 shifting to right handles the
+                # latter too
+                assert segmentation.dtype == np.uint8
+                segmentation = segmentation + 1
+
+            processed = self.preprocessor(image=image, segmentation=segmentation)
+
+            image, segmentation = processed["image"], processed["segmentation"]
+        else:
+            image = self.preprocessor(image=image,)['image']
+
+        image = (image / 127.5 - 1.0).astype(np.float32)
+        if segmentation_path:
+            if self.onehot:
+                assert segmentation.dtype == np.uint8
+                # make it one hot
+                n_labels = 183
+                flatseg = np.ravel(segmentation)
+                onehot = np.zeros((flatseg.size, n_labels), dtype=np.bool)
+                onehot[np.arange(flatseg.size), flatseg] = True
+                onehot = onehot.reshape(segmentation.shape + (n_labels,)).astype(int)
+                segmentation = onehot
+            else:
+                segmentation = (segmentation / 127.5 - 1.0).astype(np.float32)
+            return image, segmentation
+        else:
+            return image
+
+    def __getitem__(self, i):
+        img_path = self.img_id_to_filepath[self.labels["image_ids"][i]]
+        if self.use_segmenation:
+            seg_path = self.img_id_to_segmentation_filepath[self.labels["image_ids"][i]]
+            image, segmentation = self.preprocess_image(img_path, seg_path)
+        else:
+            image = self.preprocess_image(img_path)
+        captions = self.img_id_to_captions[self.labels["image_ids"][i]]
+        # randomly draw one of all available captions per image
+        caption = captions[np.random.randint(0, len(captions))]
+        example = {"image": image,
+                   #"caption": [str(caption[0])],
+                   "caption": caption,
+                   "img_path": img_path,
+                   "filename_": img_path.split(os.sep)[-1]
+                    }
+        if self.use_segmenation:
+            example.update({"seg_path": seg_path, 'segmentation': segmentation})
+        return example
+
+
+class CocoImagesAndCaptionsTrain2017(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,):
+        super().__init__(size=size,
+                         dataroot="data/coco/train2017",
+                         datajson="data/coco/annotations/captions_train2017.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop)
+
+    def get_split(self):
+        return "train"
+
+    def year(self):
+        return '2017'
+
+
+class CocoImagesAndCaptionsValidation2017(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,
+                 given_files=None):
+        super().__init__(size=size,
+                         dataroot="data/coco/val2017",
+                         datajson="data/coco/annotations/captions_val2017.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
+                         given_files=given_files)
+
+    def get_split(self):
+        return "validation"
+
+    def year(self):
+        return '2017'
+
+
+
+class CocoImagesAndCaptionsTrain2014(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,crop_type='random'):
+        super().__init__(size=size,
+                         dataroot="data/coco/train2014",
+                         datajson="data/coco/annotations2014/annotations/captions_train2014.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
+                         use_segmentation=False,
+                         crop_type=crop_type)
+
+    def get_split(self):
+        return "train"
+
+    def year(self):
+        return '2014'
+
+class CocoImagesAndCaptionsValidation2014(CocoBase):
+    """returns a pair of (image, caption)"""
+    def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False,
+                 given_files=None,crop_type='center',**kwargs):
+        super().__init__(size=size,
+                         dataroot="data/coco/val2014",
+                         datajson="data/coco/annotations2014/annotations/captions_val2014.json",
+                         onehot_segmentation=onehot_segmentation,
+                         use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop,
+                         given_files=given_files,
+                         use_segmentation=False,
+                         crop_type=crop_type)
+
+    def get_split(self):
+        return "validation"
+
+    def year(self):
+        return '2014'
+
+if __name__ == '__main__':
+    with open("data/coco/annotations2014/annotations/captions_val2014.json", "r") as json_file:
+        json_data = json.load(json_file)
+        capdirs = json_data["annotations"]
+        import pudb; pudb.set_trace()
+    #d2 = CocoImagesAndCaptionsTrain2014(size=256)
+    d2 = CocoImagesAndCaptionsValidation2014(size=256)
+    print("constructed dataset.")
+    print(f"length of {d2.__class__.__name__}: {len(d2)}")
+
+    ex2 = d2[0]
+    # ex3 = d3[0]
+    # print(ex1["image"].shape)
+    print(ex2["image"].shape)
+    # print(ex3["image"].shape)
+    # print(ex1["segmentation"].shape)
+    print(ex2["caption"].__class__.__name__)

ldm/data/dummy.py

@@ -0,0 +1,34 @@
+import numpy as np
+import random
+import string
+from torch.utils.data import Dataset, Subset
+
+class DummyData(Dataset):
+    def __init__(self, length, size):
+        self.length = length
+        self.size = size
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i):
+        x = np.random.randn(*self.size)
+        letters = string.ascii_lowercase
+        y = ''.join(random.choice(string.ascii_lowercase) for i in range(10))
+        return {"jpg": x, "txt": y}
+
+
+class DummyDataWithEmbeddings(Dataset):
+    def __init__(self, length, size, emb_size):
+        self.length = length
+        self.size = size
+        self.emb_size = emb_size
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i):
+        x = np.random.randn(*self.size)
+        y = np.random.randn(*self.emb_size).astype(np.float32)
+        return {"jpg": x, "txt": y}
+

ldm/data/imagenet.py

@@ -0,0 +1,394 @@
+import os, yaml, pickle, shutil, tarfile, glob
+import cv2
+import albumentations
+import PIL
+import numpy as np
+import torchvision.transforms.functional as TF
+from omegaconf import OmegaConf
+from functools import partial
+from PIL import Image
+from tqdm import tqdm
+from torch.utils.data import Dataset, Subset
+
+import taming.data.utils as tdu
+from taming.data.imagenet import str_to_indices, give_synsets_from_indices, download, retrieve
+from taming.data.imagenet import ImagePaths
+
+from ldm.modules.image_degradation import degradation_fn_bsr, degradation_fn_bsr_light
+
+
+def synset2idx(path_to_yaml="data/index_synset.yaml"):
+    with open(path_to_yaml) as f:
+        di2s = yaml.load(f)
+    return dict((v,k) for k,v in di2s.items())
+
+
+class ImageNetBase(Dataset):
+    def __init__(self, config=None):
+        self.config = config or OmegaConf.create()
+        if not type(self.config)==dict:
+            self.config = OmegaConf.to_container(self.config)
+        self.keep_orig_class_label = self.config.get("keep_orig_class_label", False)
+        self.process_images = True  # if False we skip loading & processing images and self.data contains filepaths
+        self._prepare()
+        self._prepare_synset_to_human()
+        self._prepare_idx_to_synset()
+        self._prepare_human_to_integer_label()
+        self._load()
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, i):
+        return self.data[i]
+
+    def _prepare(self):
+        raise NotImplementedError()
+
+    def _filter_relpaths(self, relpaths):
+        ignore = set([
+            "n06596364_9591.JPEG",
+        ])
+        relpaths = [rpath for rpath in relpaths if not rpath.split("/")[-1] in ignore]
+        if "sub_indices" in self.config:
+            indices = str_to_indices(self.config["sub_indices"])
+            synsets = give_synsets_from_indices(indices, path_to_yaml=self.idx2syn)  # returns a list of strings
+            self.synset2idx = synset2idx(path_to_yaml=self.idx2syn)
+            files = []
+            for rpath in relpaths:
+                syn = rpath.split("/")[0]
+                if syn in synsets:
+                    files.append(rpath)
+            return files
+        else:
+            return relpaths
+
+    def _prepare_synset_to_human(self):
+        SIZE = 2655750
+        URL = "https://heibox.uni-heidelberg.de/f/9f28e956cd304264bb82/?dl=1"
+        self.human_dict = os.path.join(self.root, "synset_human.txt")
+        if (not os.path.exists(self.human_dict) or
+                not os.path.getsize(self.human_dict)==SIZE):
+            download(URL, self.human_dict)
+
+    def _prepare_idx_to_synset(self):
+        URL = "https://heibox.uni-heidelberg.de/f/d835d5b6ceda4d3aa910/?dl=1"
+        self.idx2syn = os.path.join(self.root, "index_synset.yaml")
+        if (not os.path.exists(self.idx2syn)):
+            download(URL, self.idx2syn)
+
+    def _prepare_human_to_integer_label(self):
+        URL = "https://heibox.uni-heidelberg.de/f/2362b797d5be43b883f6/?dl=1"
+        self.human2integer = os.path.join(self.root, "imagenet1000_clsidx_to_labels.txt")
+        if (not os.path.exists(self.human2integer)):
+            download(URL, self.human2integer)
+        with open(self.human2integer, "r") as f:
+            lines = f.read().splitlines()
+            assert len(lines) == 1000
+            self.human2integer_dict = dict()
+            for line in lines:
+                value, key = line.split(":")
+                self.human2integer_dict[key] = int(value)
+
+    def _load(self):
+        with open(self.txt_filelist, "r") as f:
+            self.relpaths = f.read().splitlines()
+            l1 = len(self.relpaths)
+            self.relpaths = self._filter_relpaths(self.relpaths)
+            print("Removed {} files from filelist during filtering.".format(l1 - len(self.relpaths)))
+
+        self.synsets = [p.split("/")[0] for p in self.relpaths]
+        self.abspaths = [os.path.join(self.datadir, p) for p in self.relpaths]
+
+        unique_synsets = np.unique(self.synsets)
+        class_dict = dict((synset, i) for i, synset in enumerate(unique_synsets))
+        if not self.keep_orig_class_label:
+            self.class_labels = [class_dict[s] for s in self.synsets]
+        else:
+            self.class_labels = [self.synset2idx[s] for s in self.synsets]
+
+        with open(self.human_dict, "r") as f:
+            human_dict = f.read().splitlines()
+            human_dict = dict(line.split(maxsplit=1) for line in human_dict)
+
+        self.human_labels = [human_dict[s] for s in self.synsets]
+
+        labels = {
+            "relpath": np.array(self.relpaths),
+            "synsets": np.array(self.synsets),
+            "class_label": np.array(self.class_labels),
+            "human_label": np.array(self.human_labels),
+        }
+
+        if self.process_images:
+            self.size = retrieve(self.config, "size", default=256)
+            self.data = ImagePaths(self.abspaths,
+                                   labels=labels,
+                                   size=self.size,
+                                   random_crop=self.random_crop,
+                                   )
+        else:
+            self.data = self.abspaths
+
+
+class ImageNetTrain(ImageNetBase):
+    NAME = "ILSVRC2012_train"
+    URL = "http://www.image-net.org/challenges/LSVRC/2012/"
+    AT_HASH = "a306397ccf9c2ead27155983c254227c0fd938e2"
+    FILES = [
+        "ILSVRC2012_img_train.tar",
+    ]
+    SIZES = [
+        147897477120,
+    ]
+
+    def __init__(self, process_images=True, data_root=None, **kwargs):
+        self.process_images = process_images
+        self.data_root = data_root
+        super().__init__(**kwargs)
+
+    def _prepare(self):
+        if self.data_root:
+            self.root = os.path.join(self.data_root, self.NAME)
+        else:
+            cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
+            self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
+
+        self.datadir = os.path.join(self.root, "data")
+        self.txt_filelist = os.path.join(self.root, "filelist.txt")
+        self.expected_length = 1281167
+        self.random_crop = retrieve(self.config, "ImageNetTrain/random_crop",
+                                    default=True)
+        if not tdu.is_prepared(self.root):
+            # prep
+            print("Preparing dataset {} in {}".format(self.NAME, self.root))
+
+            datadir = self.datadir
+            if not os.path.exists(datadir):
+                path = os.path.join(self.root, self.FILES[0])
+                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
+                    import academictorrents as at
+                    atpath = at.get(self.AT_HASH, datastore=self.root)
+                    assert atpath == path
+
+                print("Extracting {} to {}".format(path, datadir))
+                os.makedirs(datadir, exist_ok=True)
+                with tarfile.open(path, "r:") as tar:
+                    tar.extractall(path=datadir)
+
+                print("Extracting sub-tars.")
+                subpaths = sorted(glob.glob(os.path.join(datadir, "*.tar")))
+                for subpath in tqdm(subpaths):
+                    subdir = subpath[:-len(".tar")]
+                    os.makedirs(subdir, exist_ok=True)
+                    with tarfile.open(subpath, "r:") as tar:
+                        tar.extractall(path=subdir)
+
+            filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
+            filelist = [os.path.relpath(p, start=datadir) for p in filelist]
+            filelist = sorted(filelist)
+            filelist = "\n".join(filelist)+"\n"
+            with open(self.txt_filelist, "w") as f:
+                f.write(filelist)
+
+            tdu.mark_prepared(self.root)
+
+
+class ImageNetValidation(ImageNetBase):
+    NAME = "ILSVRC2012_validation"
+    URL = "http://www.image-net.org/challenges/LSVRC/2012/"
+    AT_HASH = "5d6d0df7ed81efd49ca99ea4737e0ae5e3a5f2e5"
+    VS_URL = "https://heibox.uni-heidelberg.de/f/3e0f6e9c624e45f2bd73/?dl=1"
+    FILES = [
+        "ILSVRC2012_img_val.tar",
+        "validation_synset.txt",
+    ]
+    SIZES = [
+        6744924160,
+        1950000,
+    ]
+
+    def __init__(self, process_images=True, data_root=None, **kwargs):
+        self.data_root = data_root
+        self.process_images = process_images
+        super().__init__(**kwargs)
+
+    def _prepare(self):
+        if self.data_root:
+            self.root = os.path.join(self.data_root, self.NAME)
+        else:
+            cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
+            self.root = os.path.join(cachedir, "autoencoders/data", self.NAME)
+        self.datadir = os.path.join(self.root, "data")
+        self.txt_filelist = os.path.join(self.root, "filelist.txt")
+        self.expected_length = 50000
+        self.random_crop = retrieve(self.config, "ImageNetValidation/random_crop",
+                                    default=False)
+        if not tdu.is_prepared(self.root):
+            # prep
+            print("Preparing dataset {} in {}".format(self.NAME, self.root))
+
+            datadir = self.datadir
+            if not os.path.exists(datadir):
+                path = os.path.join(self.root, self.FILES[0])
+                if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]:
+                    import academictorrents as at
+                    atpath = at.get(self.AT_HASH, datastore=self.root)
+                    assert atpath == path
+
+                print("Extracting {} to {}".format(path, datadir))
+                os.makedirs(datadir, exist_ok=True)
+                with tarfile.open(path, "r:") as tar:
+                    tar.extractall(path=datadir)
+
+                vspath = os.path.join(self.root, self.FILES[1])
+                if not os.path.exists(vspath) or not os.path.getsize(vspath)==self.SIZES[1]:
+                    download(self.VS_URL, vspath)
+
+                with open(vspath, "r") as f:
+                    synset_dict = f.read().splitlines()
+                    synset_dict = dict(line.split() for line in synset_dict)
+
+                print("Reorganizing into synset folders")
+                synsets = np.unique(list(synset_dict.values()))
+                for s in synsets:
+                    os.makedirs(os.path.join(datadir, s), exist_ok=True)
+                for k, v in synset_dict.items():
+                    src = os.path.join(datadir, k)
+                    dst = os.path.join(datadir, v)
+                    shutil.move(src, dst)
+
+            filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG"))
+            filelist = [os.path.relpath(p, start=datadir) for p in filelist]
+            filelist = sorted(filelist)
+            filelist = "\n".join(filelist)+"\n"
+            with open(self.txt_filelist, "w") as f:
+                f.write(filelist)
+
+            tdu.mark_prepared(self.root)
+
+
+
+class ImageNetSR(Dataset):
+    def __init__(self, size=None,
+                 degradation=None, downscale_f=4, min_crop_f=0.5, max_crop_f=1.,
+                 random_crop=True):
+        """
+        Imagenet Superresolution Dataloader
+        Performs following ops in order:
+        1.  crops a crop of size s from image either as random or center crop
+        2.  resizes crop to size with cv2.area_interpolation
+        3.  degrades resized crop with degradation_fn
+
+        :param size: resizing to size after cropping
+        :param degradation: degradation_fn, e.g. cv_bicubic or bsrgan_light
+        :param downscale_f: Low Resolution Downsample factor
+        :param min_crop_f: determines crop size s,
+          where s = c * min_img_side_len with c sampled from interval (min_crop_f, max_crop_f)
+        :param max_crop_f: ""
+        :param data_root:
+        :param random_crop:
+        """
+        self.base = self.get_base()
+        assert size
+        assert (size / downscale_f).is_integer()
+        self.size = size
+        self.LR_size = int(size / downscale_f)
+        self.min_crop_f = min_crop_f
+        self.max_crop_f = max_crop_f
+        assert(max_crop_f <= 1.)
+        self.center_crop = not random_crop
+
+        self.image_rescaler = albumentations.SmallestMaxSize(max_size=size, interpolation=cv2.INTER_AREA)
+
+        self.pil_interpolation = False # gets reset later if incase interp_op is from pillow
+
+        if degradation == "bsrgan":
+            self.degradation_process = partial(degradation_fn_bsr, sf=downscale_f)
+
+        elif degradation == "bsrgan_light":
+            self.degradation_process = partial(degradation_fn_bsr_light, sf=downscale_f)
+
+        else:
+            interpolation_fn = {
+            "cv_nearest": cv2.INTER_NEAREST,
+            "cv_bilinear": cv2.INTER_LINEAR,
+            "cv_bicubic": cv2.INTER_CUBIC,
+            "cv_area": cv2.INTER_AREA,
+            "cv_lanczos": cv2.INTER_LANCZOS4,
+            "pil_nearest": PIL.Image.NEAREST,
+            "pil_bilinear": PIL.Image.BILINEAR,
+            "pil_bicubic": PIL.Image.BICUBIC,
+            "pil_box": PIL.Image.BOX,
+            "pil_hamming": PIL.Image.HAMMING,
+            "pil_lanczos": PIL.Image.LANCZOS,
+            }[degradation]
+
+            self.pil_interpolation = degradation.startswith("pil_")
+
+            if self.pil_interpolation:
+                self.degradation_process = partial(TF.resize, size=self.LR_size, interpolation=interpolation_fn)
+
+            else:
+                self.degradation_process = albumentations.SmallestMaxSize(max_size=self.LR_size,
+                                                                          interpolation=interpolation_fn)
+
+    def __len__(self):
+        return len(self.base)
+
+    def __getitem__(self, i):
+        example = self.base[i]
+        image = Image.open(example["file_path_"])
+
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+
+        image = np.array(image).astype(np.uint8)
+
+        min_side_len = min(image.shape[:2])
+        crop_side_len = min_side_len * np.random.uniform(self.min_crop_f, self.max_crop_f, size=None)
+        crop_side_len = int(crop_side_len)
+
+        if self.center_crop:
+            self.cropper = albumentations.CenterCrop(height=crop_side_len, width=crop_side_len)
+
+        else:
+            self.cropper = albumentations.RandomCrop(height=crop_side_len, width=crop_side_len)
+
+        image = self.cropper(image=image)["image"]
+        image = self.image_rescaler(image=image)["image"]
+
+        if self.pil_interpolation:
+            image_pil = PIL.Image.fromarray(image)
+            LR_image = self.degradation_process(image_pil)
+            LR_image = np.array(LR_image).astype(np.uint8)
+
+        else:
+            LR_image = self.degradation_process(image=image)["image"]
+
+        example["image"] = (image/127.5 - 1.0).astype(np.float32)
+        example["LR_image"] = (LR_image/127.5 - 1.0).astype(np.float32)
+        example["caption"] = example["human_label"]  # dummy caption
+        return example
+
+
+class ImageNetSRTrain(ImageNetSR):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def get_base(self):
+        with open("data/imagenet_train_hr_indices.p", "rb") as f:
+            indices = pickle.load(f)
+        dset = ImageNetTrain(process_images=False,)
+        return Subset(dset, indices)
+
+
+class ImageNetSRValidation(ImageNetSR):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def get_base(self):
+        with open("data/imagenet_val_hr_indices.p", "rb") as f:
+            indices = pickle.load(f)
+        dset = ImageNetValidation(process_images=False,)
+        return Subset(dset, indices)

ldm/data/inpainting/synthetic_mask.py

@@ -0,0 +1,166 @@
+from PIL import Image, ImageDraw
+import numpy as np
+
+settings = {
+    "256narrow": {
+        "p_irr": 1,
+        "min_n_irr": 4,
+        "max_n_irr": 50,
+        "max_l_irr": 40,
+        "max_w_irr": 10,
+        "min_n_box": None,
+        "max_n_box": None,
+        "min_s_box": None,
+        "max_s_box": None,
+        "marg": None,
+    },
+    "256train": {
+        "p_irr": 0.5,
+        "min_n_irr": 1,
+        "max_n_irr": 5,
+        "max_l_irr": 200,
+        "max_w_irr": 100,
+        "min_n_box": 1,
+        "max_n_box": 4,
+        "min_s_box": 30,
+        "max_s_box": 150,
+        "marg": 10,
+    },
+    "512train": {    # TODO: experimental
+            "p_irr": 0.5,
+            "min_n_irr": 1,
+            "max_n_irr": 5,
+            "max_l_irr": 450,
+            "max_w_irr": 250,
+            "min_n_box": 1,
+            "max_n_box": 4,
+            "min_s_box": 30,
+            "max_s_box": 300,
+            "marg": 10,
+        },
+    "512train-large": {    # TODO: experimental
+            "p_irr": 0.5,
+            "min_n_irr": 1,
+            "max_n_irr": 5,
+            "max_l_irr": 450,
+            "max_w_irr": 400,
+            "min_n_box": 1,
+            "max_n_box": 4,
+            "min_s_box": 75,
+            "max_s_box": 450,
+            "marg": 10,
+        },
+}
+
+
+def gen_segment_mask(mask, start, end, brush_width):
+    mask = mask > 0
+    mask = (255 * mask).astype(np.uint8)
+    mask = Image.fromarray(mask)
+    draw = ImageDraw.Draw(mask)
+    draw.line([start, end], fill=255, width=brush_width, joint="curve")
+    mask = np.array(mask) / 255
+    return mask
+
+
+def gen_box_mask(mask, masked):
+    x_0, y_0, w, h = masked
+    mask[y_0:y_0 + h, x_0:x_0 + w] = 1
+    return mask
+
+
+def gen_round_mask(mask, masked, radius):
+    x_0, y_0, w, h = masked
+    xy = [(x_0, y_0), (x_0 + w, y_0 + w)]
+
+    mask = mask > 0
+    mask = (255 * mask).astype(np.uint8)
+    mask = Image.fromarray(mask)
+    draw = ImageDraw.Draw(mask)
+    draw.rounded_rectangle(xy, radius=radius, fill=255)
+    mask = np.array(mask) / 255
+    return mask
+
+
+def gen_large_mask(prng, img_h, img_w,
+                   marg, p_irr, min_n_irr, max_n_irr, max_l_irr, max_w_irr,
+                   min_n_box, max_n_box, min_s_box, max_s_box):
+    """
+    img_h: int, an image height
+    img_w: int, an image width
+    marg: int, a margin for a box starting coordinate
+    p_irr: float, 0 <= p_irr <= 1, a probability of a polygonal chain mask
+
+    min_n_irr: int, min number of segments
+    max_n_irr: int, max number of segments
+    max_l_irr: max length of a segment in polygonal chain
+    max_w_irr: max width of a segment in polygonal chain
+
+    min_n_box: int, min bound for the number of box primitives
+    max_n_box: int, max bound for the number of box primitives
+    min_s_box: int, min length of a box side
+    max_s_box: int, max length of a box side
+    """
+
+    mask = np.zeros((img_h, img_w))
+    uniform = prng.randint
+
+    if np.random.uniform(0, 1) < p_irr:  # generate polygonal chain
+        n = uniform(min_n_irr, max_n_irr)  # sample number of segments
+
+        for _ in range(n):
+            y = uniform(0, img_h)  # sample a starting point
+            x = uniform(0, img_w)
+
+            a = uniform(0, 360)  # sample angle
+            l = uniform(10, max_l_irr)  # sample segment length
+            w = uniform(5, max_w_irr)  # sample a segment width
+
+            # draw segment starting from (x,y) to (x_,y_) using brush of width w
+            x_ = x + l * np.sin(a)
+            y_ = y + l * np.cos(a)
+
+            mask = gen_segment_mask(mask, start=(x, y), end=(x_, y_), brush_width=w)
+            x, y = x_, y_
+    else:  # generate Box masks
+        n = uniform(min_n_box, max_n_box)  # sample number of rectangles
+
+        for _ in range(n):
+            h = uniform(min_s_box, max_s_box)  # sample box shape
+            w = uniform(min_s_box, max_s_box)
+
+            x_0 = uniform(marg, img_w - marg - w)  # sample upper-left coordinates of box
+            y_0 = uniform(marg, img_h - marg - h)
+
+            if np.random.uniform(0, 1) < 0.5:
+                mask = gen_box_mask(mask, masked=(x_0, y_0, w, h))
+            else:
+                r = uniform(0, 60)  # sample radius
+                mask = gen_round_mask(mask, masked=(x_0, y_0, w, h), radius=r)
+    return mask
+
+
+make_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256train"])
+make_narrow_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256narrow"])
+make_512_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train"])
+make_512_lama_mask_large = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train-large"])
+
+
+MASK_MODES = {
+    "256train": make_lama_mask,
+    "256narrow": make_narrow_lama_mask,
+    "512train": make_512_lama_mask,
+    "512train-large": make_512_lama_mask_large
+}
+
+if __name__ == "__main__":
+    import sys
+
+    out = sys.argv[1]
+
+    prng = np.random.RandomState(1)
+    kwargs = settings["256train"]
+    mask = gen_large_mask(prng, 256, 256, **kwargs)
+    mask = (255 * mask).astype(np.uint8)
+    mask = Image.fromarray(mask)
+    mask.save(out)

ldm/data/laion.py

@@ -0,0 +1,537 @@
+import webdataset as wds
+import kornia
+from PIL import Image
+import io
+import os
+import torchvision
+from PIL import Image
+import glob
+import random
+import numpy as np
+import pytorch_lightning as pl
+from tqdm import tqdm
+from omegaconf import OmegaConf
+from einops import rearrange
+import torch
+from webdataset.handlers import warn_and_continue
+
+
+from ldm.util import instantiate_from_config
+from ldm.data.inpainting.synthetic_mask import gen_large_mask, MASK_MODES
+from ldm.data.base import PRNGMixin
+
+
+class DataWithWings(torch.utils.data.IterableDataset):
+    def __init__(self, min_size, transform=None, target_transform=None):
+        self.min_size = min_size
+        self.transform = transform if transform is not None else nn.Identity()
+        self.target_transform = target_transform if target_transform is not None else nn.Identity()
+        self.kv = OnDiskKV(file='/home/ubuntu/laion5B-watermark-safety-ordered', key_format='q', value_format='ee')
+        self.kv_aesthetic = OnDiskKV(file='/home/ubuntu/laion5B-aesthetic-tags-kv', key_format='q', value_format='e')
+        self.pwatermark_threshold = 0.8
+        self.punsafe_threshold = 0.5
+        self.aesthetic_threshold = 5.
+        self.total_samples = 0
+        self.samples = 0
+        location = 'pipe:aws s3 cp --quiet s3://s-datasets/laion5b/laion2B-data/{000000..231349}.tar -'
+
+        self.inner_dataset = wds.DataPipeline(
+            wds.ResampledShards(location),
+            wds.tarfile_to_samples(handler=wds.warn_and_continue),
+            wds.shuffle(1000, handler=wds.warn_and_continue),
+            wds.decode('pilrgb', handler=wds.warn_and_continue),
+            wds.map(self._add_tags, handler=wds.ignore_and_continue),
+            wds.select(self._filter_predicate),
+            wds.map_dict(jpg=self.transform, txt=self.target_transform, punsafe=self._punsafe_to_class, handler=wds.warn_and_continue),
+            wds.to_tuple('jpg', 'txt', 'punsafe', handler=wds.warn_and_continue),
+        )
+
+    @staticmethod
+    def _compute_hash(url, text):
+        if url is None:
+            url = ''
+        if text is None:
+            text = ''
+        total = (url + text).encode('utf-8')
+        return mmh3.hash64(total)[0]
+
+    def _add_tags(self, x):
+        hsh = self._compute_hash(x['json']['url'], x['txt'])
+        pwatermark, punsafe = self.kv[hsh]
+        aesthetic = self.kv_aesthetic[hsh][0]
+        return {**x, 'pwatermark': pwatermark, 'punsafe': punsafe, 'aesthetic': aesthetic}
+
+    def _punsafe_to_class(self, punsafe):
+        return torch.tensor(punsafe >= self.punsafe_threshold).long()
+
+    def _filter_predicate(self, x):
+        try:
+            return x['pwatermark'] < self.pwatermark_threshold and x['aesthetic'] >= self.aesthetic_threshold and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size
+        except:
+            return False
+
+    def __iter__(self):
+        return iter(self.inner_dataset)
+
+
+def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True):
+    """Take a list  of samples (as dictionary) and create a batch, preserving the keys.
+    If `tensors` is True, `ndarray` objects are combined into
+    tensor batches.
+    :param dict samples: list of samples
+    :param bool tensors: whether to turn lists of ndarrays into a single ndarray
+    :returns: single sample consisting of a batch
+    :rtype: dict
+    """
+    keys = set.intersection(*[set(sample.keys()) for sample in samples])
+    batched = {key: [] for key in keys}
+
+    for s in samples:
+        [batched[key].append(s[key]) for key in batched]
+
+    result = {}
+    for key in batched:
+        if isinstance(batched[key][0], (int, float)):
+            if combine_scalars:
+                result[key] = np.array(list(batched[key]))
+        elif isinstance(batched[key][0], torch.Tensor):
+            if combine_tensors:
+                result[key] = torch.stack(list(batched[key]))
+        elif isinstance(batched[key][0], np.ndarray):
+            if combine_tensors:
+                result[key] = np.array(list(batched[key]))
+        else:
+            result[key] = list(batched[key])
+    return result
+
+
+class WebDataModuleFromConfig(pl.LightningDataModule):
+    def __init__(self, tar_base, batch_size, train=None, validation=None,
+                 test=None, num_workers=4, multinode=True, min_size=None,
+                 max_pwatermark=1.0,
+                 **kwargs):
+        super().__init__(self)
+        print(f'Setting tar base to {tar_base}')
+        self.tar_base = tar_base
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.train = train
+        self.validation = validation
+        self.test = test
+        self.multinode = multinode
+        self.min_size = min_size  # filter out very small images
+        self.max_pwatermark = max_pwatermark # filter out watermarked images
+
+    def make_loader(self, dataset_config, train=True):
+        if 'image_transforms' in dataset_config:
+            image_transforms = [instantiate_from_config(tt) for tt in dataset_config.image_transforms]
+        else:
+            image_transforms = []
+
+        image_transforms.extend([torchvision.transforms.ToTensor(),
+                                 torchvision.transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
+        image_transforms = torchvision.transforms.Compose(image_transforms)
+
+        if 'transforms' in dataset_config:
+            transforms_config = OmegaConf.to_container(dataset_config.transforms)
+        else:
+            transforms_config = dict()
+
+        transform_dict = {dkey: load_partial_from_config(transforms_config[dkey])
+                if transforms_config[dkey] != 'identity' else identity
+                for dkey in transforms_config}
+        img_key = dataset_config.get('image_key', 'jpeg')
+        transform_dict.update({img_key: image_transforms})
+
+        if 'postprocess' in dataset_config:
+            postprocess = instantiate_from_config(dataset_config['postprocess'])
+        else:
+            postprocess = None
+
+        shuffle = dataset_config.get('shuffle', 0)
+        shardshuffle = shuffle > 0
+
+        nodesplitter = wds.shardlists.split_by_node if self.multinode else wds.shardlists.single_node_only
+
+        if self.tar_base == "__improvedaesthetic__":
+            print("## Warning, loading the same improved aesthetic dataset "
+                    "for all splits and ignoring shards parameter.")
+            tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
+        else:
+            tars = os.path.join(self.tar_base, dataset_config.shards)
+
+        dset = wds.WebDataset(
+                tars,
+                nodesplitter=nodesplitter,
+                shardshuffle=shardshuffle,
+                handler=wds.warn_and_continue).repeat().shuffle(shuffle)
+        print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.')
+
+        dset = (dset
+                .select(self.filter_keys)
+                .decode('pil', handler=wds.warn_and_continue)
+                .select(self.filter_size)
+                .map_dict(**transform_dict, handler=wds.warn_and_continue)
+                )
+        if postprocess is not None:
+            dset = dset.map(postprocess)
+        dset = (dset
+                .batched(self.batch_size, partial=False,
+                    collation_fn=dict_collation_fn)
+                )
+
+        loader = wds.WebLoader(dset, batch_size=None, shuffle=False,
+                               num_workers=self.num_workers)
+
+        return loader
+
+    def filter_size(self, x):
+        try:
+            valid = True
+            if self.min_size is not None and self.min_size > 1:
+                try:
+                    valid = valid and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size
+                except Exception:
+                    valid = False
+            if self.max_pwatermark is not None and self.max_pwatermark < 1.0:
+                try:
+                    valid = valid and  x['json']['pwatermark'] <= self.max_pwatermark
+                except Exception:
+                    valid = False
+            return valid
+        except Exception:
+            return False
+
+    def filter_keys(self, x):
+        try:
+            return ("jpg" in x) and ("txt" in x)
+        except Exception:
+            return False
+
+    def train_dataloader(self):
+        return self.make_loader(self.train)
+
+    def val_dataloader(self):
+        return self.make_loader(self.validation, train=False)
+
+    def test_dataloader(self):
+        return self.make_loader(self.test, train=False)
+
+
+from ldm.modules.image_degradation import degradation_fn_bsr_light
+import cv2
+
+class AddLR(object):
+    def __init__(self, factor, output_size, initial_size=None, image_key="jpg"):
+        self.factor = factor
+        self.output_size = output_size
+        self.image_key = image_key
+        self.initial_size = initial_size
+
+    def pt2np(self, x):
+        x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy()
+        return x
+
+    def np2pt(self, x):
+        x = torch.from_numpy(x)/127.5-1.0
+        return x
+
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = self.pt2np(sample[self.image_key])
+        if self.initial_size is not None:
+            x = cv2.resize(x, (self.initial_size, self.initial_size), interpolation=2)
+        x = degradation_fn_bsr_light(x, sf=self.factor)['image']
+        x = cv2.resize(x, (self.output_size, self.output_size), interpolation=2)
+        x = self.np2pt(x)
+        sample['lr'] = x
+        return sample
+
+class AddBW(object):
+    def __init__(self, image_key="jpg"):
+        self.image_key = image_key
+
+    def pt2np(self, x):
+        x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy()
+        return x
+
+    def np2pt(self, x):
+        x = torch.from_numpy(x)/127.5-1.0
+        return x
+
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = sample[self.image_key]
+        w = torch.rand(3, device=x.device)
+        w /= w.sum()
+        out = torch.einsum('hwc,c->hw', x, w)
+
+        # Keep as 3ch so we can pass to encoder, also we might want to add hints
+        sample['lr'] = out.unsqueeze(-1).tile(1,1,3)
+        return sample
+
+class AddMask(PRNGMixin):
+    def __init__(self, mode="512train", p_drop=0.):
+        super().__init__()
+        assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"'
+        self.make_mask = MASK_MODES[mode]
+        self.p_drop = p_drop
+
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = sample['jpg']
+        mask = self.make_mask(self.prng, x.shape[0], x.shape[1])
+        if self.prng.choice(2, p=[1 - self.p_drop, self.p_drop]):
+            mask = np.ones_like(mask)
+        mask[mask < 0.5] = 0
+        mask[mask > 0.5] = 1
+        mask = torch.from_numpy(mask[..., None])
+        sample['mask'] = mask
+        sample['masked_image'] = x * (mask < 0.5)
+        return sample
+
+
+class AddEdge(PRNGMixin):
+    def __init__(self, mode="512train", mask_edges=True):
+        super().__init__()
+        assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"'
+        self.make_mask = MASK_MODES[mode]
+        self.n_down_choices = [0]
+        self.sigma_choices = [1, 2]
+        self.mask_edges = mask_edges
+
+    @torch.no_grad()
+    def __call__(self, sample):
+        # sample['jpg'] is tensor hwc in [-1, 1] at this point
+        x = sample['jpg']
+
+        mask = self.make_mask(self.prng, x.shape[0], x.shape[1])
+        mask[mask < 0.5] = 0
+        mask[mask > 0.5] = 1
+        mask = torch.from_numpy(mask[..., None])
+        sample['mask'] = mask
+
+        n_down_idx = self.prng.choice(len(self.n_down_choices))
+        sigma_idx = self.prng.choice(len(self.sigma_choices))
+
+        n_choices = len(self.n_down_choices)*len(self.sigma_choices)
+        raveled_idx = np.ravel_multi_index((n_down_idx, sigma_idx),
+                                           (len(self.n_down_choices), len(self.sigma_choices)))
+        normalized_idx = raveled_idx/max(1, n_choices-1)
+
+        n_down = self.n_down_choices[n_down_idx]
+        sigma = self.sigma_choices[sigma_idx]
+
+        kernel_size = 4*sigma+1
+        kernel_size = (kernel_size, kernel_size)
+        sigma = (sigma, sigma)
+        canny = kornia.filters.Canny(
+                low_threshold=0.1,
+                high_threshold=0.2,
+                kernel_size=kernel_size,
+                sigma=sigma,
+                hysteresis=True,
+                )
+        y = (x+1.0)/2.0 # in 01
+        y = y.unsqueeze(0).permute(0, 3, 1, 2).contiguous()
+
+        # down
+        for i_down in range(n_down):
+            size = min(y.shape[-2], y.shape[-1])//2
+            y = kornia.geometry.transform.resize(y, size, antialias=True)
+
+        # edge
+        _, y = canny(y)
+
+        if n_down > 0:
+            size = x.shape[0], x.shape[1]
+            y = kornia.geometry.transform.resize(y, size, interpolation="nearest")
+
+        y = y.permute(0, 2, 3, 1)[0].expand(-1, -1, 3).contiguous()
+        y = y*2.0-1.0
+
+        if self.mask_edges:
+            sample['masked_image'] = y * (mask < 0.5)
+        else:
+            sample['masked_image'] = y
+            sample['mask'] = torch.zeros_like(sample['mask'])
+
+        # concat normalized idx
+        sample['smoothing_strength'] = torch.ones_like(sample['mask'])*normalized_idx
+
+        return sample
+
+
+def example00():
+    url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/000000.tar -"
+    dataset = wds.WebDataset(url)
+    example = next(iter(dataset))
+    for k in example:
+        print(k, type(example[k]))
+
+    print(example["__key__"])
+    for k in ["json", "txt"]:
+        print(example[k].decode())
+
+    image = Image.open(io.BytesIO(example["jpg"]))
+    outdir = "tmp"
+    os.makedirs(outdir, exist_ok=True)
+    image.save(os.path.join(outdir, example["__key__"] + ".png"))
+
+
+    def load_example(example):
+        return {
+            "key": example["__key__"],
+            "image": Image.open(io.BytesIO(example["jpg"])),
+            "text": example["txt"].decode(),
+        }
+
+
+    for i, example in tqdm(enumerate(dataset)):
+        ex = load_example(example)
+        print(ex["image"].size, ex["text"])
+        if i >= 100:
+            break
+
+
+def example01():
+    # the first laion shards contain ~10k examples each
+    url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/{000000..000002}.tar -"
+
+    batch_size = 3
+    shuffle_buffer = 10000
+    dset = wds.WebDataset(
+            url,
+            nodesplitter=wds.shardlists.split_by_node,
+            shardshuffle=True,
+            )
+    dset = (dset
+            .shuffle(shuffle_buffer, initial=shuffle_buffer)
+            .decode('pil', handler=warn_and_continue)
+            .batched(batch_size, partial=False,
+                collation_fn=dict_collation_fn)
+            )
+
+    num_workers = 2
+    loader = wds.WebLoader(dset, batch_size=None, shuffle=False, num_workers=num_workers)
+
+    batch_sizes = list()
+    keys_per_epoch = list()
+    for epoch in range(5):
+        keys = list()
+        for batch in tqdm(loader):
+            batch_sizes.append(len(batch["__key__"]))
+            keys.append(batch["__key__"])
+
+        for bs in batch_sizes:
+            assert bs==batch_size
+        print(f"{len(batch_sizes)} batches of size {batch_size}.")
+        batch_sizes = list()
+
+        keys_per_epoch.append(keys)
+        for i_batch in [0, 1, -1]:
+            print(f"Batch {i_batch} of epoch {epoch}:")
+            print(keys[i_batch])
+        print("next epoch.")
+
+
+def example02():
+    from omegaconf import OmegaConf
+    from torch.utils.data.distributed import DistributedSampler
+    from torch.utils.data import IterableDataset
+    from torch.utils.data import DataLoader, RandomSampler, Sampler, SequentialSampler
+    from pytorch_lightning.trainer.supporters import CombinedLoader, CycleIterator
+
+    #config = OmegaConf.load("configs/stable-diffusion/txt2img-1p4B-multinode-clip-encoder-high-res-512.yaml")
+    #config = OmegaConf.load("configs/stable-diffusion/txt2img-upscale-clip-encoder-f16-1024.yaml")
+    config = OmegaConf.load("configs/stable-diffusion/txt2img-v2-clip-encoder-improved_aesthetics-256.yaml")
+    datamod = WebDataModuleFromConfig(**config["data"]["params"])
+    dataloader = datamod.train_dataloader()
+
+    for batch in dataloader:
+        print(batch.keys())
+        print(batch["jpg"].shape)
+        break
+
+
+def example03():
+    # improved aesthetics
+    tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
+    dataset = wds.WebDataset(tars)
+
+    def filter_keys(x):
+        try:
+            return ("jpg" in x) and ("txt" in x)
+        except Exception:
+            return False
+
+    def filter_size(x):
+        try:
+            return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
+        except Exception:
+            return False
+
+    def filter_watermark(x):
+        try:
+            return x['json']['pwatermark'] < 0.5
+        except Exception:
+            return False
+
+    dataset = (dataset
+                .select(filter_keys)
+                .decode('pil', handler=wds.warn_and_continue))
+    n_save = 20
+    n_total = 0
+    n_large = 0
+    n_large_nowm = 0
+    for i, example in enumerate(dataset):
+        n_total += 1
+        if filter_size(example):
+            n_large += 1
+            if filter_watermark(example):
+                n_large_nowm += 1
+                if n_large_nowm < n_save+1:
+                    image = example["jpg"]
+                    image.save(os.path.join("tmp", f"{n_large_nowm-1:06}.png"))
+
+        if i%500 == 0:
+            print(i)
+            print(f"Large: {n_large}/{n_total} | {n_large/n_total*100:.2f}%")
+            if n_large > 0:
+                print(f"No Watermark: {n_large_nowm}/{n_large} | {n_large_nowm/n_large*100:.2f}%")
+
+
+
+def example04():
+    # improved aesthetics
+    for i_shard in range(60208)[::-1]:
+        print(i_shard)
+        tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{:06}.tar -".format(i_shard)
+        dataset = wds.WebDataset(tars)
+
+        def filter_keys(x):
+            try:
+                return ("jpg" in x) and ("txt" in x)
+            except Exception:
+                return False
+
+        def filter_size(x):
+            try:
+                return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
+            except Exception:
+                return False
+
+        dataset = (dataset
+                    .select(filter_keys)
+                    .decode('pil', handler=wds.warn_and_continue))
+        try:
+            example = next(iter(dataset))
+        except Exception:
+            print(f"Error @ {i_shard}")
+
+
+if __name__ == "__main__":
+    #example01()
+    #example02()
+    example03()
+    #example04()

ldm/data/lsun.py

@@ -0,0 +1,92 @@
+import os
+import numpy as np
+import PIL
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+class LSUNBase(Dataset):
+    def __init__(self,
+                 txt_file,
+                 data_root,
+                 size=None,
+                 interpolation="bicubic",
+                 flip_p=0.5
+                 ):
+        self.data_paths = txt_file
+        self.data_root = data_root
+        with open(self.data_paths, "r") as f:
+            self.image_paths = f.read().splitlines()
+        self._length = len(self.image_paths)
+        self.labels = {
+            "relative_file_path_": [l for l in self.image_paths],
+            "file_path_": [os.path.join(self.data_root, l)
+                           for l in self.image_paths],
+        }
+
+        self.size = size
+        self.interpolation = {"linear": PIL.Image.LINEAR,
+                              "bilinear": PIL.Image.BILINEAR,
+                              "bicubic": PIL.Image.BICUBIC,
+                              "lanczos": PIL.Image.LANCZOS,
+                              }[interpolation]
+        self.flip = transforms.RandomHorizontalFlip(p=flip_p)
+
+    def __len__(self):
+        return self._length
+
+    def __getitem__(self, i):
+        example = dict((k, self.labels[k][i]) for k in self.labels)
+        image = Image.open(example["file_path_"])
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+
+        # default to score-sde preprocessing
+        img = np.array(image).astype(np.uint8)
+        crop = min(img.shape[0], img.shape[1])
+        h, w, = img.shape[0], img.shape[1]
+        img = img[(h - crop) // 2:(h + crop) // 2,
+              (w - crop) // 2:(w + crop) // 2]
+
+        image = Image.fromarray(img)
+        if self.size is not None:
+            image = image.resize((self.size, self.size), resample=self.interpolation)
+
+        image = self.flip(image)
+        image = np.array(image).astype(np.uint8)
+        example["image"] = (image / 127.5 - 1.0).astype(np.float32)
+        return example
+
+
+class LSUNChurchesTrain(LSUNBase):
+    def __init__(self, **kwargs):
+        super().__init__(txt_file="data/lsun/church_outdoor_train.txt", data_root="data/lsun/churches", **kwargs)
+
+
+class LSUNChurchesValidation(LSUNBase):
+    def __init__(self, flip_p=0., **kwargs):
+        super().__init__(txt_file="data/lsun/church_outdoor_val.txt", data_root="data/lsun/churches",
+                         flip_p=flip_p, **kwargs)
+
+
+class LSUNBedroomsTrain(LSUNBase):
+    def __init__(self, **kwargs):
+        super().__init__(txt_file="data/lsun/bedrooms_train.txt", data_root="data/lsun/bedrooms", **kwargs)
+
+
+class LSUNBedroomsValidation(LSUNBase):
+    def __init__(self, flip_p=0.0, **kwargs):
+        super().__init__(txt_file="data/lsun/bedrooms_val.txt", data_root="data/lsun/bedrooms",
+                         flip_p=flip_p, **kwargs)
+
+
+class LSUNCatsTrain(LSUNBase):
+    def __init__(self, **kwargs):
+        super().__init__(txt_file="data/lsun/cat_train.txt", data_root="data/lsun/cats", **kwargs)
+
+
+class LSUNCatsValidation(LSUNBase):
+    def __init__(self, flip_p=0., **kwargs):
+        super().__init__(txt_file="data/lsun/cat_val.txt", data_root="data/lsun/cats",
+                         flip_p=flip_p, **kwargs)

ldm/data/nerf_like.py

@@ -0,0 +1,165 @@
+from torch.utils.data import Dataset
+import os
+import json
+import numpy as np
+import torch
+import imageio
+import math
+import cv2
+from torchvision import transforms
+
+def cartesian_to_spherical(xyz):
+    ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
+    xy = xyz[:,0]**2 + xyz[:,1]**2
+    z = np.sqrt(xy + xyz[:,2]**2)
+    theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
+    #ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
+    azimuth = np.arctan2(xyz[:,1], xyz[:,0])
+    return np.array([theta, azimuth, z])
+
+
+def get_T(T_target, T_cond):
+    theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :])
+    theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
+    
+    d_theta = theta_target - theta_cond
+    d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
+    d_z = z_target - z_cond
+    
+    d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
+    return d_T
+
+def get_spherical(T_target, T_cond):
+    theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :])
+    theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :])
+    
+    d_theta = theta_target - theta_cond
+    d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
+    d_z = z_target - z_cond
+    
+    d_T = torch.tensor([math.degrees(d_theta.item()), math.degrees(d_azimuth.item()), d_z.item()])
+    return d_T
+
+class RTMV(Dataset):
+    def __init__(self, root_dir='datasets/RTMV/google_scanned',\
+                 first_K=64, resolution=256, load_target=False):
+        self.root_dir = root_dir
+        self.scene_list = sorted(next(os.walk(root_dir))[1])
+        self.resolution = resolution
+        self.first_K = first_K
+        self.load_target = load_target
+
+    def __len__(self):
+        return len(self.scene_list)
+
+    def __getitem__(self, idx):
+        scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
+        with open(os.path.join(scene_dir, 'transforms.json'), "r") as f:
+            meta = json.load(f)
+        imgs = []
+        poses = []
+        for i_img in range(self.first_K):
+            meta_img = meta['frames'][i_img]
+
+            if i_img == 0 or self.load_target:
+                img_path = os.path.join(scene_dir, meta_img['file_path'])
+                img = imageio.imread(img_path)
+                img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
+                imgs.append(img)
+            
+            c2w = meta_img['transform_matrix']
+            poses.append(c2w)
+            
+        imgs = (np.array(imgs) / 255.).astype(np.float32)  # (RGBA) imgs
+        imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
+        imgs = imgs * 2 - 1. # convert to stable diffusion range
+        poses = torch.tensor(np.array(poses).astype(np.float32))
+        return imgs, poses
+                
+    def blend_rgba(self, img):
+        img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:])  # blend A to RGB
+        return img
+            
+
+class GSO(Dataset):
+    def __init__(self, root_dir='datasets/GoogleScannedObjects',\
+                 split='val', first_K=5, resolution=256, load_target=False, name='render_mvs'):
+        self.root_dir = root_dir
+        with open(os.path.join(root_dir, '%s.json' % split), "r") as f:
+            self.scene_list = json.load(f)
+        self.resolution = resolution
+        self.first_K = first_K
+        self.load_target = load_target
+        self.name = name
+
+    def __len__(self):
+        return len(self.scene_list)
+
+    def __getitem__(self, idx):
+        scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
+        with open(os.path.join(scene_dir, 'transforms_%s.json' % self.name), "r") as f:
+            meta = json.load(f)
+        imgs = []
+        poses = []
+        for i_img in range(self.first_K):
+            meta_img = meta['frames'][i_img]
+
+            if i_img == 0 or self.load_target:
+                img_path = os.path.join(scene_dir, meta_img['file_path'])
+                img = imageio.imread(img_path)
+                img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
+                imgs.append(img)
+            
+            c2w = meta_img['transform_matrix']
+            poses.append(c2w)
+            
+        imgs = (np.array(imgs) / 255.).astype(np.float32)  # (RGBA) imgs
+        mask = imgs[:, :, :, -1]
+        imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
+        imgs = imgs * 2 - 1. # convert to stable diffusion range
+        poses = torch.tensor(np.array(poses).astype(np.float32))
+        return imgs, poses
+                
+    def blend_rgba(self, img):
+        img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:])  # blend A to RGB
+        return img
+             
+class WILD(Dataset):
+    def __init__(self, root_dir='data/nerf_wild',\
+                 first_K=33, resolution=256, load_target=False):
+        self.root_dir = root_dir
+        self.scene_list = sorted(next(os.walk(root_dir))[1])
+        self.resolution = resolution
+        self.first_K = first_K
+        self.load_target = load_target
+
+    def __len__(self):
+        return len(self.scene_list)
+
+    def __getitem__(self, idx):
+        scene_dir = os.path.join(self.root_dir, self.scene_list[idx])
+        with open(os.path.join(scene_dir, 'transforms_train.json'), "r") as f:
+            meta = json.load(f)
+        imgs = []
+        poses = []
+        for i_img in range(self.first_K):
+            meta_img = meta['frames'][i_img]
+
+            if i_img == 0 or self.load_target:
+                img_path = os.path.join(scene_dir, meta_img['file_path'])
+                img = imageio.imread(img_path + '.png')
+                img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR)
+                imgs.append(img)
+            
+            c2w = meta_img['transform_matrix']
+            poses.append(c2w)
+            
+        imgs = (np.array(imgs) / 255.).astype(np.float32)  # (RGBA) imgs
+        imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2)
+        imgs = imgs * 2 - 1. # convert to stable diffusion range
+        poses = torch.tensor(np.array(poses).astype(np.float32))
+        return imgs, poses
+                
+    def blend_rgba(self, img):
+        img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:])  # blend A to RGB
+        return img

ldm/data/simple.py

@@ -0,0 +1,526 @@
+from typing import Dict
+import webdataset as wds
+import numpy as np
+from omegaconf import DictConfig, ListConfig
+import torch
+from torch.utils.data import Dataset
+from pathlib import Path
+import json
+from PIL import Image
+from torchvision import transforms
+import torchvision
+from einops import rearrange
+from ldm.util import instantiate_from_config
+from datasets import load_dataset
+import pytorch_lightning as pl
+import copy
+import csv
+import cv2
+import random
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+import json
+import os, sys
+import webdataset as wds
+import math
+from torch.utils.data.distributed import DistributedSampler
+
+# Some hacky things to make experimentation easier
+def make_transform_multi_folder_data(paths, caption_files=None, **kwargs):
+    ds = make_multi_folder_data(paths, caption_files, **kwargs)
+    return TransformDataset(ds)
+
+def make_nfp_data(base_path):
+    dirs = list(Path(base_path).glob("*/"))
+    print(f"Found {len(dirs)} folders")
+    print(dirs)
+    tforms = [transforms.Resize(512), transforms.CenterCrop(512)]
+    datasets = [NfpDataset(x, image_transforms=copy.copy(tforms), default_caption="A view from a train window") for x in dirs]
+    return torch.utils.data.ConcatDataset(datasets)
+
+
+class VideoDataset(Dataset):
+    def __init__(self, root_dir, image_transforms, caption_file, offset=8, n=2):
+        self.root_dir = Path(root_dir)
+        self.caption_file = caption_file
+        self.n = n
+        ext = "mp4"
+        self.paths = sorted(list(self.root_dir.rglob(f"*.{ext}")))
+        self.offset = offset
+
+        if isinstance(image_transforms, ListConfig):
+            image_transforms = [instantiate_from_config(tt) for tt in image_transforms]
+        image_transforms.extend([transforms.ToTensor(),
+                                 transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
+        image_transforms = transforms.Compose(image_transforms)
+        self.tform = image_transforms
+        with open(self.caption_file) as f:
+            reader = csv.reader(f)
+            rows = [row for row in reader]
+        self.captions = dict(rows)
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, index):
+        for i in range(10):
+            try:
+                return self._load_sample(index)
+            except Exception:
+                # Not really good enough but...
+                print("uh oh")
+
+    def _load_sample(self, index):
+        n = self.n
+        filename = self.paths[index]
+        min_frame = 2*self.offset + 2
+        vid = cv2.VideoCapture(str(filename))
+        max_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
+        curr_frame_n = random.randint(min_frame, max_frames)
+        vid.set(cv2.CAP_PROP_POS_FRAMES,curr_frame_n)
+        _, curr_frame = vid.read()
+
+        prev_frames = []
+        for i in range(n):
+            prev_frame_n = curr_frame_n - (i+1)*self.offset
+            vid.set(cv2.CAP_PROP_POS_FRAMES,prev_frame_n)
+            _, prev_frame = vid.read()
+            prev_frame = self.tform(Image.fromarray(prev_frame[...,::-1]))
+            prev_frames.append(prev_frame)
+
+        vid.release()
+        caption = self.captions[filename.name]
+        data = {
+            "image": self.tform(Image.fromarray(curr_frame[...,::-1])),
+            "prev": torch.cat(prev_frames, dim=-1),
+            "txt": caption
+        }
+        return data
+
+# end hacky things
+
+
+def make_tranforms(image_transforms):
+    # if isinstance(image_transforms, ListConfig):
+    #     image_transforms = [instantiate_from_config(tt) for tt in image_transforms]
+    image_transforms = []
+    image_transforms.extend([transforms.ToTensor(),
+                                transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
+    image_transforms = transforms.Compose(image_transforms)
+    return image_transforms
+
+
+def make_multi_folder_data(paths, caption_files=None, **kwargs):
+    """Make a concat dataset from multiple folders
+    Don't suport captions yet
+
+    If paths is a list, that's ok, if it's a Dict interpret it as:
+    k=folder v=n_times to repeat that
+    """
+    list_of_paths = []
+    if isinstance(paths, (Dict, DictConfig)):
+        assert caption_files is None, \
+            "Caption files not yet supported for repeats"
+        for folder_path, repeats in paths.items():
+            list_of_paths.extend([folder_path]*repeats)
+        paths = list_of_paths
+
+    if caption_files is not None:
+        datasets = [FolderData(p, caption_file=c, **kwargs) for (p, c) in zip(paths, caption_files)]
+    else:
+        datasets = [FolderData(p, **kwargs) for p in paths]
+    return torch.utils.data.ConcatDataset(datasets)
+
+
+
+class NfpDataset(Dataset):
+    def __init__(self,
+        root_dir,
+        image_transforms=[],
+        ext="jpg",
+        default_caption="",
+        ) -> None:
+        """assume sequential frames and a deterministic transform"""
+
+        self.root_dir = Path(root_dir)
+        self.default_caption = default_caption
+
+        self.paths = sorted(list(self.root_dir.rglob(f"*.{ext}")))
+        self.tform = make_tranforms(image_transforms)
+
+    def __len__(self):
+        return len(self.paths) - 1
+
+
+    def __getitem__(self, index):
+        prev = self.paths[index]
+        curr = self.paths[index+1]
+        data = {}
+        data["image"] = self._load_im(curr)
+        data["prev"] = self._load_im(prev)
+        data["txt"] = self.default_caption
+        return data
+
+    def _load_im(self, filename):
+        im = Image.open(filename).convert("RGB")
+        return self.tform(im)
+
+class ObjaverseDataModuleFromConfig(pl.LightningDataModule):
+    def __init__(self, root_dir, batch_size, total_view, train=None, validation=None,
+                 test=None, num_workers=4, **kwargs):
+        super().__init__(self)
+        self.root_dir = root_dir
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.total_view = total_view
+
+        if train is not None:
+            dataset_config = train
+        if validation is not None:
+            dataset_config = validation
+
+        if 'image_transforms' in dataset_config:
+            image_transforms = [torchvision.transforms.Resize(dataset_config.image_transforms.size)]
+        else:
+            image_transforms = []
+        image_transforms.extend([transforms.ToTensor(),
+                                transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))])
+        self.image_transforms = torchvision.transforms.Compose(image_transforms)
+
+
+    def train_dataloader(self):
+        dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=False, \
+                                image_transforms=self.image_transforms)
+        sampler = DistributedSampler(dataset)
+        return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, sampler=sampler)
+
+    def val_dataloader(self):
+        dataset = ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=True, \
+                                image_transforms=self.image_transforms)
+        sampler = DistributedSampler(dataset)
+        return wds.WebLoader(dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
+    
+    def test_dataloader(self):
+        return wds.WebLoader(ObjaverseData(root_dir=self.root_dir, total_view=self.total_view, validation=self.validation),\
+                          batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
+
+
+class ObjaverseData(Dataset):
+    def __init__(self,
+        root_dir='.objaverse/hf-objaverse-v1/views',
+        image_transforms=[],
+        ext="png",
+        default_trans=torch.zeros(3),
+        postprocess=None,
+        return_paths=False,
+        total_view=4,
+        validation=False
+        ) -> None:
+        """Create a dataset from a folder of images.
+        If you pass in a root directory it will be searched for images
+        ending in ext (ext can be a list)
+        """
+        self.root_dir = Path(root_dir)
+        self.default_trans = default_trans
+        self.return_paths = return_paths
+        if isinstance(postprocess, DictConfig):
+            postprocess = instantiate_from_config(postprocess)
+        self.postprocess = postprocess
+        self.total_view = total_view
+
+        if not isinstance(ext, (tuple, list, ListConfig)):
+            ext = [ext]
+
+        with open(os.path.join(root_dir, 'valid_paths.json')) as f:
+            self.paths = json.load(f)
+            
+        total_objects = len(self.paths)
+        if validation:
+            self.paths = self.paths[math.floor(total_objects / 100. * 99.):] # used last 1% as validation
+        else:
+            self.paths = self.paths[:math.floor(total_objects / 100. * 99.)] # used first 99% as training
+        print('============= length of dataset %d =============' % len(self.paths))
+        self.tform = image_transforms
+
+    def __len__(self):
+        return len(self.paths)
+        
+    def cartesian_to_spherical(self, xyz):
+        ptsnew = np.hstack((xyz, np.zeros(xyz.shape)))
+        xy = xyz[:,0]**2 + xyz[:,1]**2
+        z = np.sqrt(xy + xyz[:,2]**2)
+        theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down
+        #ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up
+        azimuth = np.arctan2(xyz[:,1], xyz[:,0])
+        return np.array([theta, azimuth, z])
+
+    def get_T(self, target_RT, cond_RT):
+        R, T = target_RT[:3, :3], target_RT[:, -1]
+        T_target = -R.T @ T
+
+        R, T = cond_RT[:3, :3], cond_RT[:, -1]
+        T_cond = -R.T @ T
+
+        theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
+        theta_target, azimuth_target, z_target = self.cartesian_to_spherical(T_target[None, :])
+        
+        d_theta = theta_target - theta_cond
+        d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi)
+        d_z = z_target - z_cond
+        
+        d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()])
+        return d_T
+
+    def load_im(self, path, color):
+        '''
+        replace background pixel with random color in rendering
+        '''
+        try:
+            img = plt.imread(path)
+        except:
+            print(path)
+            sys.exit()
+        img[img[:, :, -1] == 0.] = color
+        img = Image.fromarray(np.uint8(img[:, :, :3] * 255.))
+        return img
+
+    def __getitem__(self, index):
+
+        data = {}
+        if self.paths[index][-2:] == '_1': # dirty fix for rendering dataset twice
+            total_view = 8
+        else:
+            total_view = 4
+        index_target, index_cond = random.sample(range(total_view), 2) # without replacement
+        filename = os.path.join(self.root_dir, self.paths[index])
+
+        # print(self.paths[index])
+
+        if self.return_paths:
+            data["path"] = str(filename)
+        
+        color = [1., 1., 1., 1.]
+
+        try:
+            target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
+            cond_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_cond), color))
+            target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
+            cond_RT = np.load(os.path.join(filename, '%03d.npy' % index_cond))
+        except:
+            # very hacky solution, sorry about this
+            filename = os.path.join(self.root_dir, '692db5f2d3a04bb286cb977a7dba903e_1') # this one we know is valid
+            target_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_target), color))
+            cond_im = self.process_im(self.load_im(os.path.join(filename, '%03d.png' % index_cond), color))
+            target_RT = np.load(os.path.join(filename, '%03d.npy' % index_target))
+            cond_RT = np.load(os.path.join(filename, '%03d.npy' % index_cond))
+            target_im = torch.zeros_like(target_im)
+            cond_im = torch.zeros_like(cond_im)
+
+        data["image_target"] = target_im
+        data["image_cond"] = cond_im
+        data["T"] = self.get_T(target_RT, cond_RT)
+
+        if self.postprocess is not None:
+            data = self.postprocess(data)
+
+        return data
+
+    def process_im(self, im):
+        im = im.convert("RGB")
+        return self.tform(im)
+
+class FolderData(Dataset):
+    def __init__(self,
+        root_dir,
+        caption_file=None,
+        image_transforms=[],
+        ext="jpg",
+        default_caption="",
+        postprocess=None,
+        return_paths=False,
+        ) -> None:
+        """Create a dataset from a folder of images.
+        If you pass in a root directory it will be searched for images
+        ending in ext (ext can be a list)
+        """
+        self.root_dir = Path(root_dir)
+        self.default_caption = default_caption
+        self.return_paths = return_paths
+        if isinstance(postprocess, DictConfig):
+            postprocess = instantiate_from_config(postprocess)
+        self.postprocess = postprocess
+        if caption_file is not None:
+            with open(caption_file, "rt") as f:
+                ext = Path(caption_file).suffix.lower()
+                if ext == ".json":
+                    captions = json.load(f)
+                elif ext == ".jsonl":
+                    lines = f.readlines()
+                    lines = [json.loads(x) for x in lines]
+                    captions = {x["file_name"]: x["text"].strip("\n") for x in lines}
+                else:
+                    raise ValueError(f"Unrecognised format: {ext}")
+            self.captions = captions
+        else:
+            self.captions = None
+
+        if not isinstance(ext, (tuple, list, ListConfig)):
+            ext = [ext]
+
+        # Only used if there is no caption file
+        self.paths = []
+        for e in ext:
+            self.paths.extend(sorted(list(self.root_dir.rglob(f"*.{e}"))))
+        self.tform = make_tranforms(image_transforms)
+
+    def __len__(self):
+        if self.captions is not None:
+            return len(self.captions.keys())
+        else:
+            return len(self.paths)
+
+    def __getitem__(self, index):
+        data = {}
+        if self.captions is not None:
+            chosen = list(self.captions.keys())[index]
+            caption = self.captions.get(chosen, None)
+            if caption is None:
+                caption = self.default_caption
+            filename = self.root_dir/chosen
+        else:
+            filename = self.paths[index]
+
+        if self.return_paths:
+            data["path"] = str(filename)
+
+        im = Image.open(filename).convert("RGB")
+        im = self.process_im(im)
+        data["image"] = im
+
+        if self.captions is not None:
+            data["txt"] = caption
+        else:
+            data["txt"] = self.default_caption
+
+        if self.postprocess is not None:
+            data = self.postprocess(data)
+
+        return data
+
+    def process_im(self, im):
+        im = im.convert("RGB")
+        return self.tform(im)
+import random
+
+class TransformDataset():
+    def __init__(self, ds, extra_label="sksbspic"):
+        self.ds = ds
+        self.extra_label = extra_label
+        self.transforms = {
+            "align": transforms.Resize(768),
+            "centerzoom": transforms.CenterCrop(768),
+            "randzoom": transforms.RandomCrop(768),
+        }
+
+
+    def __getitem__(self, index):
+        data = self.ds[index]
+
+        im = data['image']
+        im = im.permute(2,0,1)
+        # In case data is smaller than expected
+        im = transforms.Resize(1024)(im)
+
+        tform_name = random.choice(list(self.transforms.keys()))
+        im = self.transforms[tform_name](im)
+
+        im = im.permute(1,2,0)
+
+        data['image'] = im
+        data['txt'] = data['txt'] + f" {self.extra_label} {tform_name}"
+
+        return data
+
+    def __len__(self):
+        return len(self.ds)
+
+def hf_dataset(
+    name,
+    image_transforms=[],
+    image_column="image",
+    text_column="text",
+    split='train',
+    image_key='image',
+    caption_key='txt',
+    ):
+    """Make huggingface dataset with appropriate list of transforms applied
+    """
+    ds = load_dataset(name, split=split)
+    tform = make_tranforms(image_transforms)
+
+    assert image_column in ds.column_names, f"Didn't find column {image_column} in {ds.column_names}"
+    assert text_column in ds.column_names, f"Didn't find column {text_column} in {ds.column_names}"
+
+    def pre_process(examples):
+        processed = {}
+        processed[image_key] = [tform(im) for im in examples[image_column]]
+        processed[caption_key] = examples[text_column]
+        return processed
+
+    ds.set_transform(pre_process)
+    return ds
+
+class TextOnly(Dataset):
+    def __init__(self, captions, output_size, image_key="image", caption_key="txt", n_gpus=1):
+        """Returns only captions with dummy images"""
+        self.output_size = output_size
+        self.image_key = image_key
+        self.caption_key = caption_key
+        if isinstance(captions, Path):
+            self.captions = self._load_caption_file(captions)
+        else:
+            self.captions = captions
+
+        if n_gpus > 1:
+            # hack to make sure that all the captions appear on each gpu
+            repeated = [n_gpus*[x] for x in self.captions]
+            self.captions = []
+            [self.captions.extend(x) for x in repeated]
+
+    def __len__(self):
+        return len(self.captions)
+
+    def __getitem__(self, index):
+        dummy_im = torch.zeros(3, self.output_size, self.output_size)
+        dummy_im = rearrange(dummy_im * 2. - 1., 'c h w -> h w c')
+        return {self.image_key: dummy_im, self.caption_key: self.captions[index]}
+
+    def _load_caption_file(self, filename):
+        with open(filename, 'rt') as f:
+            captions = f.readlines()
+        return [x.strip('\n') for x in captions]
+
+
+
+import random
+import json
+class IdRetreivalDataset(FolderData):
+    def __init__(self, ret_file, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        with open(ret_file, "rt") as f:
+            self.ret = json.load(f)
+
+    def __getitem__(self, index):
+        data = super().__getitem__(index)
+        key = self.paths[index].name
+        matches = self.ret[key]
+        if len(matches) > 0:
+            retreived = random.choice(matches)
+        else:
+            retreived = key
+        filename = self.root_dir/retreived
+        im = Image.open(filename).convert("RGB")
+        im = self.process_im(im)
+        # data["match"] = im
+        data["match"] = torch.cat((data["image"], im), dim=-1)
+        return data

ldm/extras.py

@@ -0,0 +1,77 @@
+from pathlib import Path
+from omegaconf import OmegaConf
+import torch
+from ldm.util import instantiate_from_config
+import logging
+from contextlib import contextmanager
+
+from contextlib import contextmanager
+import logging
+
+@contextmanager
+def all_logging_disabled(highest_level=logging.CRITICAL):
+    """
+    A context manager that will prevent any logging messages
+    triggered during the body from being processed.
+
+    :param highest_level: the maximum logging level in use.
+      This would only need to be changed if a custom level greater than CRITICAL
+      is defined.
+
+    https://gist.github.com/simon-weber/7853144
+    """
+    # two kind-of hacks here:
+    #    * can't get the highest logging level in effect => delegate to the user
+    #    * can't get the current module-level override => use an undocumented
+    #       (but non-private!) interface
+
+    previous_level = logging.root.manager.disable
+
+    logging.disable(highest_level)
+
+    try:
+        yield
+    finally:
+        logging.disable(previous_level)
+
+def load_training_dir(train_dir, device, epoch="last"):
+    """Load a checkpoint and config from training directory"""
+    train_dir = Path(train_dir)
+    ckpt = list(train_dir.rglob(f"*{epoch}.ckpt"))
+    assert len(ckpt) == 1, f"found {len(ckpt)} matching ckpt files"
+    config = list(train_dir.rglob(f"*-project.yaml"))
+    assert len(ckpt) > 0, f"didn't find any config in {train_dir}"
+    if len(config) > 1:
+        print(f"found {len(config)} matching config files")
+        config = sorted(config)[-1]
+        print(f"selecting {config}")
+    else:
+        config = config[0]
+
+
+    config = OmegaConf.load(config)
+    return load_model_from_config(config, ckpt[0], device)
+
+def load_model_from_config(config, ckpt, device="cpu", verbose=False):
+    """Loads a model from config and a ckpt
+    if config is a path will use omegaconf to load
+    """
+    if isinstance(config, (str, Path)):
+        config = OmegaConf.load(config)
+
+    with all_logging_disabled():
+        print(f"Loading model from {ckpt}")
+        pl_sd = torch.load(ckpt, map_location="cpu")
+        global_step = pl_sd["global_step"]
+        sd = pl_sd["state_dict"]
+        model = instantiate_from_config(config.model)
+        m, u = model.load_state_dict(sd, strict=False)
+        if len(m) > 0 and verbose:
+            print("missing keys:")
+            print(m)
+        if len(u) > 0 and verbose:
+            print("unexpected keys:")
+        model.to(device)
+        model.eval()
+        model.cond_stage_model.device = device
+        return model

ldm/guidance.py

@@ -0,0 +1,96 @@
+from typing import List, Tuple
+from scipy import interpolate
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from IPython.display import clear_output
+import abc
+
+
+class GuideModel(torch.nn.Module, abc.ABC):
+    def __init__(self) -> None:
+        super().__init__()
+
+    @abc.abstractmethod
+    def preprocess(self, x_img):
+        pass
+
+    @abc.abstractmethod
+    def compute_loss(self, inp):
+        pass
+
+
+class Guider(torch.nn.Module):
+    def __init__(self, sampler, guide_model, scale=1.0, verbose=False):
+        """Apply classifier guidance
+
+        Specify a guidance scale as either a scalar
+        Or a schedule as a list of tuples t = 0->1 and scale, e.g.
+        [(0, 10), (0.5, 20), (1, 50)]
+        """
+        super().__init__()
+        self.sampler = sampler
+        self.index = 0
+        self.show = verbose
+        self.guide_model = guide_model
+        self.history = []
+
+        if isinstance(scale, (Tuple, List)):
+            times = np.array([x[0] for x in scale])
+            values = np.array([x[1] for x in scale])
+            self.scale_schedule = {"times": times, "values": values}
+        else:
+            self.scale_schedule = float(scale)
+
+        self.ddim_timesteps = sampler.ddim_timesteps
+        self.ddpm_num_timesteps = sampler.ddpm_num_timesteps
+
+
+    def get_scales(self):
+        if isinstance(self.scale_schedule, float):
+            return len(self.ddim_timesteps)*[self.scale_schedule]
+
+        interpolater = interpolate.interp1d(self.scale_schedule["times"], self.scale_schedule["values"])
+        fractional_steps = np.array(self.ddim_timesteps)/self.ddpm_num_timesteps
+        return interpolater(fractional_steps)
+
+    def modify_score(self, model, e_t, x, t, c):
+
+        # TODO look up index by t
+        scale = self.get_scales()[self.index]
+
+        if (scale == 0):
+            return e_t
+
+        sqrt_1ma = self.sampler.ddim_sqrt_one_minus_alphas[self.index].to(x.device)
+        with torch.enable_grad():
+            x_in = x.detach().requires_grad_(True)
+            pred_x0 = model.predict_start_from_noise(x_in, t=t, noise=e_t)
+            x_img = model.first_stage_model.decode((1/0.18215)*pred_x0)
+
+            inp = self.guide_model.preprocess(x_img)
+            loss = self.guide_model.compute_loss(inp)
+            grads = torch.autograd.grad(loss.sum(), x_in)[0]
+            correction = grads * scale
+
+            if self.show:
+                clear_output(wait=True)
+                print(loss.item(), scale, correction.abs().max().item(), e_t.abs().max().item())
+                self.history.append([loss.item(), scale, correction.min().item(), correction.max().item()])
+                plt.imshow((inp[0].detach().permute(1,2,0).clamp(-1,1).cpu()+1)/2)
+                plt.axis('off')
+                plt.show()
+                plt.imshow(correction[0][0].detach().cpu())
+                plt.axis('off')
+                plt.show()
+
+
+        e_t_mod = e_t - sqrt_1ma*correction
+        if self.show:
+            fig, axs = plt.subplots(1, 3)
+            axs[0].imshow(e_t[0][0].detach().cpu(), vmin=-2, vmax=+2)
+            axs[1].imshow(e_t_mod[0][0].detach().cpu(), vmin=-2, vmax=+2)
+            axs[2].imshow(correction[0][0].detach().cpu(), vmin=-2, vmax=+2)
+            plt.show()
+        self.index += 1
+        return e_t_mod

ldm/lr_scheduler.py

@@ -0,0 +1,98 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+    """
+    note: use with a base_lr of 1.0
+    """
+    def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n, **kwargs):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+                    1 + np.cos(t * np.pi))
+            self.last_lr = lr
+            return lr
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n,**kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+    """
+    supports repeated iterations, configurable via lists
+    note: use with a base_lr of 1.0.
+    """
+    def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
+        assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
+        self.lr_warm_up_steps = warm_up_steps
+        self.f_start = f_start
+        self.f_min = f_min
+        self.f_max = f_max
+        self.cycle_lengths = cycle_lengths
+        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+        self.last_f = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def find_in_interval(self, n):
+        interval = 0
+        for cl in self.cum_cycles[1:]:
+            if n <= cl:
+                return interval
+            interval += 1
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                                                       f"current cycle {cycle}")
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
+            t = min(t, 1.0)
+            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+                    1 + np.cos(t * np.pi))
+            self.last_f = f
+            return f
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                                                       f"current cycle {cycle}")
+
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle])
+            self.last_f = f
+            return f
+