MeshDiffusion：基于分数的生成式3D网格建模

简介

这是 MeshDiffusion（ICLR 2023 Spotlight）的官方实现。

MeshDiffusion 是一种用于生成3D网格的扩散模型，它直接使用深度Marching四面体（DMTet）进行参数化。更多详细信息和交互式演示，请参阅我们的项目页面。

快速入门

要求

• Python >= 3.8
• CUDA 11.6
• Pytorch >= 1.6
• Pytorch3D
• ml_collections

请按照 nvdiffrec 的说明安装所需依赖。

预训练模型

下载我们预训练的MeshDiffusion模型（分辨率为64），适用于椅子、汽车、飞机、桌子 和 步枪。作为备用方案，您也可以从 Google Drive 下载汽车和椅子的模型。

在此下载分辨率为128的模型：汽车 和 椅子。

数据集

我们提供了分辨率为64的处理后数据集（以立方体网格形式），可在此链接获取。数据集的形变尺度设置为3.0，所有非网格生成四面体顶点的SDF值根据其符号被设置为1或-1，如论文所述。移除边界的立方体网格大小为63x63x63，并为了方便起见在右侧填充至64x64x64。请查看 eval.py 以了解如何从3D立方体网格中提取DMTet表示。

推理

无条件生成

运行以下命令：

python main_diffusion.py --config=$DIFFUSION_CONFIG --mode=uncond_gen \
--config.eval.eval_dir=$OUTPUT_PATH \
--config.eval.ckpt_path=$CKPT_PATH

随后运行：

cd nvdiffrec
python eval.py --config $DMTET_CONFIG --out-dir $OUT_DIR --sample-path $SAMPLE_PATH \
--deform-scale $DEFORM_SCALE [--angle-ind $ANGLE_INDEX] [--num-smoothing-steps $NUM_SMOOTHING_STEPS]

其中 $SAMPLE_PATH 是 $OUTPUT_PATH 中生成的 .npy 样本文件，而 $DEFORM_SCALE 是为DMTet数据集设置的四面体顶点形变尺度（默认情况下，分辨率为64时我们使用3.0；请根据您自己的数据集调整该值）。如果需要从不同角度渲染图像，请将 $ANGLE_INDEX 更改为0到50之间的某个数字。如果后期处理中的拉普拉斯平滑效果过于强烈，可以调整 $NUM_SMOOTHING_STEPS（默认值为3）。例如，对于飞机类别，我们建议将 $NUM_SMOOTHING_STEPS 设置为1或2。

最终会保存一个 .obj 格式的网格文件到指定文件夹中，您可以在MeshLab等工具中查看。保存的图像来自未经后期处理的原始网格，因此仅用于快速验证。

单视角条件生成

首先，从处于规范姿态的网格单视图中拟合一个DMTet：

cd nvdiffrec
python fit_singleview.py --config $DMTET_CONFIG --mesh-path $MESH_PATH --angle-ind $ANGLE_IND --out-dir $OUT_DIR --validate $VALIDATE

其中 $ANGLE_IND 是控制物体绕z轴旋转的整数（0到50）。如果需要可视化拟合后的DMTet，请将 $VALIDATE 设置为1。

然后使用训练好的扩散模型来补全遮挡区域：

cd ..

python main_diffusion.py --mode=cond_gen --config=$DIFFUSION_CONFIG \
--config.eval.eval_dir=$EVAL_DIR \
--config.eval.ckpt_path=$CKPT_PATH \
--config.eval.partial_dmtet_path=$OUT_DIR/tets/dmtet.pt \
--config.eval.tet_path=$TET_PATH \
--config.eval.batch_size=$EVAL_BATCH_SIZE

其中 $TET_PATH 是位于 nvdiffrec/data/tets 目录下的均匀四面体网格文件（分辨率为64或128）。

现在将完成的网格以 .obj 文件的形式存储在 $SAMPLE_PATH 中：

cd nvdiffrec
python eval.py --config $DMTET_CONFIG --sample-path $SAMPLE_PATH \
--deform-scale $DEFORM_SCALE [--angle-ind $ANGLE_INDEX] [--num-smoothing-steps $NUM_SMOOTHING_STEPS]

注意：单视图拟合和扩散模型所使用的形变尺度应保持一致。在进行条件生成之前请务必检查。

训练

对于ShapeNet数据集，首先创建所有真实标签网格路径的列表，并将其保存为 ./nvdiffrec/data/shapenet_json 下的json文件。

然后运行以下命令：

cd nvdiffrec
python fit_dmtets.py --config $DMTET_CONFIG --meta-path $META_PATH --out-dir $DMTET_DATA_PATH --index 0 --split-size 100000

其中 split_size 可设置为大于数据集大小的任意大数值。如果采用多任务批量拟合，请将 split_size 调整为合适的数值，并为不同任务分配不同的 index。如有必要，可在配置文件中调整第一轮和第二轮拟合的分辨率。$META_PATH 是用于存储网格路径列表的json文件。

现在将DMTet数据集（以Python字典形式存储）转换为3D立方体网格数据集：

cd ../data/
python tets_to_3dgrid.py --resolution $RESOLUTION --root $DMTET_DICT_FOLDER --source $SOURCE_FOLDER --target grid --index 0

假设DMTet字典数据集存储在 $DMTET_DICT_FOLDER/$SOURCE_FOLDER 中，我们将生成的立方体网格数据集保存在 $DMTET_DICT_FOLDER/grid 中。

为扩散模型训练创建所有DMTet 3D立方体网格文件位置的元数据文件：

cd ../metadata/
python save_meta.py --data_path $DMTET_DICT_FOLDER/grid --json_path $META_FILE

训练扩散模型：

cd ..

python main_diffusion.py --mode=train --config=$DIFFUSION_CONFIG \
--config.data.meta_path=$META_FILE \
--config.data.filter_meta_path=$TRAIN_SPLIT_FILE

其中 $TRAIN_SPLIT_FILE 是包含要纳入训练集的索引的json列表。示例可在 metadata/train_split/ 中找到。关于扩散模型的配置文件，请参考 configs/res64.py 或 configs/res128.py。

使用我们提供的数据集进行训练

所提供的数据集以 .npy 格式存储，而非 .pt。请改用以下命令：

cd ../metadata/
python save_meta.py --data_path $DMTET_NPY_FOLDER --json_path $META_FILE --extension npy

训练扩散模型：

cd ..

python main_diffusion.py --mode=train --config=$DIFFUSION_CONFIG \
--config.data.meta_path=$META_FILE \
--config.data.filter_meta_path=$TRAIN_SPLIT_FILE \
--config.data.extension=npy

纹理生成

请按照 https://github.com/TEXTurePaper/TEXTurePaper 中的说明，为生成的网格创建文本条件纹理。

其他

如果需要更高分辨率的四面体网格，首先请遵循 nvdiffrec/data/tets 目录下的 README 文件，使用 quartet 生成均匀的四面体网格。然后运行 nvdiffrec/data/tets/crop_tets.py 脚本以移除边界（从而使结果网格满足平移对称性）。

Blender 可视化

如需使用 Blender 可视化生成的网格，请参阅 blender_viz/ 目录以获取更多详细信息。

引用

如果您认为我们的工作对您的研究有所帮助，请考虑引用以下文献：

@InProceedings{Liu2023MeshDiffusion,
    title={MeshDiffusion: 基于分数模型的三维网格生成},
    author={刘振、冯瑶、迈克尔·J·布莱克、德里克·诺鲁泽扎赖、利亚姆·保尔、刘伟阳},
    booktitle={国际学习表征会议},
    year={2023},
    url={https://openreview.net/forum?id=0cpM2ApF9p6}
}

致谢

本仓库基于 https://github.com/NVlabs/nvdiffrec 和 https://github.com/yang-song/score_sde_pytorch 进行了改编。

MeshDiffusion 快速上手指南

MeshDiffusion 是一个基于扩散模型的 3D 网格生成工具，采用深度行进四面体（DMTet）直接参数化技术。本指南将帮助你快速配置环境并运行预训练模型生成 3D 网格。

环境准备

在开始之前，请确保你的系统满足以下要求：

• 操作系统: Linux (推荐)
• Python: >= 3.8
• CUDA: 11.6
• PyTorch: >= 1.6
• 关键依赖:
  • Pytorch3D
  • ml_collections
  • nvdiffrec (核心渲染与拟合后端)

注意：本项目强依赖 NVIDIA 的 nvdiffrec 库，安装过程较为复杂，需严格遵循其官方安装指引以确保证渲染器编译成功。

安装步骤

1. 克隆仓库

git clone https://github.com/your-repo/MeshDiffusion.git
cd MeshDiffusion

2. 安装基础依赖

首先安装 Python 层面的依赖包。建议创建虚拟环境：

conda create -n meshdiffusion python=3.8
conda activate meshdiffusion
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu116
pip install ml_collections

3. 安装 Pytorch3D 和 nvdiffrec

这是最关键的一步。请严格按照 nvdiffrec 官方安装说明 进行操作。通常包括安装头文件和编译 CUDA 扩展。

简要步骤参考（具体请以 nvdiffrec README 为准）：

# 安装 Pytorch3D (可能需要从源码编译以匹配 CUDA 版本)
pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"

# 进入 nvdiffrec 子模块或目录进行安装
cd nvdiffrec
pip install -r requirements.txt
# 执行特定的构建命令（如有）
python setup.py build_ext --inplace
cd ..

4. 下载预训练模型

项目提供了椅子、汽车、飞机、桌子和步枪等类别的预训练模型（分辨率 64 和 128）。

国内加速下载推荐： 部分模型已托管在 HuggingFace，国内用户可直接使用镜像或加速链接下载：

• Table (64): HuggingFace 下载
• Rifle (64): HuggingFace 下载
• Car (128): HuggingFace 下载
• Chair (128): HuggingFace 下载

其他模型（Car, Chair, Airplane 64 版）可通过 MPDL Keeper 或 Google Drive 备份 获取。

下载后，请将 .pt 文件放置在项目适当位置或在运行时指定路径。

基本使用

以下演示如何使用预训练模型进行无条件生成（Unconditional Generation），即随机生成一个 3D 网格。

第一步：运行扩散模型生成样本

使用下载的配置文件和检查点生成原始数据（.npy 格式）。

python main_diffusion.py --config=configs/res64.py --mode=uncond_gen \
--config.eval.eval_dir=./outputs \
--config.eval.ckpt_path=./pretrained_models/chair_res64.pt

注：请根据实际下载的模型类别修改 --config 和 --config.eval.ckpt_path 路径。

第二步：转换为网格文件 (.obj)

生成的 .npy 文件需要经由 nvdiffrec 模块处理才能转换为可视化的 .obj 网格文件。

cd nvdiffrec
python eval.py --config configs/dmtet_chair.py --out-dir ../outputs/meshes \
--sample-path ../outputs/samples.npy \
--deform-scale 3.0

参数说明：

• --sample-path: 指向第一步生成的 .npy 文件路径。
• --deform-scale: 变形尺度，分辨率 64 的默认数据集通常为 3.0。若使用自定义数据集请调整此值。
• --angle-ind: (可选) 0-50 之间的整数，用于从不同角度渲染预览图。
• --num-smoothing-steps: (可选) 拉普拉斯平滑步数，默认为 3。对于飞机等尖锐物体，建议设为 1 或 2 以保留细节。

执行完成后，可在 ../outputs/meshes 目录下找到生成的 .obj 文件，使用 MeshLab 或 Blender 即可查看。

---

进阶提示： 若需进行单视图条件生成（根据一张图片补全 3D 模型），需先运行 fit_singleview.py 拟合局部 DMTet，再运行 main_diffusion.py --mode=cond_gen 进行补全，最后再次运行 eval.py 导出网格。请确保拟合阶段与扩散模型阶段的 deform-scale 参数保持一致。