Mesh Diffusion
简介
一直想把PCG和AIGC结合起来，但网上此类相关论文和应用颇少，难以结合。近日深入学习了扩散模型，便觉扩散模型类似于泛化能力较强的GAN，适合训练大量数据进而生成新数据。于是产生了使用扩散模型生成三维模型的想法。 本次研究的目的是，根据已有的岩石模型，将岩石模型转换为一张张的图片并输入扩散模型中进行训练，训练完成后生成新的岩石模型。
准备训练集
模型to图像
因为Mesh的顶点数据不太好直接转换成图片，所以想了另一种直接将模型映射到图像或者说将方形图像映射到模型上的方法。 整个过程算法很简单，主要分为两个阶段，从方形映射到球体、从球体映射到模型。
从方形映射到球体
单纯的图片相当于一个方形面片，方形面片没有体积，不便于直接映射到模型上，因此先将方形映射到有体积的单位球体，而后将球体映射到模型。 方形映射到球体的主流算法有如下几个：
- 经纬球。最简单的映射方法，但经纬球的精度损失极大，图像的上下边最终只能映射到球体的两个极点上。
- Cubemap。只能展开成六个正方形，而不能只展开成一个方形面片，除非损失方形的一部分内容。
- 八面体映射。将一个正方形以折纸的形式映射到正八面体上，然后映射到单位球上(将位置进行归一化即可)。此方法非常简单，且精度损失相对较小，且为方形到球体的满射，因此选择此方法。
作者使用Houdini批量生成模型的“图像”，因此贴出Vex的实现代码：
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]
可见八面体映射确实简单异常，基本就是
从球体映射到模型
这里的算法作者就设计得比较敷衍了，基本上就是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的原理，可以参考我之前的文章(笔记)扩散模型。 模型推理代码Differs上也有提供，也就不再赘述。
数据精度
需要注意的是图像精度和颜色空间问题，若使用uint8也就是8bpc位深度保存图像，重新加载图像然后恢复模型时，由于精度问题会产生锯齿。
因此，无论是生成数据集，还是训练加载模型，还是推理后保存图像，都需要在float32的精度下进行。然而Diffusers的DDPM训练代码默认导入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精度的图像。
图像还原为模型
不足与展望
- 模型到图像的映射算法太简单，不能对带有孔洞、凹面的模型进行映射，且映射的网格非常不均匀。
- 训练时的DDPM模型并没有根据三维模型的特点进行改造，例如使用八面体映射产生的图像具有特殊的对称性，如果据此进行约束，那么神经网络模型生成的三维模型不会产生接缝。
- DDPM在推理时分辨率需要与训练时分辨率相同，对训练和推理的硬件要求都比较高，使用LDM(Latent Diffusion Model)能够提高推理效率。但后来尝试训练LDM的UNet而使用ldm-celebahq-256的VQVAE，但效果并不好，后续可以尝试。 终极目标当然是爬取和生成更多的数据，以训练类似Stable Diffusion的文生图甚至图生图、模型生图的模型。
- 训练精度和训练次数不足，训练时使用的是batch_size=8，resolution=192的低分辨率参数训练了150个epoch，效果非常有限。