3D 模型生成在计算机视觉领域有着广泛的应用，从虚拟现实到自动驾驶，基于单张图像的 3D 重建技术正在迅速发展。这篇博客将带你深入探索如何使用 Zero-1-2-3 框架进行多视图 3D 重建，通过详细解析该框架中的代码结构和功能，帮助你理解其工作原理并实际应用到 3D 模型生成中。
一、项目概述
Zero-1-2-3 是一个通过扩散模型（Diffusion Models）将单张 2D 图像转换为 3D 网格模型的项目。与传统的 3D 重建技术不同，Zero-1-2-3 采用了一个基于图像生成的多阶段流程，利用不同视角的多张图片推理，生成高精度的 3D 模型。
该项目使用了以下几个重要的组件：
- Zero-1-2-3 模型：基于预训练的扩散模型，生成目标图像的多视角图像。
- Segment Anything Model (SAM)：用于处理图像中的前景和背景。
- 多视角生成：从单张输入图像开始，通过多阶段生成多个视角下的图片，再进行 3D 重建。
二、代码解析
Zero-1-2-3 的代码逻辑分为多个部分，包括图像预处理、多视角生成、推理过程、3D 模型重建等。在本博客中，我们会详细解析每一部分的核心功能。
1. 图像预处理
在 3D 重建过程中，输入图像的预处理尤为重要。首先我们使用 Segment Anything Model (SAM) 对图像进行分割，去除背景噪声，将前景进行提取，并对图像进行缩放，使其适应网络的输入。
def preprocess(predictor, raw_im, lower_contrast=False):
raw_im.thumbnail([512, 512], Image.Resampling.LANCZOS)
image_sam = sam_out_nosave(predictor, raw_im.convert("RGB"), pred_bbox(raw_im))
input_256 = image_preprocess_nosave(image_sam, lower_contrast=lower_contrast, rescale=True)
torch.cuda.empty_cache()
return input_256
在这个函数中，我们首先对图像进行了缩放，确保其尺寸适应 512x512 的输入格式。然后通过 SAM 模型预测并分割图像中的前景部分，最后将图像数据转换为张量供后续的模型使用。
2. 第一阶段视角生成
在第一阶段，我们从输入图像生成多视角的图片。通过生成不同角度的图像，Zero-1-2-3 可以在生成过程中推理出物体的深度和形状。
def stage1_run(model, device, exp_dir, input_im, scale, ddim_steps):
stage1_dir = os.path.join(exp_dir, "stage1_8")
os.makedirs(stage1_dir, exist_ok=True)
output_ims = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4)), device=device, ddim_steps=ddim_steps, scale=scale)
stage2_steps = 50
zero123_infer(model, exp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
polar_angle = estimate_elev(exp_dir)
gen_poses(exp_dir, polar_angle)
if polar_angle <= 75:
output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4,8)), device=device, ddim_steps=ddim_steps, scale=scale)
else:
output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(8,12)), device=device, ddim_steps=ddim_steps, scale=scale)
torch.cuda.empty_cache()
return 90 - polar_angle, output_ims + output_ims_2
第一阶段使用 predict_stage1_gradio 函数生成四个不同视角下的图像，并通过极角估计 (estimate_elev) 计算输入图像的极角位置，进而根据极角生成更多视角的图片。这个过程有助于模拟物体在不同视角下的外观。
3. 第二阶段视角推理
第二阶段我们为每一个视角进一步生成四张图片，从而获得更详细的 3D 信息。这使得物体的表面和轮廓可以通过多个角度的图像捕获。
def stage2_run(model, device, exp_dir, elev, scale, stage2_steps=50):
if 90 - elev <= 75:
zero123_infer(model, exp_dir, indices=list(range(1,8)), device=device, ddim_steps=stage2_steps, scale=scale)
else:
zero123_infer(model, exp_dir, indices=list(range(1,4)) + list(range(8,12)), device=device, ddim_steps=stage2_steps, scale=scale)
第二阶段将为每个生成的视角图像推理更多细节，从而为最终的 3D 重建提供更多的信息。
4. 3D 重建
在视角生成完成之后，我们进入最终的 3D 重建阶段。该过程基于生成的多视角图像，通过扩展方法将它们重建为 3D 网格模型。
def reconstruct(exp_dir, output_format=".ply", device_idx=0, resolution=256):
exp_dir = os.path.abspath(exp_dir)
main_dir_path = os.path.abspath(os.path.dirname("./"))
os.chdir('reconstruction/')
bash_script = f'CUDA_VISIBLE_DEVICES={device_idx} python exp_runner_generic_blender_val.py \
--specific_dataset_name {exp_dir} \
--mode export_mesh \
--conf confs/one2345_lod0_val_demo.conf \
--resolution {resolution}'
print(bash_script)
os.system(bash_script)
os.chdir(main_dir_path)
ply_path = os.path.join(exp_dir, f"mesh.ply")
if output_format == ".ply":
return ply_path
return convert_mesh_format(exp_dir, output_format=output_format)
在这个函数中，使用 Blender 等工具进行 3D 网格重建，并将生成的 3D 模型导出为 .ply 格式。用户可以选择将其转换为 .obj 或 .glb 格式，方便在不同平台上查看和使用。
三、整体流程
- 图像输入：首先，用户输入一张 2D 图像，模型会对其进行预处理，提取图像中的前景部分。
- 多视角生成：通过多阶段推理，模型生成该图像在不同视角下的表现，涵盖物体的多个角度。
- 极角估计：根据生成的图像，模型估算出输入图像的极角，并根据极角生成更多细节图像。
- 3D 重建：最终，所有视角图像被用来进行 3D 重建，生成 3D 网格模型。
四、总结
Zero-1-2-3 项目展示了基于扩散模型的 3D 重建技术，其核心优势在于无需复杂的每个形状优化过程，通过简单的前向推理即可生成高质量的 3D 模型。对于希望快速进行 3D 模型生成的开发者，该框架提供了强大的功能和便捷的使用方式。希望本文能帮助你理解该项目的工作原理，并激发你对 3D 重建领域的兴趣。
