NeuRBF：基于自适应径向基函数的神经场表示

ICCV 2023（口头报告）

张晨1，李忠1，宋良辰2，陈乐乐1，余荧奕3，袁骏松2，徐毅1
1OPPO美国研究中心，2布法罗大学，3上海科技大学

项目主页 | 论文 | 视频

本仓库是ICCV 2023论文"NeuRBF：基于自适应径向基函数的神经场表示"的官方PyTorch实现。我们的工作提出了一种新型神经场，具有高表示精度和模型紧凑性。该仓库包含图像拟合、SDF拟合和神经辐射场的代码。

安装

克隆仓库

git clone https://github.com/oppo-us-research/NeuRBF.git
cd NeuRBF

安装conda环境

# 创建conda环境
conda create -n neurbf python=3.9 -y
conda activate neurbf

# 安装CuPy
pip install cupy-cuda11x
python -m cupyx.tools.install_library --cuda 11.x --library cutensor
python -m cupyx.tools.install_library --cuda 11.x --library cudnn
python -m cupyx.tools.install_library --cuda 11.x --library nccl

# 安装PyTorch
pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117

# 用于图像/SDF拟合和合成NeRF数据集上的NeRF任务
pip install einops matplotlib kornia imageio imageio-ffmpeg opencv-python pysdf PyMCubes trimesh plotly scipy GPUtil scikit-image scikit-learn pykdtree commentjson tqdm configargparse lpips tensorboard torch-ema ninja tensorboardX numpy pandas rich packaging scipy torchmetrics jax pillow plyfile omegaconf

# 用于LLFF数据集上的NeRF任务
pip install jax tqdm pillow opencv-python pandas lpips imageio torchmetrics scikit-image tensorboard matplotlib
pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch

构建torch-ngp扩展（用于图像/SDF拟合）

cd thirdparty/torch_ngp/gridencoder
pip install .
cd ../../../

测试环境

Ubuntu 18.04，PyTorch 1.13.0 & CUDA 11.6，RTX A6000。
Ubuntu 18.04，PyTorch 2.0.1 & CUDA 11.7，RTX A6000。
Ubuntu 18.04，PyTorch 2.0.1 & CUDA 11.7，RTX 3090。
Windows 11，PyTorch 2.0.1 & CUDA 11.7，RTX 3080 Ti 16G。

使用方法

图像拟合

拟合单张图像

下载示例冥王星图像。将其放在data/img中并重命名为pluto.png。然后运行

python main.py --config configs/img.py --path ./data/img/pluto.png --alias pluto

结果和tensorboard日志将位于log/img中。要调整模型大小，可以指定--log2_hashmap_size_ref参数的值。

注意：为减少GPU内存使用，可以在上述命令中添加--ds_device cpu。这将把一些数据放在CPU而不是GPU上，但也会使训练变慢。

拟合DIV2K数据集中的所有图像

下载DIV2K数据集的验证集并将其放在data/img/div2k中。每张图像的路径应为data/img/div2k/DIV2K_valid_HR/xxxx.png。然后运行

python main_img_div2k.py

该脚本将分别拟合每张图像，结果将存储在log/img_div2k中。

SDF拟合

下载斯坦福3D扫描库中的示例犰狳网格。解压后将其放在data/sdf中，并重命名为armadillo.ply。

运行以下预处理脚本，该脚本将对网格进行归一化并采样额外的评估点。

python preproc_mesh.py --path ./data/sdf/armadillo.ply

然后运行

python main.py --config configs/sdf.py --path ./data/sdf/armadillo_nrml.obj --alias armadillo

结果和tensorboard日志将位于log/sdf中。要调整模型大小，可以指定--log2_hashmap_size_ref参数的值。注意：为减少GPU内存使用，您可以类似地在上述命令中添加--ds_device cpu。

NeRF

数据集

在合成NeRF数据集上运行

下载数据集并解压到data目录。例如，lego场景的路径应为data/nerf_synthetic/lego。

对于训练，使用以下命令

python main_nerf.py --config_init configs/nerf_tensorf/nerf_synthetic_init.py --config configs/nerf_tensorf/nerf_synthetic.py --data_name lego

它将首先提取场景信息以初始化RBF位置和形状参数，然后训练完整模型。结果和tensorboard日志将位于log/nerf_synthetic。更改--data_name以在其他场景上运行。

注意：为减少GPU内存使用，您可以添加--batch_size_init 2048（已在RTX 3090 24G上测试可用）。如有需要，可以进一步降低数字2048。

要使用训练好的模型渲染图像，请运行

python main_nerf.py --config configs/nerf_tensorf/nerf_synthetic.py --data_name lego --ckpt [checkpoint路径] --render_only 1 --[要渲染的内容] 1

其中[checkpoint路径]是检查点文件的路径，[要渲染的内容]可以是render_test或render_train。请注意，训练命令已在训练结束时执行render_test。

要从训练好的模型中提取网格，请运行

python main_nerf.py --config configs/nerf_tensorf/nerf_synthetic.py --data_name lego --ckpt [checkpoint路径] --export_mesh 1

在LLFF数据集上运行

下载数据集并解压到data目录。例如，room场景的路径应为data/nerf_llff_data/room。

对于训练，使用以下命令

python main_nerf_kplanes.py --config_init configs/nerf_kplanes/llff_init.py --config configs/nerf_kplanes/llff.py --data_name room

它将首先提取场景信息以初始化RBF位置和形状参数，然后训练完整模型。结果和tensorboard日志将位于log/llff。更改--data_name以在其他场景上运行。

要使用训练好的模型渲染螺旋路径视频，请运行

python main_nerf_kplanes.py --config configs/nerf_kplanes/llff.py --log-dir [模型文件夹] --render-only

其中[模型文件夹]是包含训练好模型的文件夹路径。

致谢

我们诚挚感谢以下仓库的作者，他们的工作在我们发布的代码中得到了引用：torch-ngp、TensoRF、nrff、K-Planes、MINER_pl和siren。

引用

如果您觉得我们的工作有用，请考虑引用：

@inproceedings{chen2023neurbf,
    author    = {Chen, Zhang and Li, Zhong and Song, Liangchen and Chen, Lele and Yu, Jingyi and Yuan, Junsong and Xu, Yi},
    title     = {NeuRBF: A Neural Fields Representation with Adaptive Radial Basis Functions},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2023},
    pages     = {4182-4194}
}