Mesh Diffusion
原文链接在此https://www.vis.dolag.work/深度学习/扩散模型/实践之mesh-diffusion.html。
一直想把PCG和AIGC结合起来，但网上此类相关论文和应用颇少，难以结合。近日深入学习了扩散模型，便觉DDPM类似于泛化能力较强的GAN，适合训练大量数据进而生成新数据，于是产生了使用扩散模型生成三维模型的想法。其实主要是没有精力和能力设计三维模型的神经网络，就干脆将模型生成问题退化为图像生成问题。
本次研究的目的是，根据已有的岩石模型，将岩石模型转换为一张张的图片并输入DDPM中进行训练，训练完成后生成新的岩石模型。 准备训练集
模型to图像 因为Mesh的顶点数据不太好直接转换成图片，所以想了另一种直接将模型映射到图像或者说将方形图像映射到模型上的方法。 整个过程算法很简单，主要分为两个阶段，从方形映射到球体、从球体映射到模型。 从方形映射到球体 单纯的图片相当于一个方形面片，方形面片没有体积，不便于直接映射到模型上，因此先将方形映射到有体积的单位球体，而后将球体映射到模型。 方形映射到球体的主流算法有如下几个：
经纬球。最简单的映射方法，但经纬球的精度损失极大，图像的上下边最终只能映射到球体的两个极点上。
Cubemap。只能展开成六个正方形，而不能只展开成一个方形面片，除非损失方形的一部分内容。
八面体映射。将一个正方形以折纸的形式映射到正八面体上，然后映射到单位球上(将位置进行归一化即可)。此方法非常简单，且精度损失相对较小，且为方形到球体的满射，因此选择此方法。
作者使用Houdini批量生成模型的“图像”，因此贴出Vex的实现代码(b站还是不能贴代码块)： v@uv = set(u, v, 0); vector uv = set(u, v, 0); uv = 2 * uv - 1; vector p = set(uv.x, 1 - abs(uv.x) - abs(uv.y), uv.y); if(p.y < 0){ float x = p.x; float z = p.z; p.x = (1 - abs(z)) * sign(p.x); p.z = (1 - abs(x)) * sign(p.z); } @P = normalize(p); @P *= 0.5; @P += 0.5; // to [0, 1]
可见八面体映射确实简单异常，基本就是，然后将坐标进行归一化即可。这里将球体的位置范围给重映射到了[0, 1]的范围内，方便后续进行处理。
从球体映射到模型 这里的算法作者就设计得比较敷衍了，基本上就是smooth之后然后吸附到离模型最近的点上，并进行循环迭代。
因此算法对带有凹面、孔洞的模型结果不好，但由于岩石模型大多都是无孔洞的模型，因此效果还算可以。下图中左边是原始模型，右边是映射后的模型，在平坦表面处模型比较还原，但在曲率较大的沟壑处效果就较为一般。
需要注意的是，在上一步中，我们的球体位置坐标范围为[0, 1]，因此也需要将模型位置坐标等比缩放到[0, 1]的范围内。这样我们将方形映射到原始模型后的模型，其范围也必然在[0, 1]内，因此就可以将XYZ轴坐标分别作为图片的RGB通道，从而导出为图片。
数据增强 笔者在收集了所有Bridge上的较有特点的岩石模型后，总共也只能生成不到三百张图像，这对于DDPM的训练来说非常不利，因此需要进行一些数据增强的操作，扩充训练集。 传统数据增强方法 传统的数据增强方法包括了位移，缩放，旋转，翻转等，由于我们的模型需要完整地映射到一张图里，而位移、缩放方法会让图像损失像素，因此只能进行90°、180°、270°旋转和上下左右翻转。 但传统的数据增强方法增强效果有限，能够生成的数据集数量也有限，因此需要有新方法进行数据增强。 基于SDF(符号距离场)增强数据 在三维中，SDF就是点距离模型最近的位置的距离，其正负代表了在模型外或内部，在模型位置时SDF值为0。下图展示了一个球体模型的SDF场。下图中蓝色代表值为-1，红色代表值为0，黄色代表了值为1。
SDF可以根据模型生成，反过来也可以从SDF生成模型，使用Marching Cubes算法即可。
除了SDF可以用于生成模型之外，SDF还有许多美妙的性质——例如两个SDF之间可以进行取最大值、取最小值、线性插值等各种运算，进行运算后转模型的结果都能很好地保留两个SDF的特征。 下三张图分别是取最大值、取最小值、线性插值的结果。
可以说这三种运算的运算结果都能得到“岩石”模型。
其中线性插值运算能够生成无数的岩石模型，这就让我们能够生成极多质量较好的数据。笔者通过这种方法，随机选取两个岩石模型进行随机权重的线性插值，一共生成了2000多个数据以供训练。
那么问题来了，如果我能生成数据集，那我为什么还要训练函数来生成岩石呢(
训练模型&模型推理
模型使用Huggingface开源的Diffusers，此仓库集成了各种扩散模型，非常方便进行训练和推理，此外还有accelerator进行加速，比自己手写的DDPM又快又节约显存(。由于我们的目标只是生成岩石模型，不需要有文字、图像等条件，因此选择最简单的DDPM(Denoising Diffusion Probabilistic Models)即可，在Diffusers中DDPM叫做unconditional image generation。关于DDPM的原理，可以参考我之前的文章https://www.vis.dolag.work/深度学习/扩散模型/基础.html
代码默认导入uint8类型的图像，PIL保存图像也没法保存float32精度，因此我们需要对训练和推理代码进行一些改造。 笔者在映射模型到贴图阶段使用exr格式进行图像输出，因此需要用imageio来加载exr格式的图像。 import imageio import datasets from PIL import Image exr_files = [f for f in os.listdir(args.train_data_dir) if f.endswith(".exr")] images = [(imageio.imread(os.path.join(args.train_data_dir, f))).astype(np.float32) for f in exr_files] images = [Image.fromarray(img, mode='RGB') for img in images] features = {"image": images} dataset = datasets.Dataset.from_dict(features) 其中Dataset是huggingface定义的数据集类，想要替换原本代码中的数据集，就需要创建一个Dataset对象。 同样的，在进行推理时也需要用imageio或tifffile库保存为float32精度的图像。 图像还原为模型 笔者同样在Houdini中加载图像并还原为模型。 还原图像极其简单，加载图像到颜色属性后，直接将颜色属性作为位置属性即可。(b站专栏投稿编辑工具，到这里就没法使用居中对齐了，下面的图都是左对齐的，看着好难受)
而后，笔者使用简单的方法对模型进行处理，包括了补上接缝、平滑模型、强化边缘等操作，将其还原为岩石。
下面几张图皆为DDPM生成效果。
不足与展望
模型到图像的映射算法太简单，不能对带有孔洞、凹面的模型进行映射，且映射的网格非常不均匀。
训练时的DDPM模型并没有根据三维模型的特点进行改造，例如使用八面体映射产生的图像具有特殊的对称性，如果据此进行约束，那么神经网络模型生成的三维模型不会产生接缝。
DDPM模型过于简单，推理时分辨率需要与训练时分辨率相同效果才好，比较理想的是使用LDM(Latent Diffusion Model)，然后训练VAE网络转换为实际的图像，这样在泛化能力更强的同时也能提高图像分辨率。终极目标当然是爬取和生成更多的数据，以训练类似Stable Diffusion的文生图甚至图生图、模型生图的模型。
训练精度和训练次数不足，训练时使用的是batch_size=8，resolution=128的低分辨率参数进行训练(再多一点就显存了)，且只训练了30个epoch，效果非常有限。因此笔者赶紧下单了16G 4070Ti Super(。
