基于向量量化的神经网络辐射场压缩

当前主流的可实时推理的NeRF方法，如Plenoxel、Dvgo存在模型存储消耗巨大的问题，一个模型需要上百兆，不利于NeRF渲染方案的实用化。本项目使用了向量量化技术对NeRF模型进行了压缩，大大减小了模型大小，同时保持了较好的重建效果。

Clone with HTTP

 git clone https://www.modelscope.cn/DAMOXR/cv_nerf_3d-reconstruction_vector-quantize-compression.git

方法描述

本方法通过对当前NeRF模型存储内容的分析，发现模型中存在大量冗余的体素(Voxel)，因此首先分析体素的重要性，又针对不同的场景采用自适应的阈值来划分体素的重要程度。对于重要的体素，我们使用向量量化的方法对其进行压缩，从而达到减小模型存储大小的目的。

在多个数据集上的渲染可视化效果也表明了对效果几乎没有影响。

使用方式和范围

适用范围： 当前支持两种数据类型，分别是nerf_synthetic和LLFF

• nerf-synthesis数据集人工合成的多视角数据集，包含8个场景，下载及更多信息请参考NeRF三维重建数据集
• LLFF 数据集包含真实拍摄的多视角数据集，包含8个场景,下载及更多信息请参考NeRF_LLFF真实场景数据集

运行环境：

• 模型只支持在GPU上运行，已在V100,A100,RTX3090上测试通过，具体效率与显卡性能相关
• GPU 显存需求： 8G以上

使用方法

推理代码范例

Blender 数据集：

import os
from modelscope.msdatasets import MsDataset
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from modelscope.utils.test_utils import test_level

model_id = 'DAMOXR/cv_nerf_3d-reconstruction_vector-quantize-compression'
pretrained_model = 'ficus_demo.pt'
data_dir = MsDataset.load('nerf_recon_dataset', namespace='damo',
                    split='train').config_kwargs['split_config']['train']
nerf_synthetic_dataset = os.path.join(data_dir, 'nerf_synthetic')
blender_scene = 'ficus'
data_dir = os.path.join(nerf_synthetic_dataset, blender_scene)

nerf_vq_compression = pipeline(Tasks.nerf_recon_vq_compression, 
                                 model=model_id,
                                 dataset_name='blender',
                                 data_dir=data_dir,
                                 downsample=1,
                                 ndc_ray=False,
                                 ckpt_path=pretrained_model
                                 )
render_dir = os.path.join('./exp/', blender_scene)

# For evaluation test dataset to get metrics like psnr, ssim, N_vis is the number of test images
nerf_vq_compression(dict(test_mode='evaluation_test', render_dir=render_dir, N_vis=5))

# For rendering novel views, N_vis is the number of views
nerf_vq_compression(dict(test_mode='render_path', render_dir=render_dir, N_vis=30))

LLFF 数据集：

model_id = 'DAMOXR/cv_nerf_3d-reconstruction_vector-quantize-compression'
pretrained_model = 'fern_demo.pt'
data_dir = MsDataset.load(
    'DAMOXR/nerf_llff_data',
    subset_name='default',
    split='test',
).config_kwargs['split_config']['test']
nerf_llff = os.path.join(data_dir, 'nerf_llff_data')
llff_scene = 'fern'
data_dir = os.path.join(nerf_llff, llff_scene)

nerf_vq_compression = pipeline(Tasks.nerf_recon_vq_compression, 
                              model=model_id,
                              dataset_name='llff',
                              data_dir=data_dir,
                              downsample=4,
                              ndc_ray=True,
                              ckpt_path=pretrained_model
                              )
render_dir = os.path.join('./exp/', llff_scene)

# For evaluation test dataset to get metrics like psnr, ssim, N_vis is the number of test images
nerf_vq_compression(dict(test_mode='evaluation_test', render_dir=render_dir, N_vis=5))

# For rendering novel views, N_vis is the number of views
nerf_vq_compression(dict(test_mode='render_path', render_dir=render_dir, N_vis=10))

数据评估及结果

Blender 数据集

压缩前:

VQRF压缩后:

LLFF 数据集

压缩前:

VQRF压缩后:

方法引用

@inproceedings{li2023compressing,
  title={Compressing volumetric radiance fields to 1 mb},
  author={Li, Lingzhi and Shen, Zhen and Wang, Zhongshu and Shen, Li and Bo, Liefeng},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={4222--4231},
  year={2023}
}