Implementation of BakedSDF

[pablovela5620]

https://bakedsdf.github.io/, most of the components are already there (VolSDF | MipNerf 360 | Marching cubes). The largest missing part is the appearance model that uses spherical Gaussians. Would love to see this brought into SDFStudio

[satyajitghana]

+1 would love to see this implementation. The results of BasedSDF are really amazing.

[niujinshuchong]

Hi, BakedSDF is implemented. After downloading the mipnerf360 dataset with ns-download-data mipnerf360, it can be run with

ns-train bakedsdf --vis wandb --experiment-name garden-bakedsdf mipnerf360-data --data data/nerfstudio-data-mipnerf360/garden

or

ns-train bakedsdf-mlp --vis wandb --experiment-name garden-bakedsdf-mlp mipnerf360-data --data data/nerfstudio-data-mipnerf360/garden

Here are some results:

bakedsdf:

bakedsdf-mlp

Tuning some hyper parameters would likely give better results.

[pablovela5620]

Wow, you got that implemented extremely fast! Excellent, I'll give it a test!

[reynoldscem]

To what extent is bakedSDF implemented? Do you support the full export stage?

[reynoldscem]

Looking at the code, unless I'm missing something, the quantization-aware training is not included either. Is this the case or am I missing something?

[reynoldscem]

@niujinshuchong can you advise & @pablovela5620 did you try it out?

[niujinshuchong]

@reynoldscem Training of bakedsdf and mesh extraction are supported. You could extract the mesh after training with

ns-extract-mesh --load-config CONFIG --output-path MESHFILE.ply --bounding-box-min -2.0 -2.0 -2.0 --bounding-box-max 2.0 2.0 2.0 --resolution 512 --marching_cube_threshold 0.001 --create_visibility_mask True

The baking part is missing and I don't think it will be added soon. Contributions are welcome.

[tianxiaguixin002]

@niujinshuchong

Thanks for your great work, I try to train bakedSDF in garden scene of nerfstudio-data-mipnerf360 data with following cmd:

ns-train bakedsdf --vis viewer --experiment-name garden-bakedsdf \ --data /home/data/dataset/neus_data/nerfstudio-data-mipnerf360/garden \ --pipeline.datamanager.camera-res-scale-factor 0.25 \ --pipeline.datamanager.train-num-rays-per-batch 5000 mipnerf360-data

then extrash mesh:

ns-extract-mesh --load-config outputs/garden-bakedsdf/bakedsdf/2023-03-07_110420/config.yml \ --output-path meshes/bakedsdf-garden-512.ply --bounding-box-min -2.0 -2.0 -2.0 --bounding-box-max 2.0 2.0 2.0 --resolution 512 --marching_cube_threshold 0.001 --create_visibility_mask True

The mesh like this:

The surrounding mesh is incomplete, and I adjust the bounding-box and resolution, the result doesn't improved, please help me where is the problem. besides, can you show the mesh result of mipnerf360 data? like garden scene? thank you very much!

[niujinshuchong]

@tianxiaguixin002 I think this might because we compute the visibility grid with testing images. I push a update for that. Here is the results using all training images to compute the visibility grids.

bakedsdf

bakedsdf-mlp

MLP's mesh is more smooth.

[tianxiaguixin002]

-experiment-name garden-b

Thank you very much for your reply, I will try again in a few days.

[tlightsky]

how can I convert the result to a normal glb file, which can be rendered in threejs?

[reynoldscem]

@tlightsky as far as I can tell, that part isn't included

[niujinshuchong]

@tlightsky @reynoldscem Yes, It's not included.

[tianxiaguixin002]

@niujinshuchong I use the latest code and have trained bakedsdf-mlp in mipnerf360-garden scene with the offical cmds from sdfstudio-examples.md. The mesh quality is good under normal observation position. like following picture:

But then looking down from above and adjust the zoom ratio, some mesh vertexs are radial, these points maybe in SDF grid [-2, -1] and [1,2].

So do you know what causes this phenomenon？ is it a correct result ? I think maybe due to too few view for the far region, the parallax is difficult to compute. Looking forward for your reply, thank you.

[satyajitghana]

@tianxiaguixin002 seems okay, you can download the official bakedsdf .glb file from https://dl.dropboxusercontent.com/s/7srlfofbiryehn6/gardenvase.glb and compare with this. I think they are very similar. the far region distortion you see is common, a bbox crop is required to clear those.

[gaoalexander]

@niujinshuchong What exactly is the difference between bakedsdf and bakedsdf-mlp?

[niujinshuchong]

@gaoalexander bakedsdf uses multi-res. hash grids and a small mlp as scene representation while bakedsdf-mlp uses a big MLP with 8 layers and 1024 hidden dim. They are all configurable from the command line.

[CanCanZeng]

Hi, I use ns-train bakedsdf --experiment-name garden-bakedsdf mipnerf360-data --data data/nerfstudio-data-mipnerf360/garden to test bakedsdf on my ubuntu20.04 machine with nvidia 4080 GPU, but the result is rather poor, bellow is my conda env:

# packages in environment at ~/miniconda3/envs/nerfacc:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                        main  
_openmp_mutex             5.1                       1_gnu  
absl-py                   1.4.0                    pypi_0    pypi
addict                    2.4.0                    pypi_0    pypi
aiohttp                   3.8.4                    pypi_0    pypi
aiosignal                 1.3.1                    pypi_0    pypi
anyio                     3.6.2                    pypi_0    pypi
appdirs                   1.4.4                    pypi_0    pypi
argon2-cffi               21.3.0                   pypi_0    pypi
argon2-cffi-bindings      21.2.0                   pypi_0    pypi
asttokens                 2.2.1                    pypi_0    pypi
async-timeout             4.0.2                    pypi_0    pypi
attrs                     23.1.0                   pypi_0    pypi
av                        9.2.0                    pypi_0    pypi
babel                     2.12.1                   pypi_0    pypi
backcall                  0.2.0                    pypi_0    pypi
beautifulsoup4            4.12.2                   pypi_0    pypi
bidict                    0.22.1                   pypi_0    pypi
blas                      1.0                         mkl  
bleach                    6.0.0                    pypi_0    pypi
brotlipy                  0.7.0           py39h27cfd23_1003  
bzip2                     1.0.8                h7b6447c_0  
ca-certificates           2022.10.11           h06a4308_0  
cachetools                5.3.0                    pypi_0    pypi
certifi                   2022.12.7        py39h06a4308_0  
cffi                      1.15.1           py39h5eee18b_3  
charset-normalizer        2.0.4              pyhd3eb1b0_0  
click                     8.1.3                    pypi_0    pypi
comm                      0.1.3                    pypi_0    pypi
commonmark                0.9.1                    pypi_0    pypi
configargparse            1.5.3                    pypi_0    pypi
contourpy                 1.0.7                    pypi_0    pypi
cryptography              38.0.4           py39h9ce1e76_0  
cuda                      11.7.1                        0    nvidia
cuda-cccl                 11.7.91                       0    nvidia
cuda-command-line-tools   11.7.1                        0    nvidia
cuda-compiler             11.7.1                        0    nvidia
cuda-cudart               11.7.99                       0    nvidia
cuda-cudart-dev           11.7.99                       0    nvidia
cuda-cuobjdump            11.7.91                       0    nvidia
cuda-cupti                11.7.101                      0    nvidia
cuda-cuxxfilt             11.7.91                       0    nvidia
cuda-demo-suite           12.0.76                       0    nvidia
cuda-documentation        12.0.76                       0    nvidia
cuda-driver-dev           11.7.99                       0    nvidia
cuda-gdb                  12.0.90                       0    nvidia
cuda-libraries            11.7.1                        0    nvidia
cuda-libraries-dev        11.7.1                        0    nvidia
cuda-memcheck             11.8.86                       0    nvidia
cuda-nsight               12.0.78                       0    nvidia
cuda-nsight-compute       12.0.0                        0    nvidia
cuda-nvcc                 11.7.99                       0    nvidia
cuda-nvdisasm             12.0.76                       0    nvidia
cuda-nvml-dev             11.7.91                       0    nvidia
cuda-nvprof               12.0.90                       0    nvidia
cuda-nvprune              11.7.91                       0    nvidia
cuda-nvrtc                11.7.99                       0    nvidia
cuda-nvrtc-dev            11.7.99                       0    nvidia
cuda-nvtx                 11.7.91                       0    nvidia
cuda-nvvp                 12.0.90                       0    nvidia
cuda-runtime              11.7.1                        0    nvidia
cuda-sanitizer-api        12.0.90                       0    nvidia
cuda-toolkit              11.7.1                        0    nvidia
cuda-tools                11.7.1                        0    nvidia
cuda-visual-tools         11.7.1                        0    nvidia
cycler                    0.11.0                   pypi_0    pypi
dash                      2.9.3                    pypi_0    pypi
dash-core-components      2.0.0                    pypi_0    pypi
dash-html-components      2.0.0                    pypi_0    pypi
dash-table                5.0.0                    pypi_0    pypi
debugpy                   1.6.7                    pypi_0    pypi
decorator                 5.1.1                    pypi_0    pypi
defusedxml                0.7.1                    pypi_0    pypi
descartes                 1.1.0                    pypi_0    pypi
docker-pycreds            0.4.0                    pypi_0    pypi
docstring-parser          0.14.1                   pypi_0    pypi
executing                 1.2.0                    pypi_0    pypi
faiss-cpu                 1.7.2                    pypi_0    pypi
faiss-gpu                 1.7.2                    pypi_0    pypi
fastjsonschema            2.16.3                   pypi_0    pypi
ffmpeg                    4.3                  hf484d3e_0    pytorch
filelock                  3.12.0                   pypi_0    pypi
fire                      0.5.0                    pypi_0    pypi
flask                     2.2.3                    pypi_0    pypi
flit-core                 3.6.0              pyhd3eb1b0_0  
fonttools                 4.38.0                   pypi_0    pypi
freetype                  2.12.1               h4a9f257_0  
frozendict                2.3.7                    pypi_0    pypi
frozenlist                1.3.3                    pypi_0    pypi
fsspec                    2023.4.0                 pypi_0    pypi
functorch                 0.2.1                    pypi_0    pypi
gdown                     4.5.1                    pypi_0    pypi
gds-tools                 1.5.0.59                      0    nvidia
giflib                    5.2.1                h7b6447c_0  
gitdb                     4.0.10                   pypi_0    pypi
gitpython                 3.1.31                   pypi_0    pypi
gmp                       6.2.1                h295c915_3  
gnutls                    3.6.15               he1e5248_0  
google-auth               2.17.3                   pypi_0    pypi
google-auth-oauthlib      0.4.6                    pypi_0    pypi
grpcio                    1.54.0                   pypi_0    pypi
h5py                      3.8.0                    pypi_0    pypi
huggingface-hub           0.13.4                   pypi_0    pypi
idna                      3.4              py39h06a4308_0  
imageio                   2.21.1                   pypi_0    pypi
importlib-metadata        6.6.0                    pypi_0    pypi
intel-openmp              2021.4.0          h06a4308_3561  
ipykernel                 6.22.0                   pypi_0    pypi
ipython                   8.12.0                   pypi_0    pypi
ipython-genutils          0.2.0                    pypi_0    pypi
ipywidgets                8.0.6                    pypi_0    pypi
itsdangerous              2.1.2                    pypi_0    pypi
jedi                      0.18.2                   pypi_0    pypi
jinja2                    3.1.2                    pypi_0    pypi
joblib                    1.2.0                    pypi_0    pypi
jpeg                      9e                   h7f8727e_0  
json5                     0.9.11                   pypi_0    pypi
jsonschema                4.17.3                   pypi_0    pypi
jupyter                   1.0.0                    pypi_0    pypi
jupyter-client            8.2.0                    pypi_0    pypi
jupyter-console           6.6.3                    pypi_0    pypi
jupyter-core              5.3.0                    pypi_0    pypi
jupyter-server            1.24.0                   pypi_0    pypi
jupyterlab                3.3.4                    pypi_0    pypi
jupyterlab-pygments       0.2.2                    pypi_0    pypi
jupyterlab-server         2.22.1                   pypi_0    pypi
jupyterlab-widgets        3.0.7                    pypi_0    pypi
kiwisolver                1.4.4                    pypi_0    pypi
lame                      3.100                h7b6447c_0  
lcms2                     2.12                 h3be6417_0  
ld_impl_linux-64          2.38                 h1181459_1  
lerc                      3.0                  h295c915_0  
libcublas                 11.10.3.66                    0    nvidia
libcublas-dev             11.10.3.66                    0    nvidia
libcufft                  10.7.2.124           h4fbf590_0    nvidia
libcufft-dev              10.7.2.124           h98a8f43_0    nvidia
libcufile                 1.5.0.59                      0    nvidia
libcufile-dev             1.5.0.59                      0    nvidia
libcurand                 10.3.1.50                     0    nvidia
libcurand-dev             10.3.1.50                     0    nvidia
libcusolver               11.4.0.1                      0    nvidia
libcusolver-dev           11.4.0.1                      0    nvidia
libcusparse               11.7.4.91                     0    nvidia
libcusparse-dev           11.7.4.91                     0    nvidia
libdeflate                1.8                  h7f8727e_5  
libffi                    3.4.2                h6a678d5_6  
libgcc-ng                 11.2.0               h1234567_1  
libgomp                   11.2.0               h1234567_1  
libiconv                  1.16                 h7f8727e_2  
libidn2                   2.3.2                h7f8727e_0  
libnpp                    11.7.4.75                     0    nvidia
libnpp-dev                11.7.4.75                     0    nvidia
libnvjpeg                 11.8.0.2                      0    nvidia
libnvjpeg-dev             11.8.0.2                      0    nvidia
libpng                    1.6.37               hbc83047_0  
libstdcxx-ng              11.2.0               h1234567_1  
libtasn1                  4.16.0               h27cfd23_0  
libtiff                   4.5.0                hecacb30_0  
libunistring              0.9.10               h27cfd23_0  
libwebp                   1.2.4                h11a3e52_0  
libwebp-base              1.2.4                h5eee18b_0  
lightning-utilities       0.3.0                    pypi_0    pypi
lpips                     0.1.4                    pypi_0    pypi
lz4-c                     1.9.4                h6a678d5_0  
markdown                  3.4.3                    pypi_0    pypi
markdown-it-py            2.2.0                    pypi_0    pypi
markupsafe                2.1.2                    pypi_0    pypi
matplotlib                3.5.3                    pypi_0    pypi
matplotlib-inline         0.1.6                    pypi_0    pypi
mdurl                     0.1.2                    pypi_0    pypi
mediapy                   1.1.0                    pypi_0    pypi
mistune                   2.0.5                    pypi_0    pypi
mkl                       2021.4.0           h06a4308_640  
mkl-service               2.4.0            py39h7f8727e_0  
mkl_fft                   1.3.1            py39hd3c417c_0  
mkl_random                1.2.2            py39h51133e4_0  
msgpack                   1.0.4                    pypi_0    pypi
msgpack-numpy             0.4.8                    pypi_0    pypi
multidict                 6.0.4                    pypi_0    pypi
nbclassic                 0.5.5                    pypi_0    pypi
nbclient                  0.7.3                    pypi_0    pypi
nbconvert                 7.3.1                    pypi_0    pypi
nbformat                  5.7.0                    pypi_0    pypi
ncurses                   6.3                  h5eee18b_3  
nerfacc                   0.3.5                    pypi_0    pypi
nerfstudio                0.1.12                   pypi_0    pypi
nest-asyncio              1.5.6                    pypi_0    pypi
nettle                    3.7.3                hbbd107a_1  
networkx                  3.1                      pypi_0    pypi
ninja                     1.10.2.3                 pypi_0    pypi
notebook                  6.5.4                    pypi_0    pypi
notebook-shim             0.2.3                    pypi_0    pypi
nsight-compute            2022.4.0.15                   0    nvidia
numpy                     1.23.5           py39h14f4228_0  
numpy-base                1.23.5           py39h31eccc5_0  
nuscenes-devkit           1.1.9                    pypi_0    pypi
oauthlib                  3.2.2                    pypi_0    pypi
open3d                    0.17.0                   pypi_0    pypi
opencv-python             4.6.0.66                 pypi_0    pypi
openh264                  2.1.1                h4ff587b_0  
openssl                   1.1.1s               h7f8727e_0  
packaging                 23.0                     pypi_0    pypi
pandas                    1.5.1                    pypi_0    pypi
pandocfilters             1.5.0                    pypi_0    pypi
parso                     0.8.3                    pypi_0    pypi
pathtools                 0.1.2                    pypi_0    pypi
pexpect                   4.8.0                    pypi_0    pypi
pickleshare               0.7.5                    pypi_0    pypi
pillow                    9.3.0            py39hace64e9_1  
pip                       23.1.1                   pypi_0    pypi
platformdirs              3.2.0                    pypi_0    pypi
plotly                    5.7.0                    pypi_0    pypi
prettytable               3.5.0                    pypi_0    pypi
prometheus-client         0.16.0                   pypi_0    pypi
prompt-toolkit            3.0.38                   pypi_0    pypi
protobuf                  3.20.0                   pypi_0    pypi
psutil                    5.9.5                    pypi_0    pypi
ptyprocess                0.7.0                    pypi_0    pypi
pure-eval                 0.2.2                    pypi_0    pypi
pyasn1                    0.5.0                    pypi_0    pypi
pyasn1-modules            0.3.0                    pypi_0    pypi
pycocotools               2.0.6                    pypi_0    pypi
pycparser                 2.21               pyhd3eb1b0_0  
pygments                  2.15.1                   pypi_0    pypi
pymeshlab                 2022.2.post2             pypi_0    pypi
pyngrok                   5.1.0                    pypi_0    pypi
pyopenssl                 22.0.0             pyhd3eb1b0_0  
pyparsing                 3.0.9                    pypi_0    pypi
pyquaternion              0.9.9                    pypi_0    pypi
pyrsistent                0.19.3                   pypi_0    pypi
pysocks                   1.7.1            py39h06a4308_0  
python                    3.9.16               h7a1cb2a_0  
python-dateutil           2.8.2                    pypi_0    pypi
python-engineio           4.4.1                    pypi_0    pypi
python-socketio           5.7.1                    pypi_0    pypi
pytorch-cuda              11.7                 h67b0de4_1    pytorch
pytorch-lightning         1.8.3                    pypi_0    pypi
pytorch-metric-learning   1.6.3                    pypi_0    pypi
pytorch-mutex             1.0                        cuda    pytorch
pytz                      2023.3                   pypi_0    pypi
pywavelets                1.4.1                    pypi_0    pypi
pyyaml                    6.0                      pypi_0    pypi
pyzmq                     25.0.2                   pypi_0    pypi
qtconsole                 5.4.2                    pypi_0    pypi
qtpy                      2.3.1                    pypi_0    pypi
readline                  8.2                  h5eee18b_0  
requests                  2.28.1           py39h06a4308_0  
requests-oauthlib         1.3.1                    pypi_0    pypi
rich                      12.5.1                   pypi_0    pypi
rsa                       4.9                      pypi_0    pypi
scikit-image              0.19.3                   pypi_0    pypi
scikit-learn              1.2.2                    pypi_0    pypi
scipy                     1.10.1                   pypi_0    pypi
send2trash                1.8.0                    pypi_0    pypi
sentry-sdk                1.20.0                   pypi_0    pypi
setproctitle              1.3.2                    pypi_0    pypi
setuptools                67.7.2                   pypi_0    pypi
shapely                   2.0.1                    pypi_0    pypi
shtab                     1.6.1                    pypi_0    pypi
six                       1.16.0             pyhd3eb1b0_1  
smmap                     5.0.0                    pypi_0    pypi
sniffio                   1.3.0                    pypi_0    pypi
soupsieve                 2.4.1                    pypi_0    pypi
sqlite                    3.40.1               h5082296_0  
stack-data                0.6.2                    pypi_0    pypi
tenacity                  8.2.2                    pypi_0    pypi
tensorboard               2.9.0                    pypi_0    pypi
tensorboard-data-server   0.6.1                    pypi_0    pypi
tensorboard-plugin-wit    1.8.1                    pypi_0    pypi
tensorboardx              2.6                      pypi_0    pypi
tensorrt                  8.5.2.2                  pypi_0    pypi
termcolor                 2.2.0                    pypi_0    pypi
terminado                 0.17.1                   pypi_0    pypi
threadpoolctl             3.1.0                    pypi_0    pypi
tifffile                  2023.4.12                pypi_0    pypi
timm                      0.6.12                   pypi_0    pypi
tinycss2                  1.2.1                    pypi_0    pypi
tinycudann                1.7                      pypi_0    pypi
tk                        8.6.12               h1ccaba5_0  
torch                     1.12.1+cu116             pypi_0    pypi
torch-fidelity            0.3.0                    pypi_0    pypi
torch2trt-dynamic         0.5.0                    pypi_0    pypi
torchaudio                0.12.1+cu116             pypi_0    pypi
torchmetrics              0.11.4                   pypi_0    pypi
torchtyping               0.1.4                    pypi_0    pypi
torchvision               0.13.1+cu116             pypi_0    pypi
tornado                   6.3.1                    pypi_0    pypi
tqdm                      4.64.1                   pypi_0    pypi
traitlets                 5.9.0                    pypi_0    pypi
trimesh                   3.15.8                   pypi_0    pypi
typeguard                 3.0.2                    pypi_0    pypi
typing_extensions         4.4.0            py39h06a4308_0  
tyro                      0.4.2                    pypi_0    pypi
tzdata                    2022g                h04d1e81_0  
u-msgpack-python          2.7.2                    pypi_0    pypi
urllib3                   1.26.13          py39h06a4308_0  
wandb                     0.15.0                   pypi_0    pypi
wcwidth                   0.2.6                    pypi_0    pypi
webencodings              0.5.1                    pypi_0    pypi
websocket-client          1.5.1                    pypi_0    pypi
werkzeug                  2.2.3                    pypi_0    pypi
wheel                     0.37.1             pyhd3eb1b0_0  
widgetsnbextension        4.0.7                    pypi_0    pypi
xatlas                    0.0.7                    pypi_0    pypi
xz                        5.2.10               h5eee18b_0  
yarl                      1.8.2                    pypi_0    pypi
zipp                      3.15.0                   pypi_0    pypi
zlib                      1.2.13               h5eee18b_0  
zstd                      1.5.2                ha4553b6_0

can anybody give me some advice?

[CanCanZeng]

This is the extracted mesh using TSDF method

[niujinshuchong]

@CanCanZeng Could you try to use ns-extract-mesh for mesh extraction? You could find the command here: https://github.com/autonomousvision/sdfstudio/blob/master/docs/sdfstudio-examples.md#bakedsdf-on-the-mipnerf360-dataset

[CanCanZeng]

Thank you for your advice @niujinshuchong , it works!

[flow-specter]

Hi, how long does bakedsdf to be trained to have a good result?