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 = 50zero123_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_pathreturn 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 重建领域的兴趣。